From ebd29e0959086ea950d505a8aefb9e32ca7732a7 Mon Sep 17 00:00:00 2001 From: Ilayaperumal Gopinathan Date: Mon, 25 Nov 2024 22:21:11 +0000 Subject: [PATCH] GH-1826 Fix EmbeddingModel's usage on Document#embedding - Since the Document object's reference to the `embedding` is deprecated and will be removed, the VectorStore implementations require a way to store the embedding of the corresponding Document objects - One way to fix this is, to have the EmbeddingModel#embed to return the embeddings in the same order as that of the Documents passed to it. - Since both the Document and embedding collections use the List object, their iteration operation will make sure to keep them in line with the same order. - A fix is required to preserve the order when batching strategy is applied. - Updated the Javadoc for BatchingStrategy - Fixed the Document List order in TokenCountBatchingStrategy - Refactored the vector store implementations to update this change Resolves #GH-1826 --- .../ai/openai/embedding/EmbeddingIT.java | 6 ++--- .../ai/embedding/BatchingStrategy.java | 4 +++- .../ai/embedding/EmbeddingModel.java | 14 +++++------ .../embedding/TokenCountBatchingStrategy.java | 6 +++-- .../ai/vectorstore/CosmosDBVectorStore.java | 7 +++--- .../vectorstore/azure/AzureVectorStore.java | 6 ++--- .../ai/vectorstore/CassandraVectorStore.java | 10 ++++---- .../CassandraVectorStoreConfig.java | 2 ++ .../vectorstore/CassandraVectorStoreIT.java | 11 --------- .../ai/vectorstore/ChromaVectorStore.java | 8 +++---- .../ai/vectorstore/CoherenceVectorStore.java | 4 ++-- .../vectorstore/ElasticsearchVectorStore.java | 21 +++++++++++++--- .../ai/vectorstore/GemFireVectorStore.java | 7 +++--- .../ai/vectorstore/MilvusVectorStore.java | 5 ++-- .../vectorstore/MongoDBAtlasVectorStore.java | 24 +++++++++++++++---- .../ai/vectorstore/Neo4jVectorStore.java | 12 ++++++---- .../ai/vectorstore/OracleVectorStore.java | 7 +++--- .../ai/vectorstore/PgVectorStore.java | 18 ++++++-------- .../PgVectorStoreWithChatMemoryAdvisorIT.java | 3 --- .../ai/vectorstore/PineconeVectorStore.java | 5 ++-- .../vectorstore/qdrant/QdrantVectorStore.java | 5 ++-- .../ai/vectorstore/RedisVectorStore.java | 6 ++--- .../ai/vectorstore/TypesenseVectorStore.java | 5 ++-- .../ai/vectorstore/WeaviateVectorStore.java | 16 +++++++------ 24 files changed, 118 insertions(+), 94 deletions(-) diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java index 7035405c0..52ace2598 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java @@ -66,12 +66,12 @@ class EmbeddingIT extends AbstractIT { @Test void embeddingBatchDocuments() throws Exception { assertThat(this.embeddingModel).isNotNull(); - List embedded = this.embeddingModel.embed( + List embeddings = this.embeddingModel.embed( List.of(new Document("Hello world"), new Document("Hello Spring"), new Document("Hello Spring AI!")), OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(), new TokenCountBatchingStrategy()); - assertThat(embedded.size()).isEqualTo(3); - embedded.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions())); + assertThat(embeddings.size()).isEqualTo(3); + embeddings.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions())); } @Test diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java index e354f1da8..714d681a5 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java @@ -31,7 +31,9 @@ public interface BatchingStrategy { /** * {@link EmbeddingModel} implementations can call this method to optimize embedding - * tokens. The incoming collection of {@link Document}s are split into su-batches. + * tokens. The incoming collection of {@link Document}s are split into sub-batches. It + * is important to preserve the order of the list of {@link Document}s when batching + * as they are mapped to their corresponding embeddings by their order. * @param documents to batch * @return a list of sub-batches that contain {@link Document}s. */ diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java index c40ed34d2..51a1ac035 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java @@ -78,25 +78,23 @@ public interface EmbeddingModel extends Model embed(List documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) { Assert.notNull(documents, "Documents must not be null"); - List embeddings = new ArrayList<>(); - + List embeddings = new ArrayList<>(documents.size()); List> batch = batchingStrategy.batch(documents); - for (List subBatch : batch) { List texts = subBatch.stream().map(Document::getContent).toList(); EmbeddingRequest request = new EmbeddingRequest(texts, options); EmbeddingResponse response = this.call(request); for (int i = 0; i < subBatch.size(); i++) { - Document document = subBatch.get(i); - float[] output = response.getResults().get(i).getOutput(); - embeddings.add(output); - document.setEmbedding(output); + embeddings.add(response.getResults().get(i).getOutput()); } } + Assert.isTrue(embeddings.size() == documents.size(), + "Embeddings must have the same number as that of the documents"); return embeddings; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java index 713eefb96..7ffc01fee 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java @@ -17,7 +17,7 @@ package org.springframework.ai.embedding; import java.util.ArrayList; -import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -139,7 +139,9 @@ public class TokenCountBatchingStrategy implements BatchingStrategy { List> batches = new ArrayList<>(); int currentSize = 0; List currentBatch = new ArrayList<>(); - Map documentTokens = new HashMap<>(); + // Make sure the documentTokens' entry order is preserved by making it a + // LinkedHashMap. + Map documentTokens = new LinkedHashMap<>(); for (Document document : documents) { int tokenCount = this.tokenCountEstimator diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java index 11dc70067..eb3fba4b3 100644 --- a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java @@ -204,13 +204,14 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen public void doAdd(List documents) { // Batch the documents based on the batching strategy - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); // Create a list to hold both the CosmosItemOperation and the corresponding // document ID List> itemOperationsWithIds = documents.stream().map(doc -> { - CosmosItemOperation operation = CosmosBulkOperations - .getCreateItemOperation(mapCosmosDocument(doc, doc.getEmbedding()), new PartitionKey(doc.getId())); + CosmosItemOperation operation = CosmosBulkOperations.getCreateItemOperation( + mapCosmosDocument(doc, embeddings.get(documents.indexOf(doc))), new PartitionKey(doc.getId())); return new ImmutablePair<>(doc.getId(), operation); // Pair the document ID // with the operation }).toList(); diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java index f1893bd20..788352d51 100644 --- a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java +++ b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java @@ -223,12 +223,13 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements return; // nothing to do; } - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); final var searchDocuments = documents.stream().map(document -> { SearchDocument searchDocument = new SearchDocument(); searchDocument.put(ID_FIELD_NAME, document.getId()); - searchDocument.put(EMBEDDING_FIELD_NAME, document.getEmbedding()); + searchDocument.put(EMBEDDING_FIELD_NAME, embeddings.get(documents.indexOf(document))); searchDocument.put(CONTENT_FIELD_NAME, document.getContent()); searchDocument.put(METADATA_FIELD_NAME, new JSONObject(document.getMetadata()).toJSONString()); @@ -327,7 +328,6 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements .content(entry.content) .metadata(metadata) .score(result.getScore()) - .embedding(EmbeddingUtils.toPrimitive(entry.embedding)) .build(); }) .collect(Collectors.toList()); diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java index af105efc0..df8467fa7 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -181,7 +181,8 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme public void doAdd(List documents) { var futures = new CompletableFuture[documents.size()]; - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); int i = 0; for (Document d : documents) { @@ -196,7 +197,8 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme builder = builder.setString(this.conf.schema.content(), d.getContent()) .setVector(this.conf.schema.embedding(), - CqlVector.newInstance(EmbeddingUtils.toList(d.getEmbedding())), Float.class); + CqlVector.newInstance(EmbeddingUtils.toList(embeddings.get(documents.indexOf(d)))), + Float.class); for (var metadataColumn : this.conf.schema.metadataColumns() .stream() @@ -265,10 +267,6 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme .score((double) score) .build(); - if (this.conf.returnEmbeddings) { - doc.setEmbedding(EmbeddingUtils - .toPrimitive(row.getVector(this.conf.schema.embedding(), Float.class).stream().toList())); - } documents.add(doc); } return documents; diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java index 007d2c08c..518eb622a 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java @@ -90,6 +90,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { final boolean disallowSchemaChanges; + // TODO: Remove this flag as the document no longer holds embeddings. + @Deprecated(since = "1.0.0-M5", forRemoval = true) final boolean returnEmbeddings; final DocumentIdTranslator documentIdTranslator; diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java index 13f7e7d72..d50414bb0 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java @@ -122,18 +122,12 @@ class CassandraVectorStoreIT { List documents = documents(); store.add(documents); - for (Document d : documents) { - assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(), - e -> assertThat(e).isNotEmpty()); - } List results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); assertThat(results).hasSize(1); Document resultDoc = results.get(0); assertThat(resultDoc.getId()).isEqualTo(documents().get(0).getId()); - assertThat(resultDoc.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNull(), - e -> assertThat(e).isEmpty()); assertThat(resultDoc.getContent()).contains( "Spring AI provides abstractions that serve as the foundation for developing AI applications."); @@ -159,17 +153,12 @@ class CassandraVectorStoreIT { try (CassandraVectorStore store = createTestStore(context, builder)) { List documents = documents(); store.add(documents); - for (Document d : documents) { - assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(), - e -> assertThat(e).isNotEmpty()); - } List results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); assertThat(results).hasSize(1); Document resultDoc = results.get(0); assertThat(resultDoc.getId()).isEqualTo(documents().get(0).getId()); - assertThat(resultDoc.getEmbedding()).isNotEmpty(); assertThat(resultDoc.getContent()).contains( "Spring AI provides abstractions that serve as the foundation for developing AI applications."); diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java index eafe1e37a..75e9c4808 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java @@ -145,14 +145,14 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements List contents = new ArrayList<>(); List embeddings = new ArrayList<>(); - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List documentEmbeddings = this.embeddingModel.embed(documents, + EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); for (Document document : documents) { ids.add(document.getId()); metadatas.add(document.getMetadata()); contents.add(document.getContent()); - document.setEmbedding(document.getEmbedding()); - embeddings.add(document.getEmbedding()); + embeddings.add(documentEmbeddings.get(documents.indexOf(document))); } this.chromaApi.upsertEmbeddings(this.collectionId, @@ -192,12 +192,12 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements if (metadata == null) { metadata = new HashMap<>(); } + metadata.put(DocumentMetadata.DISTANCE.value(), distance); Document document = Document.builder() .id(id) .content(content) .metadata(metadata) - .embedding(chromaEmbedding.embedding()) .score(1.0 - distance) .build(); responseDocuments.add(document); diff --git a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/CoherenceVectorStore.java b/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/CoherenceVectorStore.java index 29ee2bfdf..3bd9ddc4f 100644 --- a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/CoherenceVectorStore.java +++ b/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/CoherenceVectorStore.java @@ -167,9 +167,9 @@ public class CoherenceVectorStore implements VectorStore, InitializingBean { public void add(final List documents) { Map chunks = new HashMap<>((int) Math.ceil(documents.size() / 0.75f)); for (Document doc : documents) { - doc.setEmbedding(this.embeddingModel.embed(doc)); var id = toChunkId(doc.getId()); - var chunk = new DocumentChunk(doc.getContent(), doc.getMetadata(), toFloat32Vector(doc.getEmbedding())); + var chunk = new DocumentChunk(doc.getContent(), doc.getMetadata(), + toFloat32Vector(this.embeddingModel.embed(doc))); chunks.put(id, chunk); } this.documentChunks.putAll(chunks); diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java index ee26f4b54..f96feb9a7 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java @@ -72,6 +72,7 @@ import org.springframework.util.Assert; * @author Soby Chacko * @author Christian Tzolov * @author Thomas Vitale + * @author Ilayaperumal Gopinathan * @since 1.0.0 */ public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean { @@ -133,11 +134,14 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp } BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder(); - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); for (Document document : documents) { - bulkRequestBuilder.operations(op -> op - .index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(document))); + ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getContent(), + document.getMetadata(), embeddings.get(documents.indexOf(document))); + bulkRequestBuilder.operations( + op -> op.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(doc))); } BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build()); if (bulkRequest.errors()) { @@ -282,4 +286,15 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp return SIMILARITY_TYPE_MAPPING.get(this.options.getSimilarity()).value(); } + /** + * The representation of {@link Document} along with its embedding. + * + * @param id The id of the document + * @param content The content of the document + * @param metadata The metadata of the document + * @param embedding The vectors representing the content of the document + */ + public record ElasticSearchDocument(String id, String content, Map metadata, float[] embedding) { + } + } diff --git a/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java b/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java index 24e47aedf..494794ea6 100644 --- a/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java +++ b/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java @@ -208,10 +208,11 @@ public class GemFireVectorStore extends AbstractObservationVectorStore implement @Override public void doAdd(List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); UploadRequest upload = new UploadRequest(documents.stream() - .map(document -> new UploadRequest.Embedding(document.getId(), document.getEmbedding(), DOCUMENT_FIELD, - document.getContent(), document.getMetadata())) + .map(document -> new UploadRequest.Embedding(document.getId(), embeddings.get(documents.indexOf(document)), + DOCUMENT_FIELD, document.getContent(), document.getMetadata())) .toList()); String embeddingsJson = null; diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java index 18c5f26ea..e586797dc 100644 --- a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java @@ -161,7 +161,8 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements List> embeddingArray = new ArrayList<>(); // TODO: Need to customize how we pass the embedding options - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); for (Document document : documents) { docIdArray.add(document.getId()); @@ -169,7 +170,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements // the content used to compute the embeddings contentArray.add(document.getContent()); metadataArray.add(new JSONObject(document.getMetadata())); - embeddingArray.add(EmbeddingUtils.toList(document.getEmbedding())); + embeddingArray.add(EmbeddingUtils.toList(embeddings.get(documents.indexOf(document)))); } List fields = new ArrayList<>(); diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java index 3ab60f4cc..1ea36cb13 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java @@ -51,6 +51,7 @@ import org.springframework.util.Assert; * @author Soby Chacko * @author Christian Tzolov * @author Thomas Vitale + * @author Ilayaperumal Gopinathan * @since 1.0.0 */ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore implements InitializingBean { @@ -175,22 +176,26 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl String content = mongoDocument.getString(CONTENT_FIELD_NAME); double score = mongoDocument.getDouble(SCORE_FIELD_NAME); Map metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class); + metadata.put(DocumentMetadata.DISTANCE.value(), 1 - score); + // @formatter:off return Document.builder() .id(id) .content(content) .metadata(metadata) .score(score) - .embedding(queryEmbedding) - .build(); + .build(); // @formatter:on } @Override public void doAdd(List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); for (Document document : documents) { - this.mongoTemplate.save(document, this.config.collectionName); + MongoDBDocument mdbDocument = new MongoDBDocument(document.getId(), document.getContent(), + document.getMetadata(), embeddings.get(documents.indexOf(document))); + this.mongoTemplate.save(mdbDocument, this.config.collectionName); } } @@ -338,4 +343,15 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl } + /** + * The representation of {@link Document} along with its embedding. + * + * @param id The id of the document + * @param content The content of the document + * @param metadata The metadata of the document + * @param embedding The vectors representing the content of the document + */ + public record MongoDBDocument(String id, String content, Map metadata, float[] embedding) { + } + } diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java index 97a74a175..1a64f952d 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java @@ -108,9 +108,12 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements @Override public void doAdd(List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); - var rows = documents.stream().map(this::documentToRecord).toList(); + var rows = documents.stream() + .map(document -> documentToRecord(document, embeddings.get(documents.indexOf(document)))) + .toList(); try (var session = this.driver.session()) { var statement = """ @@ -204,8 +207,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements } } - private Map documentToRecord(Document document) { - document.setEmbedding(document.getEmbedding()); + private Map documentToRecord(Document document, float[] embedding) { var row = new HashMap(); @@ -217,7 +219,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v))); row.put("properties", properties); - row.put(this.config.embeddingProperty, Values.value(document.getEmbedding())); + row.put(this.config.embeddingProperty, Values.value(embedding)); return row; } diff --git a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java index f519024dc..53bb84737 100644 --- a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java +++ b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java @@ -206,7 +206,8 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements @Override public void doAdd(final List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); this.jdbcTemplate.batchUpdate(getIngestStatement(), new BatchPreparedStatementSetter() { @Override @@ -214,7 +215,7 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements final Document document = documents.get(i); final String content = document.getContent(); final byte[] json = toJson(document.getMetadata()); - final VECTOR embeddingVector = toVECTOR(document.getEmbedding()); + final VECTOR embeddingVector = toVECTOR(embeddings.get(documents.indexOf(document))); org.springframework.jdbc.core.StatementCreatorUtils.setParameterValue(ps, 1, Types.VARCHAR, document.getId()); @@ -653,13 +654,11 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements final Map metadata = getMap(rs.getObject(3, OracleJsonValue.class)); metadata.put(DocumentMetadata.DISTANCE.value(), rs.getDouble(5)); - final float[] embedding = rs.getObject(4, float[].class); return Document.builder() .id(rs.getString(1)) .content(rs.getString(2)) .metadata(metadata) .score(1 - rs.getDouble(5)) - .embedding(embedding) .build(); } diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index f74e3fe94..a8a057c8d 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -196,10 +196,11 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini @Override public void doAdd(List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); List> batchedDocuments = batchDocuments(documents); - batchedDocuments.forEach(this::insertOrUpdateBatch); + batchedDocuments.forEach(batchDocument -> insertOrUpdateBatch(batchDocument, documents, embeddings)); } private List> batchDocuments(List documents) { @@ -210,7 +211,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini return batches; } - private void insertOrUpdateBatch(List batch) { + private void insertOrUpdateBatch(List batch, List documents, List embeddings) { String sql = "INSERT INTO " + getFullyQualifiedTableName() + " (id, content, metadata, embedding) VALUES (?, ?, ?::jsonb, ?) " + "ON CONFLICT (id) DO " + "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? "; @@ -223,7 +224,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini var document = batch.get(i); var content = document.getContent(); var json = toJson(document.getMetadata()); - var embedding = document.getEmbedding(); + var embedding = embeddings.get(documents.indexOf(document)); var pGvector = new PGvector(embedding); StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, @@ -499,23 +500,18 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini String id = rs.getString(COLUMN_ID); String content = rs.getString(COLUMN_CONTENT); PGobject pgMetadata = rs.getObject(COLUMN_METADATA, PGobject.class); - PGobject embedding = rs.getObject(COLUMN_EMBEDDING, PGobject.class); Float distance = rs.getFloat(COLUMN_DISTANCE); Map metadata = toMap(pgMetadata); metadata.put(DocumentMetadata.DISTANCE.value(), distance); + // @formatter:off return Document.builder() .id(id) .content(content) .metadata(metadata) .score(1.0 - distance) - .embedding(toFloatArray(embedding)) - .build(); - } - - private float[] toFloatArray(PGobject embedding) throws SQLException { - return new PGvector(embedding.getValue()).toArray(); + .build(); // @formatter:on } private Map toMap(PGobject pgObject) { diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java index e2a37cd37..6325ac718 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java @@ -146,9 +146,6 @@ class PgVectorStoreWithChatMemoryAdvisorIT { EmbeddingModel embeddingModel = mock(EmbeddingModel.class); Mockito.doAnswer(invocationOnMock -> { - Object[] arguments = invocationOnMock.getArguments(); - List documents = (List) arguments[0]; - documents.forEach(d -> d.setEmbedding(this.embed)); return List.of(this.embed, this.embed); }).when(embeddingModel).embed(ArgumentMatchers.any(), any(), any()); given(embeddingModel.embed(any(String.class))).willReturn(this.embed); diff --git a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java index d9ff49721..b973ae448 100644 --- a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java +++ b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java @@ -124,11 +124,12 @@ public class PineconeVectorStore extends AbstractObservationVectorStore { * @param namespace The namespace to add the documents to */ public void add(List documents, String namespace) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); List upsertVectors = documents.stream() .map(document -> Vector.newBuilder() .setId(document.getId()) - .addAllValues(EmbeddingUtils.toList(document.getEmbedding())) + .addAllValues(EmbeddingUtils.toList(embeddings.get(documents.indexOf(document)))) .setMetadata(metadataToStruct(document)) .build()) .toList(); diff --git a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java index bb4f6b392..98c3a5c8b 100644 --- a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java +++ b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java @@ -127,12 +127,13 @@ public class QdrantVectorStore extends AbstractObservationVectorStore implements try { // Compute and assign an embedding to the document. - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); List points = documents.stream() .map(document -> PointStruct.newBuilder() .setId(io.qdrant.client.PointIdFactory.id(UUID.fromString(document.getId()))) - .setVectors(io.qdrant.client.VectorsFactory.vectors(document.getEmbedding())) + .setVectors(io.qdrant.client.VectorsFactory.vectors(embeddings.get(documents.indexOf(document)))) .putAllPayload(toPayload(document)) .build()) .toList(); diff --git a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java index 875943792..1397585d3 100644 --- a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java +++ b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java @@ -163,12 +163,12 @@ public class RedisVectorStore extends AbstractObservationVectorStore implements public void doAdd(List documents) { try (Pipeline pipeline = this.jedis.pipelined()) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); for (Document document : documents) { - document.setEmbedding(document.getEmbedding()); var fields = new HashMap(); - fields.put(this.config.embeddingFieldName, document.getEmbedding()); + fields.put(this.config.embeddingFieldName, embeddings.get(documents.indexOf(document))); fields.put(this.config.contentFieldName, document.getContent()); fields.putAll(document.getMetadata()); pipeline.jsonSetWithEscape(key(document.getId()), JSON_SET_PATH, fields); diff --git a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java index b805be621..514364f2d 100644 --- a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java +++ b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java @@ -125,14 +125,15 @@ public class TypesenseVectorStore extends AbstractObservationVectorStore impleme public void doAdd(List documents) { Assert.notNull(documents, "Documents must not be null"); - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); List> documentList = documents.stream().map(document -> { HashMap typesenseDoc = new HashMap<>(); typesenseDoc.put(DOC_ID_FIELD_NAME, document.getId()); typesenseDoc.put(CONTENT_FIELD_NAME, document.getContent()); typesenseDoc.put(METADATA_FIELD_NAME, document.getMetadata()); - typesenseDoc.put(EMBEDDING_FIELD_NAME, document.getEmbedding()); + typesenseDoc.put(EMBEDDING_FIELD_NAME, embeddings.get(documents.indexOf(document))); return typesenseDoc; }).toList(); diff --git a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java index 049d4a4f1..00a4a3520 100644 --- a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java +++ b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java @@ -197,9 +197,12 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore { return; } - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); - List weaviateObjects = documents.stream().map(this::toWeaviateObject).toList(); + List weaviateObjects = documents.stream() + .map(document -> toWeaviateObject(document, documents, embeddings)) + .toList(); Result response = this.weaviateClient.batch() .objectsBatcher() @@ -235,7 +238,7 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore { } } - private WeaviateObject toWeaviateObject(Document document) { + private WeaviateObject toWeaviateObject(Document document, List documents, List embeddings) { // https://weaviate.io/developers/weaviate/config-refs/datatypes Map fields = new HashMap<>(); @@ -259,7 +262,7 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore { return WeaviateObject.builder() .className(this.weaviateObjectClass) .id(document.getId()) - .vector(EmbeddingUtils.toFloatArray(document.getEmbedding())) + .vector(EmbeddingUtils.toFloatArray(embeddings.get(documents.indexOf(document)))) .properties(fields) .build(); } @@ -363,7 +366,6 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore { Map additional = (Map) item.get(ADDITIONAL_FIELD_NAME); double certainty = (Double) additional.get(ADDITIONAL_CERTAINTY_FIELD_NAME); String id = (String) additional.get(ADDITIONAL_ID_FIELD_NAME); - List embedding = ((List) additional.get(ADDITIONAL_VECTOR_FIELD_NAME)).stream().toList(); // Metadata Map metadata = new HashMap<>(); @@ -382,13 +384,13 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore { // Content String content = (String) item.get(CONTENT_FIELD_NAME); + // @formatter:off return Document.builder() .id(id) .content(content) .metadata(metadata) - .embedding(EmbeddingUtils.toPrimitive(EmbeddingUtils.doubleToFloat(embedding))) .score(certainty) - .build(); + .build(); // @formatter:on } @Override