Skip to content

fix(rag): normalize Milvus L2 retrieval scores #3070

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
liugy789 wants to merge 1 commit into
base: main
Choose a base branch
from
+145 −2
Open

fix(rag): normalize Milvus L2 retrieval scores #3070

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 23 additions & 2 deletions ...ntscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -588,6 +588,7 @@ private List<Document> searchDocumentsInMilvus(
.databaseName(databaseName)
.collectionName(collectionName)
.data(Collections.singletonList(queryVector))
.metricType(metricType)
.limit(limit)
.outputFields(
Arrays.asList(
Expand All @@ -609,8 +610,9 @@ private List<Document> searchDocumentsInMilvus(
if (searchResults != null && !searchResults.isEmpty()) {
for (SearchResp.SearchResult result : searchResults.get(0)) {
try {
// Get score
double score = result.getScore();
// Milvus returns a distance for L2. Convert it to a higher-is-better score
// before applying the shared threshold.
double score = normalizeScore(result.getScore());

// Apply score threshold if specified
if (scoreThreshold != null && score < scoreThreshold) {
Expand All @@ -631,6 +633,23 @@ private List<Document> searchDocumentsInMilvus(
return results;
}

/**
* Converts a raw Milvus result to the score contract used by AgentScope.
*
* <p>L2 results are distances where lower values are more similar, so map them to a
* monotonically decreasing score in the [0, 1] range. Other supported metrics already return
* scores in the expected direction and are left unchanged.
*
* @param rawScore the raw score or distance returned by Milvus
* @return a higher-is-better score
*/
private double normalizeScore(double rawScore) {
if (metricType == IndexParam.MetricType.L2) {
return 1.0 / (1.0 + rawScore);
}
return rawScore;
}

/**
* Reconstructs a Document from Milvus search result.
*
Expand Down Expand Up @@ -923,6 +942,8 @@ public Builder connectTimeoutMs(long connectTimeoutMs) {
* Sets the metric type for vector similarity search.
*
* <p>Default is COSINE. Other options include L2 (Euclidean) and IP (Inner Product).
* When opening an existing collection, this must match the collection's vector index.
* L2 distances are converted to {@code 1 / (1 + distance)} scores before thresholding.
*
* @param metricType the metric type
* @return this builder
Expand Down
122 changes: 122 additions & 0 deletions ...ope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/MilvusStoreTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.DisplayName;
Expand Down Expand Up @@ -843,6 +844,57 @@ private MilvusStore createMockStoreForSearchWithResults() throws VectorStoreExce
}
}

private MilvusStore createMockStoreForSearchWithMetric(
IndexParam.MetricType metricType, float... scores) throws VectorStoreException {
return createMockStoreForSearchWithMetric(metricType, null, scores);
}

private MilvusStore createMockStoreForSearchWithMetric(
IndexParam.MetricType metricType,
AtomicReference<SearchReq> capturedSearchRequest,
float... scores)
throws VectorStoreException {
try (MockedConstruction<MilvusClientV2> ignored =
mockConstruction(
MilvusClientV2.class,
(mock, context) -> {
when(mock.hasCollection(any(HasCollectionReq.class))).thenReturn(true);

List<SearchResp.SearchResult> mockResults = new ArrayList<>();
for (int i = 0; i < scores.length; i++) {
SearchResp.SearchResult mockResult =
mock(SearchResp.SearchResult.class);
when(mockResult.getScore()).thenReturn(scores[i]);
Map<String, Object> entity = new HashMap<>();
entity.put("doc_id", "doc-" + i);
entity.put("chunk_id", 0);
entity.put(
"content", "{\"type\":\"text\",\"text\":\"Test content\"}");
when(mockResult.getEntity()).thenReturn(entity);
mockResults.add(mockResult);
}

SearchResp searchResp = mock(SearchResp.class);
when(searchResp.getSearchResults()).thenReturn(List.of(mockResults));
when(mock.search(any(SearchReq.class)))
.thenAnswer(
invocation -> {
if (capturedSearchRequest != null) {
capturedSearchRequest.set(
invocation.getArgument(0));
}
return searchResp;
});
})) {
return MilvusStore.builder()
.uri(TEST_URI)
.collectionName(TEST_COLLECTION)
.dimensions(TEST_DIMENSIONS)
.metricType(metricType)
.build();
}
}

@Test
@DisplayName("Should return error for null query embedding")
void testSearchNullQueryEmbedding() throws VectorStoreException {
Expand Down Expand Up @@ -999,6 +1051,76 @@ void testSearchFilterByScoreThreshold() throws VectorStoreException {
.verifyComplete();
}

@Test
@DisplayName("Should normalize L2 distances to higher-is-better scores")
void testSearchNormalizesL2DistanceScores() throws VectorStoreException {
store = createMockStoreForSearchWithMetric(IndexParam.MetricType.L2, 1.0f, 3.0f);
double[] query = new double[] {1.0, 0.0, 0.0};

StepVerifier.create(
store.search(
SearchDocumentDto.builder()
.queryEmbedding(query)
.limit(10)
.scoreThreshold(0.0)
.build()))
.assertNext(
results -> {
assertEquals(2, results.size());
assertEquals(0.5, results.get(0).getScore(), 0.001);
assertEquals(0.25, results.get(1).getScore(), 0.001);
})
.verifyComplete();
}

@Test
@DisplayName("Should use configured metric type when searching Milvus")
void testSearchUsesConfiguredMetricType() throws VectorStoreException {
AtomicReference<SearchReq> capturedSearchRequest = new AtomicReference<>();
store =
createMockStoreForSearchWithMetric(
IndexParam.MetricType.L2, capturedSearchRequest, 1.0f);
double[] query = new double[] {1.0, 0.0, 0.0};

StepVerifier.create(
store.search(
SearchDocumentDto.builder()
.queryEmbedding(query)
.limit(10)
.scoreThreshold(0.0)
.build()))
.assertNext(
results -> {
assertNotNull(capturedSearchRequest.get());
assertEquals(
IndexParam.MetricType.L2,
capturedSearchRequest.get().getMetricType());
assertEquals(1, results.size());
})
.verifyComplete();
}

@Test
@DisplayName("Should apply score threshold after normalizing L2 distances")
void testSearchAppliesThresholdToNormalizedL2Scores() throws VectorStoreException {
store = createMockStoreForSearchWithMetric(IndexParam.MetricType.L2, 1.0f, 3.0f);
double[] query = new double[] {1.0, 0.0, 0.0};

StepVerifier.create(
store.search(
SearchDocumentDto.builder()
.queryEmbedding(query)
.limit(10)
.scoreThreshold(0.4)
.build()))
.assertNext(
results -> {
assertEquals(1, results.size());
assertEquals(0.5, results.get(0).getScore(), 0.001);
})
.verifyComplete();
}

@Test
@DisplayName("Should return error when store is closed")
void testSearchAfterClose() throws VectorStoreException {
Expand Down
Loading