Skip to content

Commit

Permalink
vectorstores: add doc and test for Replacement opt
Browse files Browse the repository at this point in the history
- add doc for `WithReplacement` opt
- add test in pgvector_test
  • Loading branch information
creeper12356 committed Nov 7, 2024
1 parent 8ffd9dd commit 700e66a
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 0 deletions.
4 changes: 4 additions & 0 deletions vectorstores/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,10 @@ func WithDeduplicater(fn func(ctx context.Context, doc schema.Document) bool) Op
}
}

// WithReplacement returns an Option for setting the replacement flag that could be used
// when adding documents. If set to true, the existing document with the **Metadata** will
// be replaced with the new document. If set to false(default case), the new document will
// be added to the store and the existing document will be kept.
func WithReplacement(replacement bool) Option {
return func(o *Options) {
o.Replacement = replacement
Expand Down
60 changes: 60 additions & 0 deletions vectorstores/pgvector/pgvector_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -615,6 +615,66 @@ func TestDeduplicater(t *testing.T) {
require.Equal(t, "vegetable", docs[0].Metadata["type"])
}

func TestReplacement(t *testing.T) {
t.Parallel()
pgvectorURL := preCheckEnvSetting(t)
ctx := context.Background()

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

conn, err := pgx.Connect(ctx, pgvectorURL)
require.NoError(t, err)

store, err := pgvector.New(
ctx,
pgvector.WithConn(conn),
pgvector.WithEmbedder(e),
)
require.NoError(t, err)

defer cleanupTestArtifacts(ctx, t, store, pgvectorURL)

_, err = store.AddDocuments(context.Background(), []schema.Document{
{PageContent: "Math", Metadata: map[string]any{
"courseID": "foo",
}},
{PageContent: "Physics", Metadata: map[string]any{
"courseID": "bar",
}},
})
require.NoError(t, err)

_, err = store.AddDocuments(context.Background(), []schema.Document{
{PageContent: "Discrete Math", Metadata: map[string]any{
"courseID": "foo",
}},
}, vectorstores.WithReplacement(true))
require.NoError(t, err)

docs, err := store.Search(ctx, 2, vectorstores.WithFilters(map[string]any{
"courseID": "foo",
}))
require.NoError(t, err)
require.Len(t, docs, 1)
require.Equal(t, "Discrete Math", docs[0].PageContent)

_, err = store.AddDocuments(context.Background(), []schema.Document{
{PageContent: "Calculus", Metadata: map[string]any{
"courseID": "foo",
}},
}, vectorstores.WithReplacement(false))
require.NoError(t, err)

docs, err = store.Search(ctx, 2, vectorstores.WithFilters(map[string]any{
"courseID": "foo",
}))
require.NoError(t, err)
require.Len(t, docs, 2)
}

func TestWithAllOptions(t *testing.T) {
t.Parallel()
pgvectorURL := preCheckEnvSetting(t)
Expand Down

0 comments on commit 700e66a

Please sign in to comment.