From 2ef45598277534e1be31636c24d345ca80be4728 Mon Sep 17 00:00:00 2001 From: Ilayaperumal Gopinathan Date: Fri, 6 Dec 2024 18:08:50 +0000 Subject: [PATCH] Fix OpenSearch vector store's doAdd(List documents) - Since Document's reference to its embedding is deprecated, store the embedding into OpenSearch vector store by creating an explicit OpenSearch Document type which has embedding associated with it --- .../ai/vectorstore/OpenSearchVectorStore.java | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) 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) { + } + }