feat(rag): dense+BM25 hybrid on MilvusClientV2, drop SDK path
Replace legacy MilvusServiceClient knowledge search/write with a single MilvusClientV2 hybrid store (BM25 function + dense ANN + RRFRanker). Use collection biz_hybrid and require knowledge reindex.
This commit is contained in:
zhuyongxin
@@ -1,208 +1,61 @@
|
||||
package com.superbiz.agent.service;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.service.milvus.MilvusHybridKnowledgeStore;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNull;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.doReturn;
|
||||
import static org.mockito.ArgumentMatchers.isNull;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.spy;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
class VectorSearchServiceTest {
|
||||
|
||||
@Test
|
||||
void sdkModeBypassesVectorStore() {
|
||||
VectorSearchService service = spy(new VectorSearchService());
|
||||
setMode(service, "sdk");
|
||||
VectorSearchService.SearchResult expected = result("sdk-doc", 0.2f);
|
||||
doReturn(List.of(expected))
|
||||
.when(service).searchSimilarDocumentsWithSdk("query", 3, null);
|
||||
void denseModeCallsHybridStoreDenseSearch() {
|
||||
MilvusHybridKnowledgeStore store = mock(MilvusHybridKnowledgeStore.class);
|
||||
VectorEmbeddingService embeddingService = mock(VectorEmbeddingService.class);
|
||||
when(embeddingService.generateQueryVector("query")).thenReturn(List.of(0.1f, 0.2f));
|
||||
VectorSearchService.SearchResult expected = result("doc-1", 0.2f);
|
||||
when(store.searchDense(eq("query"), eq(List.of(0.1f, 0.2f)), eq(3), isNull()))
|
||||
.thenReturn(List.of(expected));
|
||||
when(store.collectionName()).thenReturn("biz_hybrid");
|
||||
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
ReflectionTestUtils.setField(service, "knowledgeStore", store);
|
||||
ReflectionTestUtils.setField(service, "embeddingService", embeddingService);
|
||||
ReflectionTestUtils.setField(service, "searchMode", "dense");
|
||||
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("query", 3, null);
|
||||
|
||||
assertEquals(List.of(expected), results);
|
||||
verify(service, never()).searchSimilarDocumentsWithVectorStore(any(), eq(3), any());
|
||||
verify(store).searchDense(eq("query"), eq(List.of(0.1f, 0.2f)), eq(3), isNull());
|
||||
}
|
||||
|
||||
@Test
|
||||
void autoModeUsesVectorStoreWhenAvailable() {
|
||||
VectorStore vectorStore = mock(VectorStore.class);
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(vectorStore);
|
||||
when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(List.of(
|
||||
Document.builder()
|
||||
.id("spring-doc")
|
||||
.text("spring content")
|
||||
.metadata(Map.of("_source", "spring.md", "category", "api"))
|
||||
.score(0.8)
|
||||
.build()
|
||||
));
|
||||
void hybridModeCallsHybridStoreHybridSearch() {
|
||||
MilvusHybridKnowledgeStore store = mock(MilvusHybridKnowledgeStore.class);
|
||||
VectorEmbeddingService embeddingService = mock(VectorEmbeddingService.class);
|
||||
when(embeddingService.generateQueryVector("pool")).thenReturn(List.of(0.3f));
|
||||
VectorSearchService.SearchResult expected = result("doc-h", 0.4f);
|
||||
when(store.searchHybrid(eq("pool"), eq(List.of(0.3f)), eq(5), eq("mysql")))
|
||||
.thenReturn(List.of(expected));
|
||||
when(store.collectionName()).thenReturn("biz_hybrid");
|
||||
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setMode(service, "auto");
|
||||
setVectorStore(service, provider);
|
||||
ReflectionTestUtils.setField(service, "knowledgeStore", store);
|
||||
ReflectionTestUtils.setField(service, "embeddingService", embeddingService);
|
||||
ReflectionTestUtils.setField(service, "searchMode", "hybrid");
|
||||
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("query", 3, null);
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("pool", 5, "mysql");
|
||||
|
||||
assertEquals(1, results.size());
|
||||
assertEquals("spring-doc", results.get(0).getId());
|
||||
assertEquals("similarity", results.get(0).getScoreLabel());
|
||||
assertEquals(0.8, results.get(0).getRawScore(), 0.0001);
|
||||
assertEquals(0.4f, results.get(0).getScore(), 0.0001);
|
||||
assertTrue(results.get(0).getMetadata().contains("spring.md"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void vectorStoreSearchUsesDistanceMetadataAsCompatibleScore() {
|
||||
VectorStore vectorStore = mock(VectorStore.class);
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(vectorStore);
|
||||
when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(List.of(
|
||||
Document.builder()
|
||||
.id("spring-doc")
|
||||
.text("spring content")
|
||||
.metadata(Map.of("distance", 0.5659486, "category", "api"))
|
||||
.score(0.4340514)
|
||||
.build()
|
||||
));
|
||||
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setMode(service, "auto");
|
||||
setVectorStore(service, provider);
|
||||
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("query", 3, null);
|
||||
|
||||
assertEquals(1, results.size());
|
||||
assertEquals("similarity", results.get(0).getScoreLabel());
|
||||
assertEquals(0.4340514, results.get(0).getRawScore(), 0.0001);
|
||||
assertEquals(0.5659486f, results.get(0).getScore(), 0.0001);
|
||||
}
|
||||
|
||||
@Test
|
||||
void autoModeFallsBackToSdkWhenVectorStoreFails() {
|
||||
VectorStore vectorStore = mock(VectorStore.class);
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(vectorStore);
|
||||
when(vectorStore.similaritySearch(any(SearchRequest.class))).thenThrow(new RuntimeException("vectorstore down"));
|
||||
|
||||
VectorSearchService service = spy(new VectorSearchService());
|
||||
setMode(service, "auto");
|
||||
setVectorStore(service, provider);
|
||||
VectorSearchService.SearchResult fallback = result("sdk-doc", 0.3f);
|
||||
doReturn(List.of(fallback))
|
||||
.when(service).searchSimilarDocumentsWithSdk("query", 3, null);
|
||||
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("query", 3, null);
|
||||
|
||||
assertEquals(List.of(fallback), results);
|
||||
}
|
||||
|
||||
@Test
|
||||
void autoModeFallsBackToSdkWhenVectorStoreUnavailable() {
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(null);
|
||||
|
||||
VectorSearchService service = spy(new VectorSearchService());
|
||||
setMode(service, "auto");
|
||||
setVectorStore(service, provider);
|
||||
VectorSearchService.SearchResult fallback = result("sdk-doc", 0.3f);
|
||||
doReturn(List.of(fallback))
|
||||
.when(service).searchSimilarDocumentsWithSdk("query", 3, null);
|
||||
|
||||
List<VectorSearchService.SearchResult> results = service.searchSimilarDocuments("query", 3, null);
|
||||
|
||||
assertEquals(List.of(fallback), results);
|
||||
}
|
||||
|
||||
@Test
|
||||
void vectorStoreSearchUsesCategoryFilter() {
|
||||
VectorStore vectorStore = mock(VectorStore.class);
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(vectorStore);
|
||||
when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(List.of());
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setMode(service, "spring-ai");
|
||||
setVectorStore(service, provider);
|
||||
|
||||
service.searchSimilarDocuments("query", 5, "api");
|
||||
|
||||
ArgumentCaptor<SearchRequest> requestCaptor = ArgumentCaptor.forClass(SearchRequest.class);
|
||||
verify(vectorStore).similaritySearch(requestCaptor.capture());
|
||||
SearchRequest request = requestCaptor.getValue();
|
||||
assertEquals("query", request.getQuery());
|
||||
assertEquals(5, request.getTopK());
|
||||
assertTrue(request.hasFilterExpression());
|
||||
assertTrue(request.toString().contains("category"));
|
||||
assertTrue(request.toString().contains("api"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void springModeUsesVectorStoreAlias() {
|
||||
VectorStore vectorStore = mock(VectorStore.class);
|
||||
ObjectProvider<VectorStore> provider = mock(ObjectProvider.class);
|
||||
when(provider.getIfAvailable()).thenReturn(vectorStore);
|
||||
when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(List.of());
|
||||
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setMode(service, "spring");
|
||||
setVectorStore(service, provider);
|
||||
|
||||
service.searchSimilarDocuments("query", 5, "api");
|
||||
|
||||
verify(vectorStore).similaritySearch(any(SearchRequest.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void defaultScopeDoesNotAddMetadataFilter() {
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setKbScope(service, "");
|
||||
|
||||
assertNull(service.buildSpringAiFilterExpression(null));
|
||||
assertNull(service.buildSdkFilterExpression(null));
|
||||
assertEquals("category == 'api'", service.buildSpringAiFilterExpression("api"));
|
||||
assertEquals("metadata[\"category\"] == \"api\"", service.buildSdkFilterExpression("api"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void configuredScopeCombinesWithCategoryFilter() {
|
||||
VectorSearchService service = new VectorSearchService();
|
||||
setKbScope(service, "rag-eval");
|
||||
|
||||
assertEquals("kb_scope == 'rag-eval'", service.buildSpringAiFilterExpression(null));
|
||||
assertEquals("category == 'api' && kb_scope == 'rag-eval'",
|
||||
service.buildSpringAiFilterExpression("api"));
|
||||
assertEquals("metadata[\"category\"] == \"api\" && metadata[\"kb_scope\"] == \"rag-eval\"",
|
||||
service.buildSdkFilterExpression("api"));
|
||||
}
|
||||
|
||||
private static void setMode(VectorSearchService service, String mode) {
|
||||
ReflectionTestUtils.setField(service, "vectorStoreMode", mode);
|
||||
}
|
||||
|
||||
private static void setVectorStore(VectorSearchService service, ObjectProvider<VectorStore> provider) {
|
||||
ReflectionTestUtils.setField(service, "vectorStoreProvider", provider);
|
||||
ReflectionTestUtils.setField(service, "objectMapper", new ObjectMapper());
|
||||
ReflectionTestUtils.setField(service, "maxL2Distance", 2.0);
|
||||
}
|
||||
|
||||
private static void setKbScope(VectorSearchService service, String kbScope) {
|
||||
ReflectionTestUtils.setField(service, "kbScope", kbScope);
|
||||
assertEquals("doc-h", results.get(0).getId());
|
||||
verify(store).searchHybrid(eq("pool"), eq(List.of(0.3f)), eq(5), eq("mysql"));
|
||||
}
|
||||
|
||||
private static VectorSearchService.SearchResult result(String id, float score) {
|
||||
|
||||
+12
-17
@@ -3,41 +3,36 @@ package com.superbiz.agent.service.retrieval;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.service.VectorSearchService;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
class VectorKnowledgeSearchAdapterHybridTest {
|
||||
|
||||
@Test
|
||||
void hybridFusesFilteredAndUnfilteredDensePaths() {
|
||||
void adapterMapsStoreHitsWithChunkIdentity() {
|
||||
VectorSearchService vectorSearchService = mock(VectorSearchService.class);
|
||||
when(vectorSearchService.searchSimilarDocuments("pool timeout", 3, null)).thenReturn(List.of(
|
||||
result("u1", "{\"_source\":\"a.md\",\"docId\":\"a\",\"chunkIndex\":0,\"title\":\"generic\"}", "generic pool", 0.4f),
|
||||
result("u2", "{\"_source\":\"b.md\",\"docId\":\"b\",\"chunkIndex\":0,\"title\":\"other\"}", "other", 0.5f)
|
||||
));
|
||||
when(vectorSearchService.searchSimilarDocuments("pool timeout", 3, "mysql")).thenReturn(List.of(
|
||||
result("f1", "{\"_source\":\"c.md\",\"docId\":\"c\",\"chunkIndex\":0,\"title\":\"mysql pool timeout\"}", "mysql pool timeout runbook", 0.35f)
|
||||
result("id-1",
|
||||
"{\"_source\":\"c.md\",\"docId\":\"c\",\"chunkIndex\":0,\"title\":\"mysql pool timeout\"}",
|
||||
"mysql pool timeout runbook",
|
||||
0.35f)
|
||||
));
|
||||
|
||||
VectorKnowledgeSearchAdapter adapter = new VectorKnowledgeSearchAdapter(vectorSearchService, new ObjectMapper());
|
||||
ReflectionTestUtils.setField(adapter, "configuredMode", "hybrid");
|
||||
ReflectionTestUtils.setField(adapter, "rrfK", 60);
|
||||
ReflectionTestUtils.setField(adapter, "weightDenseUnfiltered", 1.0);
|
||||
ReflectionTestUtils.setField(adapter, "weightDenseFiltered", 1.0);
|
||||
ReflectionTestUtils.setField(adapter, "weightLexical", 1.0);
|
||||
VectorKnowledgeSearchAdapter adapter =
|
||||
new VectorKnowledgeSearchAdapter(vectorSearchService, new ObjectMapper());
|
||||
|
||||
List<KnowledgeSearchHit> hits = adapter.search(
|
||||
new KnowledgeSearchRequest("pool timeout", 3, "mysql", KnowledgeSearchMode.HYBRID));
|
||||
|
||||
assertEquals(3, hits.size());
|
||||
assertTrue(hits.stream().anyMatch(hit -> "c#chunk-0".equals(hit.evidenceKey())));
|
||||
assertTrue(hits.get(0).metadata().containsKey("fusedScore"));
|
||||
assertEquals(1, hits.size());
|
||||
assertEquals("c#chunk-0", hits.get(0).evidenceKey());
|
||||
assertEquals("c.md", hits.get(0).source());
|
||||
verify(vectorSearchService).searchSimilarDocuments("pool timeout", 3, "mysql");
|
||||
}
|
||||
|
||||
private static VectorSearchService.SearchResult result(String id, String metadata, String content, float score) {
|
||||
|
||||
Reference in New Issue
Block a user