change: merge archive logic into getClient

gptscript-ai · Jul 1, 2024 · 6598b78 · 6598b78
1 parent c066081
commit 6598b78
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 52 deletions.
diff --git a/pkg/cmd/client.go b/pkg/cmd/client.go
@@ -13,7 +13,8 @@ import (
 )
 
 type Client struct {
-	Server string `usage:"URL of the Knowledge API Server" default:"" env:"KNOW_SERVER_URL"`
+	Server           string `usage:"URL of the Knowledge API Server" default:"" env:"KNOW_SERVER_URL"`
+	datastoreArchive string
 	config.OpenAIConfig
 	config.DatabaseConfig
 	config.VectorDBConfig
@@ -24,30 +25,33 @@ type ClientFlowsConfig struct {
 	Flow      string `usage:"Flow name" env:"KNOW_FLOW"`
 }
 
-func (s *Client) getClientFromArchive(archive string) (client.Client, error) {
+func (s *Client) loadArchive() error {
+	if s.datastoreArchive == "" {
+		return nil
+	}
 	// unpack to tempdir
 	tmpDir, err := os.MkdirTemp(os.TempDir(), "knowledge-retrieve-*")
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer os.RemoveAll(tmpDir)
 
-	r, err := zip.OpenReader(archive)
+	r, err := zip.OpenReader(s.datastoreArchive)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer r.Close()
 
 	if len(r.File) != 2 {
-		return nil, fmt.Errorf("knowledge archive must contain exactly two files, found %d", len(r.File))
+		return fmt.Errorf("knowledge archive must contain exactly two files, found %d", len(r.File))
 	}
 
 	dbFile := ""
 	vectorStoreFile := ""
 	for _, f := range r.File {
 		rc, err := f.Open()
 		if err != nil {
-			return nil, err
+			return err
 		}
 		defer rc.Close()
 
@@ -58,12 +62,12 @@ func (s *Client) getClientFromArchive(archive string) (client.Client, error) {
 
 		f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
 		if err != nil {
-			return nil, err
+			return err
 		}
 		defer f.Close()
 
 		if _, err := io.Copy(f, rc); err != nil {
-			return nil, err
+			return err
 		}
 		_ = f.Close()
 		_ = rc.Close()
@@ -77,17 +81,21 @@ func (s *Client) getClientFromArchive(archive string) (client.Client, error) {
 	}
 
 	if dbFile == "" || vectorStoreFile == "" {
-		return nil, fmt.Errorf("knowledge archive must contain exactly one .db and one .gob file")
+		return fmt.Errorf("knowledge archive must contain exactly one .db and one .gob file")
 	}
 
 	s.DSN = types.ArchivePrefix + dbFile
 	s.VectorDBPath = types.ArchivePrefix + vectorStoreFile
 
-	return s.getClient()
+	return nil
 }
 
 func (s *Client) getClient() (client.Client, error) {
 
+	if err := s.loadArchive(); err != nil {
+		return nil, err
+	}
+
 	if s.Server == "" || s.Server == "standalone" {
 		ds, err := datastore.NewDatastore(s.DSN, s.AutoMigrate == "true", s.VectorDBConfig.VectorDBPath, s.OpenAIConfig)
 		if err != nil {

diff --git a/pkg/cmd/get_dataset.go b/pkg/cmd/get_dataset.go
@@ -3,8 +3,6 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/gptscript-ai/knowledge/pkg/client"
-
 	"github.com/spf13/cobra"
 )
 
@@ -21,18 +19,9 @@ func (s *ClientGetDataset) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientGetDataset) Run(cmd *cobra.Command, args []string) error {
-	var err error
-	var c client.Client
-	if s.Archive != "" {
-		c, err = s.getClientFromArchive(s.Archive)
-		if err != nil {
-			return err
-		}
-	} else {
-		c, err = s.getClient()
-		if err != nil {
-			return err
-		}
+	c, err := s.getClient()
+	if err != nil {
+		return err
 	}
 
 	datasetID := args[0]

diff --git a/pkg/cmd/list_datasets.go b/pkg/cmd/list_datasets.go
@@ -3,8 +3,6 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/gptscript-ai/knowledge/pkg/client"
-
 	"github.com/spf13/cobra"
 )
 
@@ -20,18 +18,10 @@ func (s *ClientListDatasets) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientListDatasets) Run(cmd *cobra.Command, args []string) error {
-	var err error
-	var c client.Client
-	if s.Archive != "" {
-		c, err = s.getClientFromArchive(s.Archive)
-		if err != nil {
-			return err
-		}
-	} else {
-		c, err = s.getClient()
-		if err != nil {
-			return err
-		}
+
+	c, err := s.getClient()
+	if err != nil {
+		return err
 	}
 
 	ds, err := c.ListDatasets(cmd.Context())

diff --git a/pkg/cmd/retrieve.go b/pkg/cmd/retrieve.go
@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/gptscript-ai/knowledge/pkg/client"
 	"github.com/gptscript-ai/knowledge/pkg/datastore"
 	flowconfig "github.com/gptscript-ai/knowledge/pkg/flows/config"
 	vserr "github.com/gptscript-ai/knowledge/pkg/vectorstore/errors"
@@ -32,18 +31,9 @@ func (s *ClientRetrieve) Customize(cmd *cobra.Command) {
 
 func (s *ClientRetrieve) Run(cmd *cobra.Command, args []string) error {
 
-	var err error
-	var c client.Client
-	if s.Archive != "" {
-		c, err = s.getClientFromArchive(s.Archive)
-		if err != nil {
-			return err
-		}
-	} else {
-		c, err = s.getClient()
-		if err != nil {
-			return err
-		}
+	c, err := s.getClient()
+	if err != nil {
+		return err
 	}
 
 	datasetID := s.Dataset