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

Set ElasticSearch size to match requested topK used in KNN search

Browse Source
This commit is contained in:
dafriz
2025-01-02 21:38:58 +11:00
committed by Ilayaperumal Gopinathan
parent cd5684a66e
commit 95675a85f8
2 changed files with 49 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java
View File

@@ -242,16 +242,15 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
final float finalThreshold = threshold;
float[] vectors = this.embeddingModel.embed(searchRequest.getQuery());
SearchResponse<Document> res = this.elasticsearchClient.search(
sr -> sr.index(this.options.getIndexName())
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
.similarity(finalThreshold)
.k((long) searchRequest.getTopK())
.field("embedding")
.numCandidates((long) (1.5 * searchRequest.getTopK()))
.filter(fl -> fl.queryString(
qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression()))))),
Document.class);
SearchResponse<Document> res = this.elasticsearchClient.search(sr -> sr.index(this.options.getIndexName())
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
.similarity(finalThreshold)
.k((long) searchRequest.getTopK())
.field("embedding")
.numCandidates((long) (1.5 * searchRequest.getTopK()))
.filter(fl -> fl
.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
.size(searchRequest.getTopK()), Document.class);
return res.hits().hits().stream().map(this::toDocument).collect(Collectors.toList());
}

40
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java
View File

@@ -20,6 +20,7 @@ import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Map;
@@ -400,6 +401,45 @@ class ElasticsearchVectorStoreIT {
});
}
@Test
public void overDefaultSizeTest() {
var overDefaultSize = 12;
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
ElasticsearchVectorStore.class);
var testDocs = new ArrayList<Document>();
for (int i = 0; i < overDefaultSize; i++) {
testDocs.add(new Document(String.valueOf(i), "Great Depression " + i, Map.of()));
}
vectorStore.add(testDocs);
Awaitility.await()
.until(() -> vectorStore.similaritySearch(
SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()),
hasSize(1));
List<Document> results = vectorStore.similaritySearch(SearchRequest.builder()
.query("Great Depression")
.topK(overDefaultSize)
.similarityThresholdAll()
.build());
assertThat(results).hasSize(overDefaultSize);
// Remove all documents from the store
vectorStore.delete(testDocs.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(
SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()),
hasSize(0));
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {