Skip to content

Commit

Permalink
weaviate: allow embedder override in vectorstores options
Browse files Browse the repository at this point in the history
The changes allow the embedder to be overridden in the vectorstores options.
This is done by changing the embedder calls in `AddDocuments` and `SimilaritySearch`
functions to use the embedder from the options instead of the store.

The default embedder is still the one from the store, but it can be overwritten
by passing a `vectorstores.WithEmbedder` option.

A new test case 'TestWeaviateWithOptionEmbedder' is added to ensure that the
embedder provided as an option to either `AddDocuments` or `SimilaritySearch` takes
precedence over the one provided when creating the Store.

This change provides more flexibility in choosing the embedder at runtime,
which can be useful in scenarios where the same store is used with
different embedders.
  • Loading branch information
corani committed Jan 18, 2024
1 parent 61358e0 commit 06a33be
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 3 deletions.
10 changes: 7 additions & 3 deletions vectorstores/weaviate/weaviate.go
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ func (s Store) AddDocuments(ctx context.Context,
texts = append(texts, doc.PageContent)
}

vectors, err := s.embedder.EmbedDocuments(ctx, texts)
vectors, err := opts.Embedder.EmbedDocuments(ctx, texts)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -157,7 +157,7 @@ func (s Store) SimilaritySearch(
return nil, err
}

vector, err := s.embedder.EmbedQuery(ctx, query)
vector, err := opts.Embedder.EmbedQuery(ctx, query)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -244,7 +244,11 @@ func (s Store) getFilters(opts vectorstores.Options) any {
}

func (s Store) getOptions(options ...vectorstores.Option) vectorstores.Options {
opts := vectorstores.Options{}
// use the embedder from the store by default, this can be overwritten by passing
// an `vectorstores.WithEmbedder` option.
opts := vectorstores.Options{
Embedder: s.embedder,
}
for _, opt := range options {
opt(&opts)
}
Expand Down
55 changes: 55 additions & 0 deletions vectorstores/weaviate/weaviate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -545,3 +545,58 @@ func TestWeaviateAsRetrieverWithMetadataFilters(t *testing.T) {
require.NotContains(t, result, "orange", "expected not orange in result")
require.NotContains(t, result, "yellow", "expected not yellow in result")
}

// TestWeaviateWithOptionEmbedder ensures that the embedder provided as an option to either
// `AddDocuments` or `SimilaritySearch` takes precedence over the one provided when creating
// the `Store`.
func TestWeaviateWithOptionEmbedder(t *testing.T) {
t.Parallel()

scheme, host := getValues(t)

llm, err := openai.New()
require.NoError(t, err)

notme, err := embeddings.NewEmbedder(
embeddings.EmbedderClientFunc(func(context.Context, []string) ([][]float32, error) {
require.FailNow(t, "wrong embedder was called")
return nil, nil
}),
)
require.NoError(t, err)

butme, err := embeddings.NewEmbedder(
embeddings.EmbedderClientFunc(func(ctx context.Context, texts []string) ([][]float32, error) {
return llm.CreateEmbedding(ctx, texts)
}),
)
require.NoError(t, err)

store, err := New(
WithScheme(scheme),
WithHost(host),
WithEmbedder(notme),
WithNameSpace(uuid.New().String()),
WithIndexName(randomizedCamelCaseClass()),
WithQueryAttrs([]string{"location"}),
)
require.NoError(t, err)

err = createTestClass(context.Background(), store)
require.NoError(t, err)

_, err = store.AddDocuments(context.Background(), []schema.Document{
{PageContent: "tokyo", Metadata: map[string]any{
"country": "japan",
}},
{PageContent: "potato"},
}, vectorstores.WithEmbedder(butme))
require.NoError(t, err)

docs, err := store.SimilaritySearch(context.Background(), "japan", 1,
vectorstores.WithEmbedder(butme))
require.NoError(t, err)
require.Len(t, docs, 1)
require.Equal(t, "tokyo", docs[0].PageContent)
require.Equal(t, "japan", docs[0].Metadata["country"])
}

0 comments on commit 06a33be

Please sign in to comment.