feat(knowledge): 智能体查询知识库时，支持基于文件名的模糊过滤功能，不支持 LightRAG

xerrors · xerrors · commit 610b60e20401 · 2025-12-30T20:10:42.000+08:00
实现知识库检索时可按文件名进行模糊匹配过滤
在Milvus知识库中支持文件名的like表达式过滤
前端展示添加文件名显示
添加相关测试用例验证过滤功能
diff --git a/src/agents/common/tools.py b/src/agents/common/tools.py
@@ -122,6 +122,11 @@ class KnowledgeRetrieverModel(BaseModel):
     )
 
 
+class CommonKnowledgeRetriever(KnowledgeRetrieverModel):
+    """Common knowledge retriever model."""
+    file_name: str = Field(description="限定文件名称，当操作类型为 'search' 时，可以指定文件名称，支持模糊匹配")
+
+
 def get_kb_based_tools(db_names: list[str] | None = None) -> list:
     """获取所有知识库基于的工具"""
     # 获取所有知识库
@@ -132,7 +137,9 @@ def get_kb_based_tools(db_names: list[str] | None = None) -> list:
     def _create_retriever_wrapper(db_id: str, retriever_info: dict[str, Any]):
         """创建检索器包装函数的工厂函数，避免闭包变量捕获问题"""
 
-        async def async_retriever_wrapper(query_text: str, operation: str = "search") -> Any:
+        async def async_retriever_wrapper(
+            query_text: str, operation: str = "search", file_name: str | None = None
+        ) -> Any:
             """异步检索器包装函数，支持检索和获取思维导图"""
 
             # 获取思维导图
@@ -173,10 +180,14 @@ def mindmap_to_text(node, level=0):
             retriever = retriever_info["retriever"]
             try:
                 logger.debug(f"Retrieving from database {db_id} with query: {query_text}")
+                kwargs = {}
+                if file_name:
+                    kwargs["file_name"] = file_name
+
                 if asyncio.iscoroutinefunction(retriever):
-                    result = await retriever(query_text)
+                    result = await retriever(query_text, **kwargs)
                 else:
-                    result = retriever(query_text)
+                    result = retriever(query_text, **kwargs)
                 logger.debug(f"Retrieved {len(result) if isinstance(result, list) else 'N/A'} results from {db_id}")
                 return result
             except Exception as e:
@@ -207,12 +218,16 @@ def mindmap_to_text(node, level=0):
 
             safename = retrieve_info["name"].replace(" ", "_")[:20]
 
+            args_schema = KnowledgeRetrieverModel
+            if retrieve_info["metadata"]["kb_type"] in ["milvus"]:
+                args_schema = CommonKnowledgeRetriever
+
             # 使用 StructuredTool.from_function 创建异步工具
             tool = StructuredTool.from_function(
                 coroutine=retriever_wrapper,
                 name=safename,
                 description=description,
-                args_schema=KnowledgeRetrieverModel,
+                args_schema=args_schema,
                 metadata=retrieve_info["metadata"] | {"tag": ["knowledgebase"]},
             )
 
diff --git a/src/knowledge/base.py b/src/knowledge/base.py
@@ -549,8 +549,8 @@ def get_retrievers(self) -> dict[str, dict]:
         for db_id, meta in self.databases_meta.items():
 
             def make_retriever(db_id):
-                async def retriever(query_text):
-                    return await self.aquery(query_text, db_id)
+                async def retriever(query_text, **kwargs):
+                    return await self.aquery(query_text, db_id, **kwargs)
 
                 return retriever
 
diff --git a/src/knowledge/implementations/milvus.py b/src/knowledge/implementations/milvus.py
@@ -458,11 +458,26 @@ async def aquery(self, query_text: str, db_id: str, **kwargs) -> list[dict]:
             query_embedding = embedding_function([query_text])
 
             search_params = {"metric_type": metric_type, "params": {"nprobe": 10}}
+
+            # 构建过滤表达式
+            expr = None
+            if file_name := kwargs.get("file_name"):
+                # 使用 like 支持模糊匹配
+                # 注意：需要转义双引号以防止注入
+                safe_file_name = file_name.replace('"', '\\"')
+                # 如果没有提供通配符，默认前后添加 %
+                if "%" not in safe_file_name:
+                    expr = f'source like "%{safe_file_name}%"'
+                else:
+                    expr = f'source like "{safe_file_name}"'
+                logger.debug(f"Using filter expression: {expr}")
+
             results = collection.search(
                 data=query_embedding,
                 anns_field="embedding",
                 param=search_params,
                 limit=recall_top_k,
+                expr=expr,
                 output_fields=["content", "source", "chunk_id", "file_id", "chunk_index"],
             )
 
diff --git a/test/test_milvus_filter.py b/test/test_milvus_filter.py
@@ -0,0 +1,137 @@
+import asyncio
+import os
+import shutil
+from unittest.mock import MagicMock, patch
+
+from src.knowledge import knowledge_base
+from src.utils import logger
+
+# Mock Embedding Model
+class MockEmbeddingModel:
+    async def abatch_encode(self, texts, batch_size=None):
+        # Return dummy vectors of dim 4
+        return [[0.1, 0.2, 0.3, 0.4] for _ in texts]
+    
+    def batch_encode(self, texts, batch_size=None):
+        return [[0.1, 0.2, 0.3, 0.4] for _ in texts]
+
+# Test function
+async def test_milvus_filter():
+    logger.info("Starting Milvus Filter Test")
+    
+    # Check if Milvus is available (pymilvus installed and connection works)
+    try:
+        from pymilvus import connections, utility
+        # Assuming Milvus is running at default location
+        connections.connect(alias="default", uri=os.getenv("MILVUS_URI", "http://localhost:19530"))
+        logger.info("Connected to Milvus")
+    except Exception as e:
+        logger.warning(f"Milvus not available or connection failed: {e}")
+        # Proceeding might fail, but let's try.
+    
+    db_id = "test_milvus_filter_db"
+    file1 = "test_file_A.txt"
+    file2 = "test_file_B.txt"
+
+    # Patch embedding model
+    with patch("src.models.embed.select_embedding_model", return_value=MockEmbeddingModel()):
+        
+        try:
+            # Cleanup if exists
+            if db_id in knowledge_base.global_databases_meta:
+                await knowledge_base.delete_database(db_id)
+                
+            # Create DB
+            logger.info("Creating database...")
+            # explicitly set dimension to 4 to match mock
+            await knowledge_base.create_database(
+                database_name="Test Milvus Filter",
+                description="Test DB",
+                kb_type="milvus",
+                embed_info={"name": "mock-embedding", "dimension": 4, "model_id": "mock"}
+            )
+            
+            # Get actual db_id
+            target_db = next((db for db in knowledge_base.get_databases()["databases"] if db["name"] == "Test Milvus Filter"), None)
+            if not target_db:
+                logger.error("Failed to create DB")
+                return
+            
+            db_id = target_db["db_id"]
+            logger.info(f"DB created with ID: {db_id}")
+            
+            # Create dummy files
+            
+            with open(file1, "w") as f:
+                f.write("Apple content.")
+            with open(file2, "w") as f:
+                f.write("Banana content.")
+                
+            # Add content
+            logger.info("Adding content...")
+            await knowledge_base.add_content(db_id, [os.path.abspath(file1), os.path.abspath(file2)])
+            
+            # Wait for data to be visible
+            logger.info("Waiting for data to be visible...")
+            await asyncio.sleep(2)
+            
+            # Query without filter
+            logger.info("Querying without filter...")
+            results = await knowledge_base.aquery("content", db_id)
+            logger.info(f"No filter results: {len(results)}")
+            
+            # Verify we have chunks from both files
+            sources = [r['metadata']['source'] for r in results]
+            logger.info(f"Sources: {sources}")
+            
+            # Query with filter A (Partial Match)
+            logger.info("Querying with filter A (file_A)...")
+            results_a = await knowledge_base.aquery("content", db_id, file_name="file_A")
+            logger.info(f"Filter A results: {len(results_a)}")
+            
+            if len(results_a) == 0:
+                logger.error("FAIL: Filter A returned 0 results")
+            
+            for r in results_a:
+                source = r['metadata']['source']
+                logger.info(f" - {source}")
+                if "test_file_A.txt" not in source:
+                        logger.error(f"FAIL: Expected test_file_A.txt, got {source}")
+                        raise AssertionError("Filter A failed")
+
+            # Query with wildcard filter
+            logger.info("Querying with wildcard filter (%B.txt)...")
+            results_b = await knowledge_base.aquery("content", db_id, file_name="%B.txt")
+            logger.info(f"Filter B results: {len(results_b)}")
+            if len(results_b) == 0:
+                    logger.error("FAIL: Wildcard filter returned 0 results")
+            
+            for r in results_b:
+                source = r['metadata']['source']
+                logger.info(f" - {source}")
+                if "test_file_B.txt" not in source:
+                        logger.error(f"FAIL: Expected test_file_B.txt, got {source}")
+                        raise AssertionError("Filter B failed")
+            
+            if len(results_a) > 0 and len(results_b) > 0:
+                logger.info("Test passed!")
+            else:
+                logger.error("Test failed: No results found for one or more queries")
+            
+        except Exception as e:
+            logger.error(f"Test failed with exception: {e}")
+            raise
+        finally:
+            # Cleanup
+            logger.info("Cleaning up...")
+            try:
+                await knowledge_base.delete_database(db_id)
+            except Exception:
+                pass
+            if os.path.exists(file1):
+                os.remove(file1)
+            if os.path.exists(file2):
+                os.remove(file2)
+
+if __name__ == "__main__":
+    asyncio.run(test_milvus_filter())
diff --git a/web/src/components/ToolCallingResult/BaseToolCall.vue b/web/src/components/ToolCallingResult/BaseToolCall.vue
@@ -307,7 +307,6 @@ const formatResultData = (data) => {
         text-overflow: ellipsis;
         white-space: nowrap;
         min-width: 0;
-        flex: 1;
       }
 
       :deep(.tag) {
diff --git a/web/src/components/ToolCallingResult/tools/KnowledgeBaseTool.vue b/web/src/components/ToolCallingResult/tools/KnowledgeBaseTool.vue
@@ -5,6 +5,8 @@
         <span class="note">{{ operationLabel }}</span>
         <span class="separator" v-if="queryText">|</span>
         <span class="description">{{ queryText }}</span>
+        <span class="separator" v-if="fileName">|</span>
+        <span class="description" v-if="fileName">文件: {{ fileName }}</span>
       </div>
     </template>
     <template #result="{ resultContent }">
@@ -169,6 +171,10 @@ const queryText = computed(() => {
   return args.value.query_text || '';
 });
 
+const fileName = computed(() => {
+  return args.value.file_name || '';
+});
+
 const parseData = (content) => {
   if (typeof content === 'string') {
     try {

Original file line number	Diff line number	Diff line change
`@@ -307,7 +307,6 @@ const formatResultData = (data) => {`
`307`	`307`	`text-overflow: ellipsis;`
`308`	`308`	`white-space: nowrap;`
`309`	`309`	`min-width: 0;`
`310`		`- flex: 1;`
`311`	`310`	`}`
`312`	`311`
`313`	`312`	`:deep(.tag) {`