add: pgvector vectorstore implementation

iwilltry42 · iwilltry42 · commit 586f9a7af949 · 2024-10-16T10:33:11.000+02:00
diff --git a/go.mod b/go.mod
@@ -28,6 +28,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/hupe1980/golc v0.0.112
 	github.com/iwilltry42/bm25-go v0.0.0-20240909111832-a928590cc9da
+	github.com/jackc/pgx/v5 v5.7.1
 	github.com/jmcarbo/stopwords v1.1.9
 	github.com/joho/godotenv v1.5.1
 	github.com/knadh/koanf/parsers/json v0.1.0
@@ -38,14 +39,15 @@ require (
 	github.com/ledongthuc/pdf v0.0.0-20240201131950-da5b75280b06
 	github.com/lu4p/cat v0.1.5
 	github.com/mitchellh/mapstructure v1.5.0
+	github.com/pgvector/pgvector-go v0.2.2
 	github.com/philippgille/chromem-go v0.6.1-0.20240811154507-a1944285b284
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
 	github.com/swaggo/files v1.0.1
 	github.com/swaggo/gin-swagger v1.6.0
 	github.com/swaggo/swag v1.16.3
 	github.com/tmc/langchaingo v0.1.12
-	golang.org/x/sync v0.7.0
+	golang.org/x/sync v0.8.0
 	gorm.io/gorm v1.25.10
 	sigs.k8s.io/yaml v1.4.0
 )
@@ -109,6 +111,8 @@ require (
 	github.com/hupe1980/go-textractor v0.0.9 // indirect
 	github.com/hupe1980/go-tiktoken v0.0.9 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/jaytaylor/html2text v0.0.0-20200412013138-3577fbdbcff7 // indirect
 	github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
@@ -158,11 +162,11 @@ require (
 	gitlab.com/golang-commonmark/mdurl v0.0.0-20191124015652-932350d1cb84 // indirect
 	gitlab.com/golang-commonmark/puny v0.0.0-20191124015043-9f83538fa04f // indirect
 	golang.org/x/arch v0.8.0 // indirect
-	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/crypto v0.27.0 // indirect
 	golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8 // indirect
 	golang.org/x/net v0.26.0 // indirect
-	golang.org/x/sys v0.21.0 // indirect
-	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/sys v0.25.0 // indirect
+	golang.org/x/text v0.18.0 // indirect
 	golang.org/x/tools v0.22.0 // indirect
 	google.golang.org/api v0.184.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240610135401-a8a62080eff3 // indirect
diff --git a/go.sum b/go.sum
diff --git a/pkg/cmd/askdir.go b/pkg/cmd/askdir.go
@@ -31,7 +31,7 @@ func (s *ClientAskDir) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientAskDir) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/client.go b/pkg/cmd/client.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"archive/zip"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -89,13 +90,13 @@ func (s *Client) loadArchive() error {
 		return fmt.Errorf("knowledge archive must contain exactly one .db and one .gob file")
 	}
 
-	s.DSN = types.ArchivePrefix + dbFile
-	s.VectorDBPath = types.ArchivePrefix + vectorStoreFile
+	s.DatabaseConfig.DSN = types.ArchivePrefix + dbFile
+	s.VectorDBConfig.DSN = types.ArchivePrefix + vectorStoreFile
 
 	return nil
 }
 
-func (s *Client) getClient() (client.Client, error) {
+func (s *Client) getClient(ctx context.Context) (client.Client, error) {
 	if err := s.loadArchive(); err != nil {
 		return nil, err
 	}
@@ -111,7 +112,7 @@ func (s *Client) getClient() (client.Client, error) {
 			return nil, err
 		}
 
-		ds, err := datastore.NewDatastore(s.DSN, s.AutoMigrate == "true", s.VectorDBConfig.VectorDBPath, provider)
+		ds, err := datastore.NewDatastore(ctx, s.DatabaseConfig.DSN, s.AutoMigrate == "true", s.VectorDBConfig.DSN, provider)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/cmd/create_dataset.go b/pkg/cmd/create_dataset.go
@@ -17,7 +17,7 @@ func (s *ClientCreateDataset) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientCreateDataset) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/delete_dataset.go b/pkg/cmd/delete_dataset.go
@@ -17,7 +17,7 @@ func (s *ClientDeleteDataset) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientDeleteDataset) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/edit_dataset.go b/pkg/cmd/edit_dataset.go
@@ -3,6 +3,7 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
+
 	"github.com/gptscript-ai/knowledge/pkg/datastore"
 	"github.com/gptscript-ai/knowledge/pkg/index"
 	"github.com/spf13/cobra"
@@ -23,7 +24,7 @@ func (s *ClientEditDataset) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientEditDataset) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/export.go b/pkg/cmd/export.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+
 	"github.com/spf13/cobra"
 )
 
@@ -17,7 +18,7 @@ func (s *ClientExportDatasets) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientExportDatasets) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/get_dataset.go b/pkg/cmd/get_dataset.go
@@ -3,6 +3,7 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
+
 	"github.com/spf13/cobra"
 )
 
@@ -19,7 +20,7 @@ func (s *ClientGetDataset) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientGetDataset) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/import.go b/pkg/cmd/import.go
@@ -22,7 +22,7 @@ func (s *ClientImportDatasets) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientImportDatasets) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/ingest.go b/pkg/cmd/ingest.go
@@ -52,7 +52,7 @@ This is a constraint of the Vector Database and Similarity Search, as different
 }
 
 func (s *ClientIngest) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/list_datasets.go b/pkg/cmd/list_datasets.go
@@ -3,6 +3,7 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
+
 	"github.com/spf13/cobra"
 )
 
@@ -18,7 +19,7 @@ func (s *ClientListDatasets) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientListDatasets) Run(cmd *cobra.Command, args []string) error {
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/load.go b/pkg/cmd/load.go
@@ -59,7 +59,6 @@ func (s *ClientLoad) Run(cmd *cobra.Command, args []string) error {
 
 	var texts []string
 	for _, doc := range docs {
-
 		if len(doc.Content) == 0 {
 			continue
 		}
diff --git a/pkg/cmd/reset.go b/pkg/cmd/reset.go
@@ -21,7 +21,7 @@ func (s *ClientResetDatastore) Customize(cmd *cobra.Command) {
 }
 
 func (s *ClientResetDatastore) Run(cmd *cobra.Command, args []string) error {
-	dsn, vectordbPath, _, err := datastore.GetDefaultDSNs(s.DSN, s.VectorDBConfig.VectorDBPath)
+	dsn, vectordbPath, _, err := datastore.GetDefaultDSNs(s.DatabaseConfig.DSN, s.VectorDBConfig.DSN)
 	if err != nil {
 		return err
 	}
@@ -34,6 +34,6 @@ func (s *ClientResetDatastore) Run(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to remove vector database directory: %w", err)
 	}
 
-	fmt.Printf("Successfully reset datastore (DSN: %q, VectorDBPath: %q)\n", dsn, vectordbPath)
+	fmt.Printf("Successfully reset datastore (DSN: %q, DSN: %q)\n", dsn, vectordbPath)
 	return nil
 }
diff --git a/pkg/cmd/retrieve.go b/pkg/cmd/retrieve.go
@@ -51,7 +51,7 @@ func (s *ClientRetrieve) Run(cmd *cobra.Command, args []string) error {
 	}
 	slog.Info("Retrieving sources for query", "query", query, "datasets", datasetIDs)
 
-	c, err := s.getClient()
+	c, err := s.getClient(cmd.Context())
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cmd/server.go b/pkg/cmd/server.go
@@ -2,12 +2,13 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/gptscript-ai/knowledge/pkg/datastore/embeddings"
-	"github.com/spf13/cobra"
 	"log/slog"
 	"os/signal"
 	"syscall"
 
+	"github.com/gptscript-ai/knowledge/pkg/datastore/embeddings"
+	"github.com/spf13/cobra"
+
 	"github.com/gptscript-ai/knowledge/pkg/config"
 	"github.com/gptscript-ai/knowledge/pkg/datastore"
 	"github.com/gptscript-ai/knowledge/pkg/server"
@@ -47,7 +48,7 @@ func (s *Server) Run(cmd *cobra.Command, _ []string) error {
 		return err
 	}
 
-	ds, err := datastore.NewDatastore(s.DSN, s.AutoMigrate == "true", s.VectorDBConfig.VectorDBPath, provider)
+	ds, err := datastore.NewDatastore(cmd.Context(), s.DatabaseConfig.DSN, s.AutoMigrate == "true", s.VectorDBConfig.DSN, provider)
 	if err != nil {
 		return fmt.Errorf("failed to initialize datastore: %w", err)
 	}
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -31,7 +31,7 @@ type DatabaseConfig struct {
 }
 
 type VectorDBConfig struct {
-	VectorDBPath string `usage:"VectorDBPath to the vector database (default \"chromem:$XDG_DATA_HOME/gptscript/knowledge/vector.db\")" default:"" env:"KNOW_VECTOR_DSN"`
+	DSN string `name:"vector-dsn" usage:"DSN to the vector database (default \"chromem:$XDG_DATA_HOME/gptscript/knowledge/vector.db\")" default:"" env:"KNOW_VECTOR_DSN"`
 }
 
 func LoadConfig(configFile string) (*Config, error) {
diff --git a/pkg/datastore/datastore.go b/pkg/datastore/datastore.go
@@ -81,7 +81,7 @@ func LogEmbeddingFunc(embeddingFunc cg.EmbeddingFunc) cg.EmbeddingFunc {
 	}
 }
 
-func NewDatastore(indexDSN string, automigrate bool, vectorDSN string, embeddingProvider etypes.EmbeddingModelProvider) (*Datastore, error) {
+func NewDatastore(ctx context.Context, indexDSN string, automigrate bool, vectorDSN string, embeddingProvider etypes.EmbeddingModelProvider) (*Datastore, error) {
 	indexDSN, vectorDSN, isArchive, err := GetDefaultDSNs(indexDSN, vectorDSN)
 	if err != nil {
 		return nil, fmt.Errorf("failed to determine datastore paths: %w", err)
@@ -98,7 +98,7 @@ func NewDatastore(indexDSN string, automigrate bool, vectorDSN string, embedding
 
 	slog.Debug("Using embedding model provider", "provider", embeddingProvider.Name(), "config", output.RedactSensitive(embeddingProvider.Config()))
 
-	vsdb, err := vectorstore.New(vectorDSN, embeddingProvider)
+	vsdb, err := vectorstore.New(ctx, vectorDSN, embeddingProvider)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/vectorstore/pgvector/pgvector.go b/pkg/vectorstore/pgvector/pgvector.go
diff --git a/pkg/vectorstore/vectorstores.go b/pkg/vectorstore/vectorstores.go

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ func (s ClientAskDir) Customize(cmd cobra.Command) {`
`31`	`31`	`}`
`32`	`32`
`33`	`33`	`func (s ClientAskDir) Run(cmd cobra.Command, args []string) error {`
`34`		`- c, err := s.getClient()`
	`34`	`+ c, err := s.getClient(cmd.Context())`
`35`	`35`	`if err != nil {`
`36`	`36`	`return err`
`37`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@ package cmd`
`2`	`2`
`3`	`3`	`import (`
`4`	`4`	`"archive/zip"`
	`5`	`+ "context"`
`5`	`6`	`"fmt"`
`6`	`7`	`"io"`
`7`	`8`	`"os"`
`@@ -89,13 +90,13 @@ func (s *Client) loadArchive() error {`
`89`	`90`	`return fmt.Errorf("knowledge archive must contain exactly one .db and one .gob file")`
`90`	`91`	`}`
`91`	`92`
`92`		`- s.DSN = types.ArchivePrefix + dbFile`
`93`		`- s.VectorDBPath = types.ArchivePrefix + vectorStoreFile`
	`93`	`+ s.DatabaseConfig.DSN = types.ArchivePrefix + dbFile`
	`94`	`+ s.VectorDBConfig.DSN = types.ArchivePrefix + vectorStoreFile`
`94`	`95`
`95`	`96`	`return nil`
`96`	`97`	`}`
`97`	`98`
`98`		`-func (s *Client) getClient() (client.Client, error) {`
	`99`	`+func (s *Client) getClient(ctx context.Context) (client.Client, error) {`
`99`	`100`	`if err := s.loadArchive(); err != nil {`
`100`	`101`	`return nil, err`
`101`	`102`	`}`
`@@ -111,7 +112,7 @@ func (s *Client) getClient() (client.Client, error) {`
`111`	`112`	`return nil, err`
`112`	`113`	`}`
`113`	`114`
`114`		`- ds, err := datastore.NewDatastore(s.DSN, s.AutoMigrate == "true", s.VectorDBConfig.VectorDBPath, provider)`
	`115`	`+ ds, err := datastore.NewDatastore(ctx, s.DatabaseConfig.DSN, s.AutoMigrate == "true", s.VectorDBConfig.DSN, provider)`
`115`	`116`	`if err != nil {`
`116`	`117`	`return nil, err`
`117`	`118`	`}`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ func (s ClientCreateDataset) Customize(cmd cobra.Command) {`
`17`	`17`	`}`
`18`	`18`
`19`	`19`	`func (s ClientCreateDataset) Run(cmd cobra.Command, args []string) error {`
`20`		`- c, err := s.getClient()`
	`20`	`+ c, err := s.getClient(cmd.Context())`
`21`	`21`	`if err != nil {`
`22`	`22`	`return err`
`23`	`23`	`}`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ func (s ClientDeleteDataset) Customize(cmd cobra.Command) {`
`17`	`17`	`}`
`18`	`18`
`19`	`19`	`func (s ClientDeleteDataset) Run(cmd cobra.Command, args []string) error {`
`20`		`- c, err := s.getClient()`
	`20`	`+ c, err := s.getClient(cmd.Context())`
`21`	`21`	`if err != nil {`
`22`	`22`	`return err`
`23`	`23`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,6 +3,7 @@ package cmd`
`3`	`3`	`import (`
`4`	`4`	`"encoding/json"`
`5`	`5`	`"fmt"`
	`6`	`+`
`6`	`7`	`"github.com/spf13/cobra"`
`7`	`8`	`)`
`8`	`9`
`@@ -19,7 +20,7 @@ func (s ClientGetDataset) Customize(cmd cobra.Command) {`
`19`	`20`	`}`
`20`	`21`
`21`	`22`	`func (s ClientGetDataset) Run(cmd cobra.Command, args []string) error {`
`22`		`- c, err := s.getClient()`
	`23`	`+ c, err := s.getClient(cmd.Context())`
`23`	`24`	`if err != nil {`
`24`	`25`	`return err`
`25`	`26`	`}`
Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ func (s ClientImportDatasets) Customize(cmd cobra.Command) {`
`22`	`22`	`}`
`23`	`23`
`24`	`24`	`func (s ClientImportDatasets) Run(cmd cobra.Command, args []string) error {`
`25`		`- c, err := s.getClient()`
	`25`	`+ c, err := s.getClient(cmd.Context())`
`26`	`26`	`if err != nil {`
`27`	`27`	`return err`
`28`	`28`	`}`