vectorstores: add pgvector

Signed-off-by: Abirdcfly <[email protected]>

Author: Abirdcfly

go.mod

@@ -61,6 +61,7 @@ require (
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.0 // indirect
 	github.com/oklog/ulid v1.3.1 // indirect
+	github.com/pgvector/pgvector-go v0.1.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.11.0 // indirect

go.sum

@@ -420,6 +420,8 @@ github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7J
 github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/pelletier/go-toml v1.7.0/go.mod h1:vwGMzjaWMwyfHwgIBhI2YUM4fB6nL6lVAvS1LBMMhTE=
+github.com/pgvector/pgvector-go v0.1.1 h1:kqJigGctFnlWvskUiYIvJRNwUtQl/aMSUZVs0YWQe+g=
+github.com/pgvector/pgvector-go v0.1.1/go.mod h1:wLJgD/ODkdtd2LJK4l6evHXTuG+8PxymYAVomKHOWac=
 github.com/pinecone-io/go-pinecone v0.3.0 h1:+t0CiYaaA+JN6YM9QRNlvfLEr2kkGzcVEj/xNmSAON4=
 github.com/pinecone-io/go-pinecone v0.3.0/go.mod h1:VdSieE1r4jT3XydjFi+iL5w9qsGRz/x8LxWach2Hnv8=
 github.com/pingcap/errors v0.11.4 h1:lFuQV/oaUMGcD2tqt+01ROSmJs75VG1ToEOkZIZ4nE4=

vectorstores/chroma/doc.go

@@ -1,2 +1,2 @@
-// Package chroma contains an implementation of the vectorStore interface that connects to an external Chroma database.
+// Package chroma contains an implementation of the VectorStore interface that connects to an external Chroma database.
 package chroma

vectorstores/pgvector/doc.go

@@ -0,0 +1,3 @@
+// Package pgvector contains an implementation of the VectorStore
+// interface using pgvector.
+package pgvector

vectorstores/pgvector/options.go

@@ -0,0 +1,92 @@
+package pgvector
+
+import (
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/tmc/langchaingo/embeddings"
+)
+
+const (
+	DefaultCollectionName           = "langchain"
+	DefaultPreDeleteCollection      = false
+	DefaultEmbeddingStoreTableName  = "langchain_pg_embedding"
+	DefaultCollectionStoreTableName = "langchain_pg_collection"
+)
+
+// ErrInvalidOptions is returned when the options given are invalid.
+var ErrInvalidOptions = errors.New("invalid options")
+
+// Option is a function type that can be used to modify the client.
+type Option func(p *Store)
+
+// WithEmbedder is an option for setting the embedder to use. Must be set.
+func WithEmbedder(e embeddings.Embedder) Option {
+	return func(p *Store) {
+		p.embedder = e
+	}
+}
+
+// WithConnectionURL is an option for specifying the Postgres connection URL. Must be set.
+func WithConnectionURL(connectionURL string) Option {
+	return func(p *Store) {
+		p.postgresConnectionURL = connectionURL
+	}
+}
+
+// WithPreDeleteCollection is an option for setting if the collection should be deleted before creating.
+func WithPreDeleteCollection(preDelete bool) Option {
+	return func(p *Store) {
+		p.preDeleteCollection = preDelete
+	}
+}
+
+// WithCollectionName is an option for specifying the collection name.
+func WithCollectionName(name string) Option {
+	return func(p *Store) {
+		p.collectionName = name
+	}
+}
+
+// WithEmbeddingTableName is an option for specifying the embedding table name.
+func WithEmbeddingTableName(name string) Option {
+	return func(p *Store) {
+		p.embeddingTableName = pgx.Identifier{name}.Sanitize()
+	}
+}
+
+// WithCollectionTableName is an option for specifying the collection table name.
+func WithCollectionTableName(name string) Option {
+	return func(p *Store) {
+		p.collectionTableName = pgx.Identifier{name}.Sanitize()
+	}
+}
+
+func applyClientOptions(opts ...Option) (Store, error) {
+	o := &Store{
+		collectionName:      DefaultCollectionName,
+		preDeleteCollection: DefaultPreDeleteCollection,
+		embeddingTableName:  DefaultEmbeddingStoreTableName,
+		collectionTableName: DefaultCollectionStoreTableName,
+	}
+
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	if o.postgresConnectionURL == "" {
+		o.postgresConnectionURL = os.Getenv("PGVECTOR_CONNECTION_STRING")
+	}
+
+	if o.postgresConnectionURL == "" {
+		return Store{}, fmt.Errorf("%w: missing postgresConnectionURL", ErrInvalidOptions)
+	}
+
+	if o.embedder == nil {
+		return Store{}, fmt.Errorf("%w: missing embedder", ErrInvalidOptions)
+	}
+
+	return *o, nil
+}