feat(knowledge): 智能体查询知识库时,支持基于文件名的模糊过滤功能,不支持 LightRAG
实现知识库检索时可按文件名进行模糊匹配过滤 在Milvus知识库中支持文件名的like表达式过滤 前端展示添加文件名显示 添加相关测试用例验证过滤功能
This commit is contained in:
parent
6d9f3add6c
commit
610b60e204
@ -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"]},
|
||||
)
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
137
test/test_milvus_filter.py
Normal 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())
|
||||
@ -307,7 +307,6 @@ const formatResultData = (data) => {
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
min-width: 0;
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
:deep(.tag) {
|
||||
|
||||
@ -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 {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user