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
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)

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Passing metricType explicitly to SearchReq is new behaviour: previously Milvus searched with the metric recorded on the collection index, now it searches with the configured metric. When a store connects to an existing collection whose index metric differs from the builder default (COSINE), this turns a previously-working search into a provider-side error or wrong ranking. Could you validate the configured metric against describeIndex/describeCollection at connect time and fail fast with a message naming both values? The javadoc note alone will be easy to miss for people who configured the store before this change.

.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

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 / (1 + distance) is a reasonable monotone map into (0, 1], and applying the threshold after normalization is the correct order. This is still a breaking change for anyone who already tuned scoreThreshold against raw L2 distances, and the resulting scores are not comparable with COSINE/IP scores. Worth a changelog entry plus a note in the RAG docs that thresholds must be re-tuned for L2 collections.

*/
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