#!/usr/bin/env python3 """ 简单测试示例 - 基础RAG功能验证 """ 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 SimpleTestRAG(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_basic_functionality(): """测试基础RAG功能""" print("🔧 基础RAG功能测试") print("=" * 50) # 清理向量数据库 db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/simple_test") if db_path.exists(): shutil.rmtree(db_path) print("🧹 已清理向量数据库") # 创建RAG实例 - 禁用图片处理用于基础测试 rag = SimpleTestRAG( vector_store_name="simple_test", 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/simple_test_status.db", image_config={"enabled": False} # 基础测试禁用图片 ) print("✅ RAG实例创建成功") print() # 测试基础文档 test_files = ["test_document.txt", "test_markdown.md", "python_basics.txt", "data_science.txt"] print("📂 处理基础文档...") processed_count = 0 for filename in test_files: file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename if not file_path.exists(): print(f"⚠️ {filename} - 文件不存在,跳过") continue print(f"📄 处理: {filename}") try: result = await rag.ingest(str(file_path)) if result and result.get('success'): print(f" ✅ 成功: {result['chunks_count']} 个片段") processed_count += 1 else: message = result.get('message', '未知错误') if "已经处理完毕" in message: print(f" ⚠️ 已存在,跳过") processed_count += 1 else: print(f" ❌ 失败: {message}") except Exception as e: print(f" ❌ 错误: {str(e)}") print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件") print() # 基础查询测试 print("🔍 基础查询测试...") test_queries = [ "Python编程语言的特点", "数据科学的核心技术", "机器学习的应用", "什么是深度学习" ] for i, question in enumerate(test_queries, 1): print(f"\n❓ 查询 {i}: {question}") try: answer = await rag.query(question) if "抱歉" not in answer: # 显示结果摘要 lines = answer.split('\n') source_line = lines[0] if lines[0].startswith('基于文档') else "来源未知" print(f" 📚 {source_line}") # 显示内容预览 content_start = answer.find('\n\n') if content_start > 0: content = answer[content_start+2:] preview = content[:150] + "..." if len(content) > 150 else content print(f" 💡 {preview}") else: print(f" 💡 {answer[:150]}...") else: print(f" 💡 {answer}") except Exception as e: print(f" ❌ 查询失败: {str(e)}") print("\n" + "=" * 50) print("🎉 基础功能测试完成!") print("✅ 验证项目:") print(" 📄 文档加载和切分") print(" 🔍 文本向量化和存储") print(" 🔎 相似性搜索") print(" 📝 查询结果整合") if __name__ == "__main__": asyncio.run(test_basic_functionality())