208 lines
6.7 KiB
Python
208 lines
6.7 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
完整的多格式文件测试 - 包含图片的 DOCX、PDF、Excel、CSV
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import asyncio
|
||
import warnings
|
||
from pathlib import Path
|
||
import shutil
|
||
|
||
# 过滤掉PyTorch的FutureWarning
|
||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||
|
||
# 添加源码路径
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||
|
||
from base_rag.core import BaseRAG
|
||
|
||
|
||
class AdvancedFormatRAG(BaseRAG):
|
||
"""高级格式文件处理的RAG实现"""
|
||
|
||
async def ingest(self, file_path: str, **kwargs):
|
||
"""实现文档导入逻辑"""
|
||
return await self.process_file_to_vector_store(file_path, **kwargs)
|
||
|
||
async def query(self, question: str) -> str:
|
||
"""实现查询逻辑"""
|
||
docs = await self.similarity_search_with_rerank(question, k=3)
|
||
|
||
if not docs:
|
||
return "抱歉,没有找到相关信息。"
|
||
|
||
# 显示搜索到的文档来源
|
||
sources = []
|
||
contexts = []
|
||
for doc in docs:
|
||
source = doc.metadata.get("source_file", "未知来源")
|
||
content = doc.page_content.strip()
|
||
|
||
if source not in sources:
|
||
sources.append(source)
|
||
contexts.append(content)
|
||
|
||
context = "\n\n".join(contexts)
|
||
sources_str = "、".join(sources)
|
||
|
||
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||
|
||
|
||
async def test_advanced_formats():
|
||
"""测试高级文件格式处理"""
|
||
print("🚀 高级多格式文件处理测试")
|
||
print("=" * 60)
|
||
|
||
# 清理旧的向量数据库
|
||
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/advanced_formats")
|
||
if db_path.exists():
|
||
shutil.rmtree(db_path)
|
||
print("🧹 已清理旧的向量数据库")
|
||
|
||
# 创建RAG实例
|
||
rag = AdvancedFormatRAG(
|
||
vector_store_name="advanced_formats",
|
||
retriever_top_k=3,
|
||
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/test_files",
|
||
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/advanced_status.db",
|
||
)
|
||
|
||
# 测试文件列表 - 包含新创建的文件
|
||
test_files = [
|
||
{
|
||
"file": "complex_data_science.docx",
|
||
"format": "DOCX",
|
||
"description": "复杂Word文档(含表格和图片)"
|
||
},
|
||
{
|
||
"file": "sales_data.csv",
|
||
"format": "CSV",
|
||
"description": "销售数据CSV文件"
|
||
},
|
||
{
|
||
"file": "company_report.xlsx",
|
||
"format": "XLSX",
|
||
"description": "多工作表Excel文件"
|
||
},
|
||
{
|
||
"file": "ai_research_report.pdf",
|
||
"format": "PDF",
|
||
"description": "AI研究报告PDF(含图片)"
|
||
}
|
||
]
|
||
|
||
print("📂 处理高级格式文件...")
|
||
processed_count = 0
|
||
|
||
for file_info in test_files:
|
||
filename = file_info["file"]
|
||
format_type = file_info["format"]
|
||
description = file_info["description"]
|
||
|
||
file_path = Path("../test_files") / filename
|
||
|
||
if not file_path.exists():
|
||
# 尝试绝对路径
|
||
file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename
|
||
|
||
if not file_path.exists():
|
||
print(f"❌ {format_type}: {filename} - 文件不存在")
|
||
continue
|
||
|
||
print(f"📄 处理 {format_type}: {filename}")
|
||
print(f" {description}")
|
||
|
||
try:
|
||
result = await rag.ingest(str(file_path))
|
||
if result and result.get('success'):
|
||
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||
processed_count += 1
|
||
else:
|
||
print(f" ⚠️ 跳过: {result.get('message', '可能已存在')}")
|
||
if "已经处理完毕" in str(result.get('message', '')):
|
||
processed_count += 1
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {str(e)}")
|
||
print()
|
||
|
||
print(f"📊 处理完成: {processed_count}/{len(test_files)} 个文件")
|
||
print()
|
||
|
||
# 测试针对性查询
|
||
print("💬 高级格式查询测试...")
|
||
|
||
queries = [
|
||
{
|
||
"question": "数据科学的核心技术有哪些?",
|
||
"expected": "complex_data_science.docx"
|
||
},
|
||
{
|
||
"question": "销售数据中哪个产品销售额最高?",
|
||
"expected": "sales_data.csv"
|
||
},
|
||
{
|
||
"question": "公司员工信息包含哪些部门?",
|
||
"expected": "company_report.xlsx"
|
||
},
|
||
{
|
||
"question": "人工智能研究面临的挑战是什么?",
|
||
"expected": "ai_research_report.pdf"
|
||
},
|
||
{
|
||
"question": "Python在数据科学中的作用?",
|
||
"expected": "多个文档"
|
||
}
|
||
]
|
||
|
||
for i, query_info in enumerate(queries, 1):
|
||
question = query_info["question"]
|
||
expected = query_info["expected"]
|
||
|
||
print(f"\n❓ 查询 {i}: {question}")
|
||
print(f" 期望来源: {expected}")
|
||
|
||
try:
|
||
answer = await rag.query(question)
|
||
if "抱歉" not in answer:
|
||
# 分离来源信息和内容
|
||
parts = answer.split('\n\n', 1)
|
||
if len(parts) == 2:
|
||
source_info = parts[0]
|
||
content = parts[1]
|
||
|
||
print(f" 📚 {source_info}")
|
||
|
||
# 显示内容摘要(前150字符)
|
||
if len(content) > 150:
|
||
content_preview = content[:150] + "..."
|
||
else:
|
||
content_preview = content
|
||
|
||
print(f" 💡 {content_preview}")
|
||
else:
|
||
print(f" 💡 {answer[:150]}...")
|
||
else:
|
||
print(f" 💡 {answer}")
|
||
except Exception as e:
|
||
print(f" ❌ 查询失败: {str(e)}")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("🎉 高级多格式文件测试完成!")
|
||
print("✅ 支持的格式:")
|
||
print(" 📄 DOCX - Word文档 (含表格、图片)")
|
||
print(" 📊 CSV - 逗号分隔值文件")
|
||
print(" 📈 XLSX - Excel工作簿 (多工作表)")
|
||
print(" 📑 PDF - 便携式文档格式 (含图片)")
|
||
print()
|
||
print("🔧 技术特性:")
|
||
print(" 🔄 异步处理 - 非阻塞I/O操作")
|
||
print(" 🧠 智能解析 - 自动识别文件格式")
|
||
print(" 🔍 跨格式查询 - 统一检索接口")
|
||
print(" 📋 表格数据提取 - 结构化信息处理")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(test_advanced_formats())
|