feat: improve Jina AI impl.

youngsecurity · Sep 16, 2024 · d2e2a65 · d2e2a65
1 parent a6fc13d
commit d2e2a65
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 108 deletions.
diff --git a/README.md b/README.md
@@ -215,6 +215,8 @@ Application Options:
       --transcript                  Grab transcript from YouTube video and send to chat
       --comments                    Grab comments from YouTube video and send to chat
       --dry-run                     Show what would be sent to the model without actually sending it
+  -u, --scrape_url=                 Scrape website URL to markdown using Jina AI
+  -q, --scrape_question=            Search question using Jina AI
 
 Help Options:
   -h, --help                        Show this help message

diff --git a/cli/cli.go b/cli/cli.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/danielmiessler/fabric/core"
 	"github.com/danielmiessler/fabric/db"
-	"github.com/danielmiessler/fabric/jina"
 )
 
 // Cli Controls the cli. It takes in the flags and runs the appropriate functions
@@ -122,11 +121,7 @@ func Cli() (message string, err error) {
 
 			fmt.Println(transcript)
 
-			if currentFlags.Message != "" {
-				currentFlags.Message = currentFlags.Message + "\n" + transcript
-			} else {
-				currentFlags.Message = transcript
-			}
+			currentFlags.AppendMessage(transcript)
 		}
 
 		if currentFlags.YouTubeComments {
@@ -139,11 +134,7 @@ func Cli() (message string, err error) {
 
 			fmt.Println(commentsString)
 
-			if currentFlags.Message != "" {
-				currentFlags.Message = currentFlags.Message + "\n" + commentsString
-			} else {
-				currentFlags.Message = commentsString
-			}
+			currentFlags.AppendMessage(commentsString)
 		}
 
 		if currentFlags.Pattern == "" {
@@ -152,35 +143,34 @@ func Cli() (message string, err error) {
 		}
 	}
 
-    // Initialize JinaClient
-    jinaClient := jina.NewJinaClient()
-
-    // Load the configuration for JinaClient, including the API key
-    if err = jinaClient.Configurable.Configure(); err != nil {
-        return "", fmt.Errorf("failed to configure JinaClient: %w", err)
-    }
-
-    // Check if the scrape_url flag is set and call ScrapeURL
-    if currentFlags.ScrapeURL != "" {
-        message, err = jinaClient.ScrapeURL(currentFlags.ScrapeURL)
-        if err != nil {
-            return "", fmt.Errorf("failed to scrape URL: %w", err)
-        }
-        fmt.Println(message)
-        return message, nil
-    }
-
-    // Check if the scrape_question flag is set and call ScrapeQuestion
-    if currentFlags.ScrapeQuestion != "" {
-        message, err = jinaClient.ScrapeQuestion(currentFlags.ScrapeQuestion)
-        if err != nil {
-            return "", fmt.Errorf("failed to scrape question: %w", err)
-        }
-        fmt.Println(message)
-        return message, nil
-    }
-
-
+	if (currentFlags.ScrapeURL != "" || currentFlags.ScrapeQuestion != "") && fabric.Jina.IsConfigured() {
+		// Check if the scrape_url flag is set and call ScrapeURL
+		if currentFlags.ScrapeURL != "" {
+			if message, err = fabric.Jina.ScrapeURL(currentFlags.ScrapeURL); err != nil {
+				return
+			}
+
+			fmt.Println(message)
+
+			currentFlags.AppendMessage(message)
+		}
+
+		// Check if the scrape_question flag is set and call ScrapeQuestion
+		if currentFlags.ScrapeQuestion != "" {
+			if message, err = fabric.Jina.ScrapeQuestion(currentFlags.ScrapeQuestion); err != nil {
+				return
+			}
+
+			fmt.Println(message)
+
+			currentFlags.AppendMessage(message)
+		}
+
+		if currentFlags.Pattern == "" {
+			// if the pattern flag is not set, we wanted only to grab the url or get the answer to the question
+			return
+		}
+	}
 
 	var chatter *core.Chatter
 	if chatter, err = fabric.GetChatter(currentFlags.Model, currentFlags.Stream, currentFlags.DryRun); err != nil {

diff --git a/cli/flags.go b/cli/flags.go
@@ -40,24 +40,23 @@ type Flags struct {
 	YouTubeTranscript       bool              `long:"transcript" description:"Grab transcript from YouTube video and send to chat"`
 	YouTubeComments         bool              `long:"comments" description:"Grab comments from YouTube video and send to chat"`
 	DryRun                  bool              `long:"dry-run" description:"Show what would be sent to the model without actually sending it"`
-    ScrapeURL               string            `short:"u" long:"scrape_url" description:"Scrape website URL to markdown using Jina AI"`
-	ScrapeQuestion		    string            `short:"q" long:"scrape_question" description:"Search question using Jina AI"`
-
+	ScrapeURL               string            `short:"u" long:"scrape_url" description:"Scrape website URL to markdown using Jina AI"`
+	ScrapeQuestion          string            `short:"q" long:"scrape_question" description:"Search question using Jina AI"`
 }
 
 // Init Initialize flags. returns a Flags struct and an error
 func Init() (ret *Flags, err error) {
-    var message string
+	var message string
 
-    ret = &Flags{}
-    parser := flags.NewParser(ret, flags.Default)
-    var args []string
-    if args, err = parser.Parse(); err != nil {
-        return
-    }
+	ret = &Flags{}
+	parser := flags.NewParser(ret, flags.Default)
+	var args []string
+	if args, err = parser.Parse(); err != nil {
+		return
+	}
 
-    info, _ := os.Stdin.Stat()
-    hasStdin := (info.Mode() & os.ModeCharDevice) == 0
+	info, _ := os.Stdin.Stat()
+	hasStdin := (info.Mode() & os.ModeCharDevice) == 0
 
 	// takes input from stdin if it exists, otherwise takes input from args (the last argument)
 	if hasStdin {
@@ -71,7 +70,7 @@ func Init() (ret *Flags, err error) {
 	}
 	ret.Message = message
 
-    return
+	return
 }
 
 // readStdin reads from stdin and returns the input as a string or an error
@@ -112,3 +111,12 @@ func (o *Flags) BuildChatRequest() (ret *common.ChatRequest) {
 	}
 	return
 }
+
+func (o *Flags) AppendMessage(message string) {
+	if o.Message != "" {
+		o.Message = o.Message + "\n" + message
+	} else {
+		o.Message = message
+	}
+	return
+}
diff --git a/core/fabric.go b/core/fabric.go
@@ -51,7 +51,7 @@ func NewFabricBase(db *db.Db) (ret *Fabric) {
 		VendorsAll:     NewVendorsManager(),
 		PatternsLoader: NewPatternsLoader(db.Patterns),
 		YouTube:        youtube.NewYouTube(),
-		Jina: 		 	jina.NewJinaClient(),
+		Jina:           jina.NewClient(),
 	}
 
 	label := "Default"
@@ -77,7 +77,7 @@ type Fabric struct {
 	VendorsAll *VendorsManager
 	*PatternsLoader
 	*youtube.YouTube
-	Jina *jina.JinaClient
+	Jina *jina.Client
 
 	Db *db.Db
 

diff --git a/jina/jina.go b/jina/jina.go
@@ -10,76 +10,61 @@ import (
 	"github.com/danielmiessler/fabric/common"
 )
 
-type JinaClient struct {
+type Client struct {
 	*common.Configurable
 	ApiKey *common.SetupQuestion
 }
 
-func NewJinaClient() *JinaClient {
+func NewClient() (ret *Client) {
 
 	label := "Jina AI"
 
-	client := &JinaClient{
+	ret = &Client{
 		Configurable: &common.Configurable{
-			Label: label,
+			Label:         label,
 			EnvNamePrefix: common.BuildEnvVariablePrefix(label),
 		},
 	}
-	client.ApiKey = client.AddSetupQuestion("API Key", false)
-    return client
-}
-
-// return the main content of a webpage in clean, LLM-friendly text.
-func (jc *JinaClient) ScrapeURL(url string) (string, error) {
-	requestURL := "https://r.jina.ai/" + url
-    req, err := http.NewRequest("GET", requestURL, nil)
-    if err != nil {
-        return "", fmt.Errorf("error creating request: %w", err)
-    }
-
-	// if api keys exist, set the header
-	if apiKey := jc.ApiKey.Value; apiKey != "" {
-        req.Header.Set("Authorization", "Bearer "+apiKey)
-    }
 
-    client := &http.Client{}
-    resp, err := client.Do(req)
-    if err != nil {
-        return "", fmt.Errorf("error sending request: %w", err)
-    }
-    defer resp.Body.Close()
+	ret.ApiKey = ret.AddSetupQuestion("API Key", false)
 
-    body, err := io.ReadAll(resp.Body)
-    if err != nil {
-        return "", fmt.Errorf("error reading response body: %w", err)
-    }
+	return
+}
 
-    return string(body), nil
+// ScrapeURL return the main content of a webpage in clean, LLM-friendly text.
+func (jc *Client) ScrapeURL(url string) (ret string, err error) {
+	return jc.request(fmt.Sprintf("https://r.jina.ai/%s", url))
 }
 
-func (jc *JinaClient) ScrapeQuestion(question string) (string, error) {
-    requestURL := "https://s.jina.ai/" + question
-    req, err := http.NewRequest("GET", requestURL, nil)
-    if err != nil {
-        return "", fmt.Errorf("error creating request: %w", err)
-    }
+func (jc *Client) ScrapeQuestion(question string) (ret string, err error) {
+	return jc.request(fmt.Sprintf("https://s.jina.ai/%s", question))
+}
 
-    // if api keys exist, set the header
-    if apiKey := jc.ApiKey.Value; apiKey != "" {
-        req.Header.Set("Authorization", "Bearer "+apiKey)
-    }
+func (jc *Client) request(requestURL string) (ret string, err error) {
+	var req *http.Request
+	if req, err = http.NewRequest("GET", requestURL, nil); err != nil {
+		err = fmt.Errorf("error creating request: %w", err)
+		return
+	}
 
-    client := &http.Client{}
-    resp, err := client.Do(req)
-    if err != nil {
-        return "", fmt.Errorf("error sending request: %w", err)
-    }
-    defer resp.Body.Close()
+	// if api keys exist, set the header
+	if jc.ApiKey.Value != "" {
+		req.Header.Set("Authorization", "Bearer "+jc.ApiKey.Value)
+	}
 
-    body, err := io.ReadAll(resp.Body)
-    if err != nil {
-        return "", fmt.Errorf("error reading response body: %w", err)
-    }
+	client := &http.Client{}
+	var resp *http.Response
+	if resp, err = client.Do(req); err != nil {
+		err = fmt.Errorf("error sending request: %w", err)
+		return
+	}
+	defer resp.Body.Close()
 
-    return string(body), nil
-}
+	var body []byte
+	if body, err = io.ReadAll(resp.Body); err != nil {
+		err = fmt.Errorf("error reading response body: %w", err)
+		return
+	}
+	ret = string(body)
+	return
+}