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