Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

vectorstores: add pgvector #377

Merged
merged 1 commit into from
Dec 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading