base_rag/examples/advanced_format_test.py

208 lines
6.7 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.

#!/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())