diff --git a/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java b/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java index 3b3c41cbe3..9cd3ae1729 100644 --- a/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java +++ b/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java @@ -588,6 +588,7 @@ private List searchDocumentsInMilvus( .databaseName(databaseName) .collectionName(collectionName) .data(Collections.singletonList(queryVector)) + .metricType(metricType) .limit(limit) .outputFields( Arrays.asList( @@ -609,8 +610,9 @@ private List 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) { @@ -631,6 +633,23 @@ private List searchDocumentsInMilvus( return results; } + /** + * Converts a raw Milvus result to the score contract used by AgentScope. + * + *

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. * @@ -923,6 +942,8 @@ public Builder connectTimeoutMs(long connectTimeoutMs) { * Sets the metric type for vector similarity search. * *

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 diff --git a/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/MilvusStoreTest.java b/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/MilvusStoreTest.java index 14f4910714..1b6feccc8a 100644 --- a/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/MilvusStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-rag/agentscope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/MilvusStoreTest.java @@ -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; @@ -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 capturedSearchRequest, + float... scores) + throws VectorStoreException { + try (MockedConstruction ignored = + mockConstruction( + MilvusClientV2.class, + (mock, context) -> { + when(mock.hasCollection(any(HasCollectionReq.class))).thenReturn(true); + + List 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 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 { @@ -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 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 {