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

实现知识库检索时可按文件名进行模糊匹配过滤
在Milvus知识库中支持文件名的like表达式过滤
前端展示添加文件名显示
添加相关测试用例验证过滤功能
This commit is contained in:
Wenjie Zhang 2025-12-30 20:10:42 +08:00
parent 6d9f3add6c
commit 610b60e204
6 changed files with 179 additions and 7 deletions

View File

@ -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 get_kb_based_tools(db_names: list[str] | None = None) -> list:
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 get_kb_based_tools(db_names: list[str] | None = None) -> list:
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"]},
)

View File

@ -549,8 +549,8 @@ class KnowledgeBase(ABC):
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

View File

@ -458,11 +458,26 @@ class MilvusKB(KnowledgeBase):
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"],
)

137
test/test_milvus_filter.py Normal file
View File

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

View File

@ -307,7 +307,6 @@ const formatResultData = (data) => {
text-overflow: ellipsis;
white-space: nowrap;
min-width: 0;
flex: 1;
}
:deep(.tag) {

View File

@ -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 {