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

Author: ilayaperumalg

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java

@@ -66,12 +66,12 @@ void defaultEmbedding() {
 	@Test
 	void embeddingBatchDocuments() throws Exception {
 		assertThat(this.embeddingModel).isNotNull();
-		List<float[]> embedded = this.embeddingModel.embed(
+		List<float[]> 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

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.
 	 */

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java

@@ -78,25 +78,23 @@ default List<float[]> embed(List<String> texts) {
 	 * @param options {@link EmbeddingOptions}.
 	 * @param batchingStrategy {@link BatchingStrategy}.
 	 * @return a list of float[] that represents the vectors for the incoming
-	 * {@link Document}s.
+	 * {@link Document}s. The returned list is expected to be in the same order of the
+	 * {@link Document} list.
 	 */
 	default List<float[]> embed(List<Document> documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
 		Assert.notNull(documents, "Documents must not be null");
-		List<float[]> embeddings = new ArrayList<>();
-
+		List<float[]> embeddings = new ArrayList<>(documents.size());
 		List<List<Document>> batch = batchingStrategy.batch(documents);
-
 		for (List<Document> subBatch : batch) {
 			List<String> 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;
 	}
 

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 List<List<Document>> batch(List<Document> documents) {
 		List<List<Document>> batches = new ArrayList<>();
 		int currentSize = 0;
 		List<Document> currentBatch = new ArrayList<>();
-		Map<Document, Integer> documentTokens = new HashMap<>();
+		// Make sure the documentTokens' entry order is preserved by making it a
+		// LinkedHashMap.
+		Map<Document, Integer> documentTokens = new LinkedHashMap<>();
 
 		for (Document document : documents) {
 			int tokenCount = this.tokenCountEstimator

vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java

@@ -203,13 +203,14 @@ private JsonNode mapCosmosDocument(Document document, float[] queryEmbedding) {
 	public void doAdd(List<Document> documents) {
 
 		// Batch the documents based on the batching strategy
-		this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
+		List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
+				this.batchingStrategy);
 
 		// Create a list to hold both the CosmosItemOperation and the corresponding
 		// document ID
 		List<ImmutablePair<String, CosmosItemOperation>> 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();

vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java

@@ -224,12 +224,13 @@ public void doAdd(List<Document> documents) {
 			return; // nothing to do;
 		}
 
-		this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
+		List<float[]> 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());
 
@@ -324,7 +325,6 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
 				metadata.put(DISTANCE_METADATA_FIELD_NAME, 1 - (float) result.getScore());
 
 				final Document doc = new Document(entry.id(), entry.content(), metadata);
-				doc.setEmbedding(EmbeddingUtils.toPrimitive(entry.embedding()));
 
 				return doc;
 

vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java

@@ -182,7 +182,8 @@ private static Float[] toFloatArray(float[] embedding) {
 	public void doAdd(List<Document> documents) {
 		var futures = new CompletableFuture[documents.size()];
 
-		this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
+		List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
+				this.batchingStrategy);
 
 		int i = 0;
 		for (Document d : documents) {
@@ -197,7 +198,8 @@ public void doAdd(List<Document> documents) {
 
 				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()
@@ -260,11 +262,6 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
 				}
 			}
 			Document doc = new Document(getDocumentId(row), row.getString(this.conf.schema.content()), docFields);
-
-			if (this.conf.returnEmbeddings) {
-				doc.setEmbedding(EmbeddingUtils
-					.toPrimitive(row.getVector(this.conf.schema.embedding(), Float.class).stream().toList()));
-			}
 			documents.add(doc);
 		}
 		return documents;

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;

vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java

@@ -121,18 +121,12 @@ void addAndSearch() {
 
 				List<Document> documents = documents();
 				store.add(documents);
-				for (Document d : documents) {
-					assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(),
-							e -> assertThat(e).isNotEmpty());
-				}
 
 				List<Document> 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.");
@@ -158,17 +152,12 @@ void addAndSearchReturnEmbeddings() {
 			try (CassandraVectorStore store = createTestStore(context, builder)) {
 				List<Document> documents = documents();
 				store.add(documents);
-				for (Document d : documents) {
-					assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(),
-							e -> assertThat(e).isNotEmpty());
-				}
 
 				List<Document> 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.");

vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java

@@ -145,14 +145,14 @@ public void doAdd(@NonNull List<Document> documents) {
 		List<String> contents = new ArrayList<>();
 		List<float[]> embeddings = new ArrayList<>();
 
-		this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
+		List<float[]> 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,
@@ -193,9 +193,7 @@ public Optional<Boolean> doDelete(@NonNull List<String> idList) {
 					metadata = new HashMap<>();
 				}
 				metadata.put(DISTANCE_FIELD_NAME, distance);
-				Document document = new Document(id, content, metadata);
-				document.setEmbedding(chromaEmbedding.embedding());
-				responseDocuments.add(document);
+				responseDocuments.add(new Document(id, content, metadata));
 			}
 		}