-
-
Notifications
You must be signed in to change notification settings - Fork 718
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Abirdcfly <fp544037857@gmail.com>
- Loading branch information
Showing
9 changed files
with
845 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
} |
Oops, something went wrong.