Merge pull request #610 from dddbliss/patch-1

chroma: update chroma-go functions to pass context
tmc · Feb 19, 2024 · 1b64e10 · 1b64e10
2 parents fcd78c9 + 9ac0e74
commit 1b64e10
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/go.mod b/go.mod
@@ -114,7 +114,7 @@ require (
 	cloud.google.com/go/vertexai v0.6.0
 	github.com/Masterminds/sprig/v3 v3.2.3
 	github.com/PuerkitoBio/goquery v1.8.1
-	github.com/amikos-tech/chroma-go v0.0.0-20231228181736-e8f5e927093e
+	github.com/amikos-tech/chroma-go v0.0.0-20240109142503-c8fb49c3e28c
 	github.com/cohere-ai/tokenizer v1.1.2
 	github.com/go-openapi/strfmt v0.21.3
 	github.com/go-sql-driver/mysql v1.7.1

diff --git a/go.sum b/go.sum
@@ -69,8 +69,8 @@ github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578/go.mod h1:uGdko
 github.com/Shopify/goreferrer v0.0.0-20181106222321-ec9c9a553398/go.mod h1:a1uqRtAwp2Xwc6WNPJEufxJ7fx3npB4UV/JOLmbu5I0=
 github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o=
 github.com/ajg/form v1.5.1/go.mod h1:uL1WgH+h2mgNtvBq0339dVnzXdBETtL2LeUXaIv25UY=
-github.com/amikos-tech/chroma-go v0.0.0-20231228181736-e8f5e927093e h1:68YtjXFDWgquXICyGFa0gZ+nqZbs4OURutAyUjH1mYs=
-github.com/amikos-tech/chroma-go v0.0.0-20231228181736-e8f5e927093e/go.mod h1:uJwgGN4rBUTMI88Rn68Xia+cTRogOo0/elcPvJYFtBU=
+github.com/amikos-tech/chroma-go v0.0.0-20240109142503-c8fb49c3e28c h1:DzSXJSUVK5ed+HhAnGA1TKgBdXqAW/iXPRvpd9rXWLc=
+github.com/amikos-tech/chroma-go v0.0.0-20240109142503-c8fb49c3e28c/go.mod h1:uJwgGN4rBUTMI88Rn68Xia+cTRogOo0/elcPvJYFtBU=
 github.com/andybalholm/cascadia v1.3.1/go.mod h1:R4bJ1UQfqADjvDa4P6HZHLh/3OxWWEqc0Sk8XGwHqvA=
 github.com/andybalholm/cascadia v1.3.2 h1:3Xi6Dw5lHF15JtdcmAHD3i1+T8plmv7BQ/nsViSLyss=
 github.com/andybalholm/cascadia v1.3.2/go.mod h1:7gtRlve5FxPPgIgX36uWBX58OdBsSS6lUvCFb+h7KvU=

diff --git a/vectorstores/chroma/chroma.go b/vectorstores/chroma/chroma.go
@@ -59,7 +59,7 @@ func New(opts ...Option) (Store, error) {
 	chromaClient := &chromago.Client{
 		ApiClient: chromaopenapi.NewAPIClient(configuration),
 	}
-	if _, errHb := chromaClient.Heartbeat(); errHb != nil {
+	if _, errHb := chromaClient.Heartbeat(context.Background()); errHb != nil {
 		return s, errHb
 	}
 	s.client = chromaClient
@@ -73,7 +73,7 @@ func New(opts ...Option) (Store, error) {
 		embeddingFunction = openai.NewOpenAIEmbeddingFunction(s.openaiAPIKey)
 	}
 
-	col, errCc := s.client.CreateCollection(s.nameSpace, map[string]any{}, true,
+	col, errCc := s.client.CreateCollection(context.Background(), s.nameSpace, map[string]any{}, true,
 		embeddingFunction, s.distanceFunction)
 	if errCc != nil {
 		return s, fmt.Errorf("%w: %w", ErrNewClient, errCc)
@@ -85,7 +85,7 @@ func New(opts ...Option) (Store, error) {
 
 // AddDocuments adds the text and metadata from the documents to the Chroma collection associated with 'Store'.
 // and returns the ids of the added documents.
-func (s Store) AddDocuments(_ context.Context,
+func (s Store) AddDocuments(ctx context.Context,
 	docs []schema.Document,
 	options ...vectorstores.Option,
 ) ([]string, error) {
@@ -114,13 +114,13 @@ func (s Store) AddDocuments(_ context.Context,
 	}
 
 	col := s.collection
-	if _, addErr := col.Add(nil, metadatas, texts, ids); addErr != nil {
+	if _, addErr := col.Add(ctx, nil, metadatas, texts, ids); addErr != nil {
 		return nil, fmt.Errorf("%w: %w", ErrAddDocument, addErr)
 	}
 	return ids, nil
 }
 
-func (s Store) SimilaritySearch(_ context.Context, query string, numDocuments int,
+func (s Store) SimilaritySearch(ctx context.Context, query string, numDocuments int,
 	options ...vectorstores.Option,
 ) ([]schema.Document, error) {
 	opts := s.getOptions(options...)
@@ -136,7 +136,7 @@ func (s Store) SimilaritySearch(_ context.Context, query string, numDocuments in
 	}
 
 	filter := s.getNamespacedFilter(opts)
-	qr, queryErr := s.collection.Query([]string{query}, int32(numDocuments), filter, nil, s.includes)
+	qr, queryErr := s.collection.Query(ctx, []string{query}, int32(numDocuments), filter, nil, s.includes)
 	if queryErr != nil {
 		return nil, queryErr
 	}
@@ -165,7 +165,7 @@ func (s Store) RemoveCollection() error {
 	if s.client == nil || s.collection == nil {
 		return fmt.Errorf("%w: no collection", ErrRemoveCollection)
 	}
-	_, errDc := s.client.DeleteCollection(s.collection.Name)
+	_, errDc := s.client.DeleteCollection(context.Background(), s.collection.Name)
 	if errDc != nil {
 		return fmt.Errorf("%w(%s): %w", ErrRemoveCollection, s.collection.Name, errDc)
 	}