Chroma: Replace existing documents with newer versions

vectorstores/chroma/chroma.go

@@ -8,7 +8,6 @@ import (
 	chromago "github.com/amikos-tech/chroma-go"
 	"github.com/amikos-tech/chroma-go/openai"
 	chromatypes "github.com/amikos-tech/chroma-go/types"
-	"github.com/google/uuid"
 	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/vectorstores"
@@ -106,7 +105,12 @@ func (s Store) AddDocuments(ctx context.Context,
 	texts := make([]string, len(docs))
 	metadatas := make([]map[string]any, len(docs))
 	for docIdx, doc := range docs {
-		ids[docIdx] = uuid.New().String() // TODO (noodnik2): find & use something more meaningful
+		if opts.Deduplicater != nil {
+			if opts.Deduplicater(ctx, doc) {
+				continue
+			}
+		}
+		ids[docIdx] = opts.GenerateDocumentID(ctx, doc, ids)
 		texts[docIdx] = doc.PageContent
 		mc := make(map[string]any, 0)
 		maps.Copy(mc, doc.Metadata)
@@ -115,9 +119,8 @@ func (s Store) AddDocuments(ctx context.Context,
 			metadatas[docIdx][s.nameSpaceKey] = nameSpace
 		}
 	}
-
 	col := s.collection
-	if _, addErr := col.Add(ctx, nil, metadatas, texts, ids); addErr != nil {
+	if _, addErr := col.Upsert(ctx, nil, metadatas, texts, ids); addErr != nil {
 		return nil, fmt.Errorf("%w: %w", ErrAddDocument, addErr)
 	}
 	return ids, nil

vectorstores/opensearch/opensearch.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"io"
 
-	"github.com/google/uuid"
 	opensearchgo "github.com/opensearch-project/opensearch-go"
 	"github.com/opensearch-project/opensearch-go/opensearchapi"
 	"github.com/tmc/langchaingo/embeddings"
@@ -75,7 +74,7 @@ func (s Store) AddDocuments(
 	}
 
 	for i, doc := range docs {
-		id := uuid.NewString()
+		id := opts.GenerateDocumentID(ctx, doc, ids)
 		_, err := s.documentIndexing(ctx, id, opts.NameSpace, doc.PageContent, vectors[i], doc.Metadata)
 		if err != nil {
 			return ids, err

vectorstores/options.go

@@ -2,7 +2,10 @@ package vectorstores
 
 import (
 	"context"
+	"log/slog"
+	"slices"
 
+	"github.com/google/uuid"
 	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/schema"
 )
@@ -12,11 +15,12 @@ type Option func(*Options)
 
 // Options is a set of options for similarity search and add documents.
 type Options struct {
-	NameSpace      string
-	ScoreThreshold float32
-	Filters        any
-	Embedder       embeddings.Embedder
-	Deduplicater   func(context.Context, schema.Document) bool
+	NameSpace           string
+	ScoreThreshold      float32
+	Filters             any
+	Embedder            embeddings.Embedder
+	Deduplicater        func(context.Context, schema.Document) bool
+	DocumentIDGenerater func(context.Context, schema.Document) string
 }
 
 // WithNameSpace returns an Option for setting the name space.
@@ -59,3 +63,32 @@ func WithDeduplicater(fn func(ctx context.Context, doc schema.Document) bool) Op
 		o.Deduplicater = fn
 	}
 }
+
+// WithIDGenerater returns an Option for setting to generate the IDS
+func WithIDGenerater(fn func(ctx context.Context, doc schema.Document) string) Option {
+	return func(o *Options) {
+		o.DocumentIDGenerater = fn
+	}
+}
+
+// generateDummyDoumentID generates a UUID
+func (o Options) generateDummyDoumentID(_ context.Context) string {
+	return uuid.NewString()
+}
+
+// GenerateDocumentID calls the provided ID generator or creates a new UUID if not provided or the generated ID is not unique
+func (o Options) GenerateDocumentID(ctx context.Context, doc schema.Document, ids []string) string {
+	if o.DocumentIDGenerater == nil {
+		return o.generateDummyDoumentID(ctx)
+	}
+	id := o.DocumentIDGenerater(ctx, doc)
+	if len(id) < 1 {
+		slog.Warn("Document ID generator did not generate an id", "id", id)
+		return o.generateDummyDoumentID(ctx)
+	}
+	if slices.Contains(ids, id) {
+		slog.Warn("Document ID generator generated a non unique id", "id", id)
+		return o.generateDummyDoumentID(ctx)
+	}
+	return id
+}

vectorstores/pgvector/pgvector.go

@@ -261,7 +261,7 @@ func (s Store) AddDocuments(
 
 	ids := make([]string, len(docs))
 	for docIdx, doc := range docs {
-		id := uuid.New().String()
+		id := opts.GenerateDocumentID(ctx, doc, ids)
 		ids[docIdx] = id
 		b.Queue(sql, id, doc.PageContent, pgvector.NewVector(vectors[docIdx]), doc.Metadata, s.collectionUUID)
 	}

vectorstores/pinecone/pinecone.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 
-	"github.com/google/uuid"
 	"github.com/pinecone-io/go-pinecone/pinecone"
 	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/schema"
@@ -104,7 +103,7 @@ func (s Store) AddDocuments(ctx context.Context,
 			return nil, err
 		}
 
-		id := uuid.New().String()
+		id := opts.GenerateDocumentID(ctx, docs[i], ids)
 		ids[i] = id
 		pineconeVectors = append(
 			pineconeVectors,

vectorstores/weaviate/weaviate.go

@@ -8,7 +8,6 @@ import (
 	"strings"
 
 	"github.com/go-openapi/strfmt"
-	"github.com/google/uuid"
 	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/vectorstores"
@@ -133,7 +132,7 @@ func (s Store) AddDocuments(ctx context.Context,
 	objects := make([]*models.Object, 0, len(docs))
 	ids := make([]string, len(docs))
 	for i := range docs {
-		id := strfmt.UUID(uuid.New().String())
+		id := strfmt.UUID(opts.GenerateDocumentID(ctx, docs[i], ids))
 		ids[i] = id.String()
 		objects = append(objects, &models.Object{
 			Class:      s.indexName,