100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
140 lines
3.7 KiB
Python
140 lines
3.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
RAG 接口测试：验证图片召回功能
|
|
|
|
运行方式：
|
|
.\venv\Scripts\python.exe scripts/test_rag_image_recall.py
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import json
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from knowledge.manager import get_kb_manager
|
|
from core.engine import get_engine
|
|
|
|
|
|
def test_rag_image_recall():
|
|
"""测试 RAG 检索的图片召回"""
|
|
print("=" * 60)
|
|
print("RAG 图片召回测试")
|
|
print("=" * 60)
|
|
|
|
# 初始化引擎
|
|
engine = get_engine()
|
|
if not engine._initialized:
|
|
engine.initialize()
|
|
|
|
# 测试查询
|
|
test_queries = [
|
|
"文档中有哪些图片？",
|
|
"有哪些图表？",
|
|
"表格内容是什么？",
|
|
]
|
|
|
|
for query in test_queries:
|
|
print(f"\n🔍 查询: {query}")
|
|
print("-" * 40)
|
|
|
|
# 执行检索
|
|
result = engine.search_knowledge(
|
|
query=query,
|
|
collection_name='public_kb',
|
|
top_k=10
|
|
)
|
|
|
|
if not result or not result.get('ids'):
|
|
print(" ❌ 无检索结果")
|
|
continue
|
|
|
|
print(f" 📊 检索到 {len(result['ids'])} 个切片")
|
|
|
|
# 检查图片信息
|
|
images_found = []
|
|
for i, meta in enumerate(result.get('metadatas', [])):
|
|
images_json = meta.get('images_json')
|
|
image_path = meta.get('image_path')
|
|
chunk_type = meta.get('chunk_type')
|
|
|
|
if images_json:
|
|
try:
|
|
imgs = json.loads(images_json)
|
|
images_found.extend(imgs)
|
|
print(f" [{i+1}] chunk_type={chunk_type}, images_json={len(imgs)}张图片")
|
|
except:
|
|
pass
|
|
|
|
if image_path:
|
|
images_found.append({'id': image_path, 'type': 'image_path'})
|
|
print(f" [{i+1}] chunk_type={chunk_type}, image_path={image_path}")
|
|
|
|
if images_found:
|
|
print(f"\n ✅ 共召回 {len(images_found)} 张图片")
|
|
else:
|
|
print("\n ⚠️ 未召回任何图片")
|
|
|
|
|
|
def test_extract_rich_media():
|
|
"""测试 _extract_rich_media 函数"""
|
|
print("\n" + "=" * 60)
|
|
print("_extract_rich_media 函数测试")
|
|
print("=" * 60)
|
|
|
|
from core.agentic import AgenticRAG
|
|
|
|
# 创建 AgenticRAG 实例
|
|
rag = AgenticRAG()
|
|
|
|
# 模拟检索结果
|
|
mock_contexts = [
|
|
{
|
|
'doc': '这是一段文本',
|
|
'meta': {
|
|
'chunk_type': 'text',
|
|
'source': 'test.pdf',
|
|
'page': 1,
|
|
}
|
|
},
|
|
{
|
|
'doc': '这是一张图片',
|
|
'meta': {
|
|
'chunk_type': 'image',
|
|
'source': 'test.pdf',
|
|
'page': 2,
|
|
'image_path': 'abc123.jpg',
|
|
}
|
|
},
|
|
{
|
|
'doc': '这是带关联图片的文本',
|
|
'meta': {
|
|
'chunk_type': 'text',
|
|
'source': 'test.pdf',
|
|
'page': 3,
|
|
'images_json': '[{"id": "img1.jpg", "order": 1}, {"id": "img2.jpg", "order": 2}]',
|
|
}
|
|
},
|
|
]
|
|
|
|
# 调用 _extract_rich_media
|
|
rich_media = rag._extract_rich_media(mock_contexts)
|
|
|
|
print(f"\n📊 提取结果:")
|
|
print(f" - 图片数量: {len(rich_media.get('images', []))}")
|
|
print(f" - 表格数量: {len(rich_media.get('tables', []))}")
|
|
|
|
for img in rich_media.get('images', []):
|
|
print(f" 📷 {img}")
|
|
|
|
if not rich_media.get('images'):
|
|
print(" ⚠️ 未提取到图片，检查 _extract_rich_media 逻辑")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if sys.platform == 'win32':
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|
|
|
test_rag_image_recall()
|
|
test_extract_rich_media()
|