#!/usr/bin/env python3 """ 文件格式支持综合测试 - 测试 txt、md、docx 文件 """ import asyncio import sys import os from pathlib import Path # 添加项目路径 sys.path.append('/Users/liruwei/Documents/code/project/demo/base_rag/src') from base_rag.core import BaseRAG class SimpleRAG(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", "未知来源") if source not in sources: sources.append(source) contexts.append(doc.page_content.strip()) context = "\n\n".join(contexts) sources_str = "、".join(sources) return f"基于以下文档({sources_str})的信息:\n\n{context}" async def comprehensive_file_format_test(): """全面的文件格式测试""" print("=" * 60) print("🚀 Base RAG 异步文件格式支持测试") print("=" * 60) print() # 初始化RAG系统 rag = SimpleRAG( vector_store_name="comprehensive_test_kb", retriever_top_k=3, persist_directory="/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db", status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/test_status.db" ) # 支持的文件格式和对应的测试文件 test_files = [ { "format": "TXT", "path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/knowledge.txt", "description": "纯文本文件 - 基础知识" }, { "format": "MD", "path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python_guide.md", "description": "Markdown文件 - Python编程指南" }, { "format": "MD", "path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/machine_learning.md", "description": "Markdown文件 - 机器学习入门" }, { "format": "DOCX", "path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx", "description": "Word文档 - 深度学习技术指南" } ] print("📂 文件处理测试\n") # 处理每个文件 processed_files = [] for file_info in test_files: file_path = file_info["path"] format_name = file_info["format"] description = file_info["description"] if not Path(file_path).exists(): print(f"❌ {format_name}: {description} - 文件不存在") continue print(f"📄 处理 {format_name} 格式: {description}") try: result = await rag.process_file_to_vector_store(file_path) if result and result.get('success'): print(f" ✅ 成功: {result['chunks_count']} 个文档片段") processed_files.append(file_info) else: print(f" ⚠️ 跳过: {result.get('message', '文件可能已存在')}") processed_files.append(file_info) except Exception as e: print(f" ❌ 失败: {str(e)}") print() print("=" * 60) print("💬 跨格式查询测试") print("=" * 60) print() # 测试跨格式查询 test_queries = [ { "question": "Python的基本语法有哪些?", "expected_format": "MD (Python指南)" }, { "question": "什么是机器学习?有哪些类型?", "expected_format": "MD (机器学习)" }, { "question": "深度学习有哪些核心组件?", "expected_format": "DOCX (深度学习)" }, { "question": "深度学习的应用领域都有哪些?", "expected_format": "DOCX (深度学习)" } ] for i, query_info in enumerate(test_queries, 1): question = query_info["question"] expected = query_info["expected_format"] print(f"❓ 查询 {i}: {question}") print(f" 期望来源: {expected}") try: response = await rag.query(question) print(f" 💡 回答: {response[:200]}...") print() except Exception as e: print(f" ❌ 查询失败: {str(e)}") print() print("=" * 60) print("📊 测试总结") print("=" * 60) format_stats = {} for file_info in processed_files: format_name = file_info["format"] if format_name not in format_stats: format_stats[format_name] = 0 format_stats[format_name] += 1 print(f"✅ 成功处理的文件格式:") for format_name, count in format_stats.items(): print(f" • {format_name}: {count} 个文件") print(f"\n🎉 异步RAG系统文件格式支持测试完成!") print(f" 支持格式: TXT, MD, DOCX") print(f" 异步处理: ✅ 完全支持") print(f" 跨格式查询: ✅ 完全支持") if __name__ == "__main__": asyncio.run(comprehensive_file_format_test())