diff --git a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java index a5fc66fa4..0492ec80f 100644 --- a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java +++ b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java @@ -19,6 +19,7 @@ package org.springframework.ai.vectorstore; import java.io.IOException; import java.io.StringReader; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.stream.Collectors; @@ -51,7 +52,6 @@ import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; -import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.InitializingBean; import org.springframework.util.Assert; @@ -144,11 +144,14 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem @Override public void doAdd(List documents) { - this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + List embedding = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), + this.batchingStrategy); BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder(); for (Document document : documents) { - bulkRequestBuilder - .operations(op -> op.index(idx -> idx.index(this.index).id(document.getId()).document(document))); + OpenSearchDocument openSearchDocument = new OpenSearchDocument(document.getId(), document.getContent(), + document.getMetadata(), embedding.get(documents.indexOf(document))); + bulkRequestBuilder.operations(op -> op + .index(idx -> idx.index(this.index).id(openSearchDocument.id()).document(openSearchDocument))); } bulkRequest(bulkRequestBuilder.build()); } @@ -292,4 +295,15 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem return this.similarityFunction; } + /** + * 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 OpenSearchDocument(String id, String content, Map metadata, float[] embedding) { + } + }