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

Include index name in OpenSearch similarity search

Browse Source 
- Resolved issue where index name was not being sent during
  similaritySearch
- Updated similaritySearch method to include index in the SearchRequest
- Implemented test to verify documents can be added and retrieved from
  two different indices using separate OpenSearchVectorStore instances
- Ensured similarity search results are correctly returned for the
  respective indices

Fixes #885
This commit is contained in:
inpink
2024-09-29 15:56:14 +09:00
committed by Mark Pollack
parent 8bea302b11
commit 5431bdd956
2 changed files with 64 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -60,6 +60,7 @@ import java.util.stream.Collectors;
* @author Soby Chacko
* @author Christian Tzolov
* @author Thomas Vitale
* @author inpink
* @since 1.0.0
*/
public class OpenSearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -178,6 +179,7 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
Filter.Expression filterExpression) {
return similaritySearch(new org.opensearch.client.opensearch.core.SearchRequest.Builder()
.query(getOpenSearchSimilarityQuery(embedding, filterExpression))
.index(this.index)
.sort(sortOptionsBuilder -> sortOptionsBuilder
.score(scoreSortBuilder -> scoreSortBuilder.order(SortOrder.Desc)))
.size(topK)

67
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/OpenSearchVectorStoreIT.java
View File

@@ -20,6 +20,7 @@ import org.apache.hc.core5.http.HttpHost;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
@@ -30,6 +31,7 @@ import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -55,9 +57,15 @@ import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasSize;
/**
* The OpenSearchVectorStoreIT class is a test class designed to validate the
* functionality of a vector store that integrates with OpenSearch. It contains multiple
* parameterized tests to ensure the correctness of storing, searching, and updating
* vectorized documents in OpenSearch.
*
* @author Jemin Huh
* @author Soby Chacko
* @author Thomas Vitale
* @author inpink
* @since 1.0.0
*/
@Testcontainers
@@ -99,8 +107,11 @@ class OpenSearchVectorStoreIT {
@BeforeEach
void cleanDatabase() {
getContextRunner().run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
vectorStore.delete(List.of("_all"));
VectorStore anotherVectorStore = context.getBean("anotherVectorStore", OpenSearchVectorStore.class);
anotherVectorStore.delete(List.of("_all"));
});
}
@@ -109,7 +120,7 @@ class OpenSearchVectorStoreIT {
public void addAndSearchTest(String similarityFunction) {
getContextRunner().run(context -> {
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
if (!DEFAULT.equals(similarityFunction)) {
vectorStore.withSimilarityFunction(similarityFunction);
@@ -148,7 +159,7 @@ class OpenSearchVectorStoreIT {
public void searchWithFilters(String similarityFunction) {
getContextRunner().run(context -> {
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
if (!DEFAULT.equals(similarityFunction)) {
vectorStore.withSimilarityFunction(similarityFunction);
@@ -246,7 +257,7 @@ class OpenSearchVectorStoreIT {
public void documentUpdateTest(String similarityFunction) {
getContextRunner().run(context -> {
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
if (!DEFAULT.equals(similarityFunction)) {
vectorStore.withSimilarityFunction(similarityFunction);
}
@@ -302,7 +313,7 @@ class OpenSearchVectorStoreIT {
public void searchThresholdTest(String similarityFunction) {
getContextRunner().run(context -> {
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
if (!DEFAULT.equals(similarityFunction)) {
vectorStore.withSimilarityFunction(similarityFunction);
}
@@ -343,11 +354,41 @@ class OpenSearchVectorStoreIT {
});
}
@Test
public void searchDocumentsInTwoIndicesTest() {
getContextRunner().run(context -> {
// given
OpenSearchVectorStore vectorStore1 = context.getBean("vectorStore", OpenSearchVectorStore.class);
OpenSearchVectorStore vectorStore2 = context.getBean("anotherVectorStore", OpenSearchVectorStore.class);
Document docInIndex1 = new Document("1", "Document in index 1", Map.of("meta", "index1"));
Document docInIndex2 = new Document("2", "Document in index 2", Map.of("meta", "index2"));
// when
vectorStore1.add(List.of(docInIndex1));
vectorStore2.add(List.of(docInIndex2));
List<Document> resultInIndex1 = vectorStore1
.similaritySearch(SearchRequest.query("Document in index 1").withTopK(1).withSimilarityThreshold(0));
List<Document> resultInIndex2 = vectorStore2
.similaritySearch(SearchRequest.query("Document in index 2").withTopK(1).withSimilarityThreshold(0));
// then
assertThat(resultInIndex1).hasSize(1);
assertThat(resultInIndex1.get(0).getId()).isEqualTo(docInIndex1.getId());
assertThat(resultInIndex2).hasSize(1);
assertThat(resultInIndex2.get(0).getId()).isEqualTo(docInIndex2.getId());
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
@Qualifier("vectorStore")
public OpenSearchVectorStore vectorStore(EmbeddingModel embeddingModel) {
try {
return new OpenSearchVectorStore(new OpenSearchClient(ApacheHttpClient5TransportBuilder
@@ -359,6 +400,22 @@ class OpenSearchVectorStoreIT {
}
}
@Bean
@Qualifier("anotherVectorStore")
public OpenSearchVectorStore anotherVectorStore(EmbeddingModel embeddingModel) {
try {
return new OpenSearchVectorStore("another_index",
new OpenSearchClient(ApacheHttpClient5TransportBuilder
.builder(HttpHost.create(opensearchContainer.getHttpHostAddress()))
.build()),
embeddingModel, OpenSearchVectorStore.DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION_1536,
true);
}
catch (URISyntaxException e) {
throw new RuntimeException(e);
}
}
@Bean
public EmbeddingModel embeddingModel() {
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));