#!/usr/bin/env python3 """ 简单的文件处理测试 """ import sys import os 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, FileStatus class SimpleRAG(BaseRAG): """简单的RAG实现示例""" def ingest(self, file_path: str, **kwargs): """实现文档导入逻辑""" return self.process_file_to_vector_store(file_path, **kwargs) def query(self, question: str) -> str: """实现简单的查询逻辑""" docs = self.similarity_search_with_rerank(question, k=2) 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}" def test_file_processing(): print("=== 文件处理功能测试 ===\n") # 创建RAG实例 rag = SimpleRAG( vector_store_name="test_kb", retriever_top_k=2, storage_directory="./test_docs", status_db_path="./test_status.db" ) # 创建测试文件 test_dir = Path("./test_files") test_dir.mkdir(exist_ok=True) # 创建多个不同主题的知识文件 # Python基础知识 python_file = test_dir / "python_basics.txt" python_file.write_text(""" Python是一种高级编程语言,由Guido van Rossum于1991年创建。 Python具有简洁易读的语法,适合初学者学习编程。 Python是解释型语言,支持面向对象、函数式等多种编程范式。 Python的设计哲学强调代码的可读性和简洁性。 """, encoding="utf-8") # Web框架知识 web_file = test_dir / "web_frameworks.txt" web_file.write_text(""" Flask是一个轻量级的Python Web框架,易于学习和使用。 Django是一个功能丰富的Python Web框架,适合大型项目开发。 FastAPI是现代的Python Web框架,专为构建API而设计。 Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。 """, encoding="utf-8") # 数据科学知识 datascience_file = test_dir / "data_science.txt" datascience_file.write_text(""" NumPy是Python中用于科学计算的基础库,提供多维数组对象。 Pandas是强大的数据分析和处理库,提供DataFrame数据结构。 Matplotlib是Python的绘图库,用于创建静态、动态和交互式图表。 Scikit-learn是机器学习库,提供各种算法和工具。 """, encoding="utf-8") print("1. 处理多个知识文件...") files_to_process = [python_file, web_file, datascience_file] for file_path in files_to_process: result = rag.ingest(str(file_path), chunk_size=200, chunk_overlap=20) print(f"处理 {file_path.name}: {result['message']} (片段数: {result.get('chunks_count', 0)})") print() print("2. 查询测试...") questions = [ "Python是谁创建的?", "Flask和Django有什么区别?", "Pandas是做什么的?", "什么是NumPy?", "FastAPI有什么特点?" ] for question in questions: print(f"问题: {question}") answer = rag.query(question) print(f"回答: {answer[:150]}...") print("-" * 50) print() print("3. 查看文件状态...") files = rag.get_file_processing_status() for file_info in files: print(f"文件: {file_info['filename']} | 状态: {file_info['status']}") print("\n=== 测试完成 ===") if __name__ == "__main__": test_file_processing()