#!/usr/bin/env python3 """ 简单的文件处理测试 """ 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, FileStatus 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=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}" async def test_file_processing(): print("=== 文件处理功能测试 ===\n") # 创建RAG实例 rag = SimpleRAG( vector_store_name="test_kb", retriever_top_k=2, storage_directory="./test_files", # 统一使用test_files目录 status_db_path="./status.db", # 统一数据库名称 ) # 使用现有的测试文件 test_dir = Path("./test_files") # 使用已有的测试文件 python_file = test_dir / "python_basics.txt" web_file = test_dir / "web_frameworks.txt" datascience_file = test_dir / "data_science.txt" print("1. 处理多个知识文件...") files_to_process = [python_file, web_file, datascience_file] for file_path in files_to_process: result = await 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 = await rag.query(question) print(f"回答: {answer[:150]}...") print("-" * 50) print() print("3. 查看文件状态...") files = await rag.get_file_processing_status() for file_info in files: print(f"文件: {file_info['filename']} | 状态: {file_info['status']}") print("\n=== 测试完成 ===") if __name__ == "__main__": asyncio.run(test_file_processing())