-
-
Notifications
You must be signed in to change notification settings - Fork 607
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 mongovector #1005
Merged
Merged
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
898ae09
GODRIVER-3305 POC
prestonvasquez fb3cf39
GODRIVER-3305 Extend test seed to be full doc
prestonvasquez ea4cd0b
GODRIVER-3305 Cont. w/ tests
prestonvasquez daff915
GODRIVER-3305 Cleanup comments and add setup test flag
prestonvasquez c5e1e20
GODRIVER-3305 Fix lintin errors
prestonvasquez 657a618
GODRIVER-3305 Fix linting errors
prestonvasquez 850fe06
GODRIVER-3305 Cont. Cleanup
prestonvasquez 922aae5
GODRIVER-3305 Update for free tier
prestonvasquez 1e99999
mongovector: Add doc.go
tmc File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 |
---|---|---|
@@ -0,0 +1,46 @@ | ||
// Package mongovector implements a vector store using MongoDB as the backend. | ||
// | ||
// The mongovector package provides a way to store and retrieve document embeddings | ||
// using MongoDB's vector search capabilities. It implements the VectorStore | ||
// interface from the vectorstores package, allowing it to be used interchangeably | ||
// with other vector store implementations. | ||
// | ||
// Key features: | ||
// - Store document embeddings in MongoDB | ||
// - Perform similarity searches on stored embeddings | ||
// - Configurable index and path settings | ||
// - Support for custom embedding functions | ||
// | ||
// Main types: | ||
// - Store: The main type that implements the VectorStore interface | ||
// - Option: A function type for configuring the Store | ||
// | ||
// Usage: | ||
// | ||
// import ( | ||
// "github.com/tmc/langchaingo/vectorstores/mongovector" | ||
// "go.mongodb.org/mongo-driver/mongo" | ||
// ) | ||
// | ||
// // Create a new Store | ||
// coll := // ... obtain a *mongo.Collection | ||
// embedder := // ... obtain an embeddings.Embedder | ||
// store := mongovector.New(coll, embedder) | ||
// | ||
// // Add documents | ||
// docs := []schema.Document{ | ||
// {PageContent: "Document 1"}, | ||
// {PageContent: "Document 2"}, | ||
// } | ||
// ids, err := store.AddDocuments(context.Background(), docs) | ||
// | ||
// // Perform similarity search | ||
// results, err := store.SimilaritySearch(context.Background(), "query", 5) | ||
// | ||
// The package also provides options for customizing the Store: | ||
// - WithIndex: Set a custom index name | ||
// - WithPath: Set a custom path for the vector field | ||
// - WithNumCandidates: Set the number of candidates for similarity search | ||
// | ||
// For more detailed information, see the documentation for individual types and functions. | ||
package mongovector |
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,207 @@ | ||
package mongovector | ||
|
||
import ( | ||
"context" | ||
"crypto/rand" | ||
"fmt" | ||
"math/big" | ||
"time" | ||
|
||
"github.com/tmc/langchaingo/embeddings" | ||
"github.com/tmc/langchaingo/schema" | ||
"github.com/tmc/langchaingo/vectorstores" | ||
) | ||
|
||
type mockEmbedder struct { | ||
queryVector []float32 | ||
docs map[string]schema.Document | ||
docVectors map[string][]float32 | ||
} | ||
|
||
var _ embeddings.Embedder = &mockEmbedder{} | ||
|
||
func newMockEmbedder(dim int) *mockEmbedder { | ||
emb := &mockEmbedder{ | ||
queryVector: newNormalizedVector(dim), | ||
docs: make(map[string]schema.Document), | ||
docVectors: make(map[string][]float32), | ||
} | ||
|
||
return emb | ||
} | ||
|
||
// mockDocuments will add the given documents to the embedder, assigning each | ||
// a vector such that similarity score = 0.5 * ( 1 + vector * queryVector). | ||
func (emb *mockEmbedder) mockDocuments(doc ...schema.Document) { | ||
for _, d := range doc { | ||
emb.docs[d.PageContent] = d | ||
} | ||
} | ||
|
||
// existingVectors returns all the vectors that have been added to the embedder. | ||
// The query vector is included in the list to maintain orthogonality. | ||
func (emb *mockEmbedder) existingVectors() [][]float32 { | ||
vectors := make([][]float32, 0, len(emb.docs)+1) | ||
for _, vec := range emb.docVectors { | ||
vectors = append(vectors, vec) | ||
} | ||
|
||
return append(vectors, emb.queryVector) | ||
} | ||
|
||
// EmbedDocuments will return the embedded vectors for the given texts. If the | ||
// text does not exist in the document set, a zero vector will be returned. | ||
func (emb *mockEmbedder) EmbedDocuments(_ context.Context, texts []string) ([][]float32, error) { | ||
vectors := make([][]float32, len(texts)) | ||
for i := range vectors { | ||
// If the text does not exist in the document set, return a zero vector. | ||
doc, ok := emb.docs[texts[i]] | ||
if !ok { | ||
vectors[i] = make([]float32, len(emb.queryVector)) | ||
} | ||
|
||
// If the vector exists, use it. | ||
existing, ok := emb.docVectors[texts[i]] | ||
if ok { | ||
vectors[i] = existing | ||
|
||
continue | ||
} | ||
|
||
// If it does not exist, make a linearly independent vector. | ||
newVectorBasis := newOrthogonalVector(len(emb.queryVector), emb.existingVectors()...) | ||
|
||
// Update the newVector to be scaled by the score. | ||
newVector := dotProductNormFn(doc.Score, emb.queryVector, newVectorBasis) | ||
|
||
vectors[i] = newVector | ||
emb.docVectors[texts[i]] = newVector | ||
} | ||
|
||
return vectors, nil | ||
} | ||
|
||
// EmbedQuery returns the query vector. | ||
func (emb *mockEmbedder) EmbedQuery(context.Context, string) ([]float32, error) { | ||
return emb.queryVector, nil | ||
} | ||
|
||
// Insert all of the mock documents collected by the embedder. | ||
func flushMockDocuments(ctx context.Context, store Store, emb *mockEmbedder) error { | ||
docs := make([]schema.Document, 0, len(emb.docs)) | ||
for _, doc := range emb.docs { | ||
docs = append(docs, doc) | ||
} | ||
|
||
_, err := store.AddDocuments(ctx, docs, vectorstores.WithEmbedder(emb)) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
// Consistency on indexes is not synchronous. | ||
// nolint:mnd | ||
time.Sleep(10 * time.Second) | ||
|
||
return nil | ||
} | ||
|
||
// newNormalizedFloat32 will generate a random float32 in [-1, 1]. | ||
// nolint:mnd | ||
func newNormalizedFloat32() (float32, error) { | ||
max := big.NewInt(1 << 24) | ||
|
||
n, err := rand.Int(rand.Reader, max) | ||
if err != nil { | ||
return 0.0, fmt.Errorf("failed to normalize float32") | ||
} | ||
|
||
return 2.0*(float32(n.Int64())/float32(1<<24)) - 1.0, nil | ||
} | ||
|
||
// dotProduct will return the dot product between two slices of f32. | ||
func dotProduct(v1, v2 []float32) float32 { | ||
var sum float32 | ||
|
||
for i := range v1 { | ||
sum += v1[i] * v2[i] | ||
} | ||
|
||
return sum | ||
} | ||
|
||
// linearlyIndependent true if the vectors are linearly independent. | ||
func linearlyIndependent(v1, v2 []float32) bool { | ||
var ratio float32 | ||
|
||
for i := range v1 { | ||
if v1[i] != 0 { | ||
r := v2[i] / v1[i] | ||
|
||
if ratio == 0 { | ||
ratio = r | ||
|
||
continue | ||
} | ||
|
||
if r == ratio { | ||
continue | ||
} | ||
|
||
return true | ||
} | ||
|
||
if v2[i] != 0 { | ||
return true | ||
} | ||
} | ||
|
||
return false | ||
} | ||
|
||
// Create a vector of values between [-1, 1] of the specified size. | ||
func newNormalizedVector(dim int) []float32 { | ||
vector := make([]float32, dim) | ||
for i := range vector { | ||
vector[i], _ = newNormalizedFloat32() | ||
} | ||
|
||
return vector | ||
} | ||
|
||
// Use Gram Schmidt to return a vector orthogonal to the basis, so long as | ||
// the vectors in the basis are linearly independent. | ||
func newOrthogonalVector(dim int, basis ...[]float32) []float32 { | ||
candidate := newNormalizedVector(dim) | ||
|
||
for _, b := range basis { | ||
dp := dotProduct(candidate, b) | ||
basisNorm := dotProduct(b, b) | ||
|
||
for i := range candidate { | ||
candidate[i] -= (dp / basisNorm) * b[i] | ||
} | ||
} | ||
|
||
return candidate | ||
} | ||
|
||
// return a new vector such that v1 * v2 = 2S - 1. | ||
func dotProductNormFn(score float32, qvector, basis []float32) []float32 { | ||
var sum float32 | ||
|
||
// Populate v2 upto dim-1. | ||
for i := range qvector[:len(qvector)-1] { | ||
sum += qvector[i] * basis[i] | ||
} | ||
|
||
// Calculate v_{2, dim} such that v1 * v2 = 2S - 1: | ||
basis[len(basis)-1] = (2*score - 1 - sum) / qvector[len(qvector)-1] | ||
|
||
// If the vectors are linearly independent, regenerate the dim-1 elements | ||
// of v2. | ||
if !linearlyIndependent(qvector, basis) { | ||
return dotProductNormFn(score, qvector, basis) | ||
} | ||
|
||
return basis | ||
} |
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,38 @@ | ||
package mongovector | ||
|
||
import ( | ||
"context" | ||
|
||
"github.com/tmc/langchaingo/embeddings" | ||
) | ||
|
||
// mockLLM will create consistent text embeddings mocking the OpenAI | ||
// text-embedding-3-small algorithm. | ||
type mockLLM struct { | ||
seen map[string][]float32 | ||
dim int | ||
} | ||
|
||
var _ embeddings.EmbedderClient = &mockLLM{} | ||
|
||
// createEmbedding will return vector embeddings for the mock LLM, maintaining | ||
// consistency. | ||
func (emb *mockLLM) CreateEmbedding(_ context.Context, texts []string) ([][]float32, error) { | ||
if emb.seen == nil { | ||
emb.seen = map[string][]float32{} | ||
} | ||
|
||
vectors := make([][]float32, len(texts)) | ||
for i, text := range texts { | ||
if f32s := emb.seen[text]; len(f32s) > 0 { | ||
vectors[i] = f32s | ||
|
||
continue | ||
} | ||
|
||
vectors[i] = newNormalizedVector(emb.dim) | ||
emb.seen[text] = vectors[i] // ensure consistency | ||
} | ||
|
||
return vectors, nil | ||
} |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you move these mocks into a separate subpackage (mocks? mongovectormocks?)