ai-robot-core/ai-service/scripts/test_kb_search_v2.py

97 lines
3.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
测试 kb_search_dynamic 工具 - 增加超时时间
"""
import asyncio
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from app.services.mid.kb_search_dynamic_tool import KbSearchDynamicTool, KbSearchDynamicConfig
from app.core.config import get_settings
async def test_kb_search():
"""测试知识库搜索 - 增加超时时间"""
settings = get_settings()
# 创建数据库会话
engine = create_async_engine(settings.database_url)
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
async with async_session() as session:
# 使用更长的超时时间
config = KbSearchDynamicConfig(
enabled=True,
top_k=5,
timeout_ms=10000, # 10秒超时
min_score_threshold=0.5,
)
tool = KbSearchDynamicTool(session=session, config=config)
# 测试参数
test_cases = [
{
"name": "完整参数含context过滤",
"params": {
"query": "三年级语文学习",
"tenant_id": "szmp@ash@2026",
"scene": "学习方案",
"top_k": 5,
"context": {"grade": "三年级", "subject": "语文"},
}
},
{
"name": "简化参数无context",
"params": {
"query": "三年级语文学习",
"tenant_id": "szmp@ash@2026",
"scene": "学习方案",
"top_k": 5,
}
},
]
for test_case in test_cases:
print(f"\n{'='*80}")
print(f"测试: {test_case['name']}")
print(f"{'='*80}")
print(f"参数: {test_case['params']}")
print(f"超时设置: {config.timeout_ms}ms")
try:
result = await tool.execute(**test_case['params'])
print(f"\n结果:")
print(f" success: {result.success}")
print(f" hits count: {len(result.hits)}")
print(f" applied_filter: {result.applied_filter}")
print(f" fallback_reason_code: {result.fallback_reason_code}")
print(f" duration_ms: {result.duration_ms}")
if result.hits:
print(f"\n 所有结果:")
for i, hit in enumerate(result.hits, 1):
text = hit.get('text', '')[:100] + '...' if hit.get('text') else 'N/A'
score = hit.get('score', 0)
metadata = hit.get('metadata', {})
collection = hit.get('collection', 'unknown')
print(f" {i}. [score={score:.4f}] [collection={collection}]")
print(f" text: {text}")
print(f" metadata: {metadata}")
else:
print(f"\n ⚠️ 没有命中任何结果")
except Exception as e:
print(f"\n ❌ 错误: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
asyncio.run(test_kb_search())