Skip to content

Commit

Permalink
vectorstores: add pgvector (#377)
Browse files Browse the repository at this point in the history
Signed-off-by: Abirdcfly <fp544037857@gmail.com>
  • Loading branch information
Abirdcfly authored Dec 13, 2023
1 parent 00f364f commit ed0536f
Show file tree
Hide file tree
Showing 9 changed files with 845 additions and 3 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
2 changes: 1 addition & 1 deletion vectorstores/chroma/doc.go
Original file line number Diff line number Diff line change
@@ -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
3 changes: 3 additions & 0 deletions vectorstores/pgvector/doc.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
// Package pgvector contains an implementation of the VectorStore
// interface using pgvector.
package pgvector
92 changes: 92 additions & 0 deletions vectorstores/pgvector/options.go
Original file line number Diff line number Diff line change
@@ -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
}
Loading

0 comments on commit ed0536f

Please sign in to comment.