init: RAG 知识库服务初始提交
- 后端 API(Flask + Gunicorn) - RAG 引擎(混合检索 + 云端 Reranker + 引用溯源) - 文档解析(MinerU + 多格式支持) - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
139
scripts/test_rag_image_recall.py
Normal file
139
scripts/test_rag_image_recall.py
Normal file
@@ -0,0 +1,139 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user