Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix OpenSearch vector store's doAdd(List<Document> documents)

Browse Source 
- 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
This commit is contained in:
Ilayaperumal Gopinathan
2024-12-06 18:08:50 +00:00
parent 3e9256fc00
commit 2ef4559827
1 changed files with 18 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java
View File

@@ -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<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> 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<String, Object> metadata, float[] embedding) {
}
}