#!/usr/bin/env python3 """ 多格式文件测试 - 测试 TXT、MD、DOCX 文件格式 """ import sys import os import asyncio import warnings from pathlib import Path # 过滤掉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 MultiFormatRAG(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_multiple_formats(): """测试多种文件格式处理""" print("🚀 多格式文件处理测试") print("=" * 50) # 创建RAG实例 rag = MultiFormatRAG( vector_store_name="multiformat_kb", retriever_top_k=3, storage_directory="../test_files", # 相对于examples目录 status_db_path="../status.db", # 相对于examples目录 ) # 测试文件列表 test_files = [ { "file": "knowledge.txt", "format": "TXT", "description": "纯文本文件" }, { "file": "python_guide.md", "format": "MD", "description": "Markdown文件" }, { "file": "machine_learning.md", "format": "MD", "description": "Markdown文件" }, { "file": "deep_learning_guide.docx", "format": "DOCX", "description": "Word文档" }, { "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": "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(): print(f"❌ {format_type}: {filename} - 文件不存在") continue print(f"📄 处理 {format_type}: {filename} ({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', '可能已存在')}") processed_count += 1 # 已存在也算处理过 except Exception as e: print(f" ❌ 失败: {str(e)}") print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件") print() # 测试跨格式查询 print("💬 跨格式查询测试...") queries = [ "Python有什么特点?", "什么是机器学习?", "深度学习的应用领域有哪些?", "数据科学的核心技术有哪些?", "销售数据中哪个产品销售额最高?", "公司员工的平均年薪是多少?", "人工智能的主要挑战是什么?", "机器学习有哪些类型?" ] for query in queries: print(f"\n❓ {query}") try: answer = await rag.query(query) 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}") # 显示内容摘要(前200字符) if len(content) > 200: content_preview = content[:200] + "..." else: content_preview = content print(f" 💡 {content_preview}") else: print(f" 💡 {answer}") else: print(f" 💡 {answer}") except Exception as e: print(f" ❌ 查询失败: {str(e)}") print("\n" + "=" * 50) print("✅ 多格式文件测试完成!") print("支持的格式: TXT, MD, DOCX, CSV, XLSX, PDF") if __name__ == "__main__": asyncio.run(test_multiple_formats())