""" 主应用程序入口 集成所有模块，构建工作流并运行自适应RAG系统 """ import time from langgraph.graph import END, StateGraph, START from pprint import pprint from config import setup_environment, validate_api_keys, ENABLE_GRAPHRAG from document_processor import initialize_document_processor from routers_and_graders import initialize_graders_and_router from workflow_nodes import WorkflowNodes, GraphState try: from knowledge_graph import initialize_knowledge_graph, initialize_community_summarizer from graph_retriever import initialize_graph_retriever except ImportError: print("⚠️ 无法导入知识图谱模块，GraphRAG功能将不可用") ENABLE_GRAPHRAG = False class AdaptiveRAGSystem: """自适应RAG系统主类""" def __init__(self): print("初始化自适应RAG系统...") # 设置环境和验证API密钥 try: setup_environment() validate_api_keys() # 验证API密钥是否正确设置 print("✅ API密钥验证成功") except ValueError as e: print(f"❌ {e}") raise # 检查 Ollama 服务是否运行 print("🔍 检查 Ollama 服务状态...") if not self._check_ollama_service(): print("\n" + "="*60) print("❌ Ollama 服务未启动！") print("="*60) print("\n请先启动 Ollama 服务：") print("\n方法1: 在终端运行") print(" $ ollama serve") print("\n方法2: 在 Kaggle Notebook 中运行") print(" import subprocess") print(" subprocess.Popen(['ollama', 'serve'])") print("\n方法3: 使用快捷脚本") print(" %run KAGGLE_LOAD_OLLAMA.py") print("="*60) raise ConnectionError("Ollama 服务未运行，请先启动服务") print("✅ Ollama 服务运行正常") # 初始化文档处理器 print("设置文档处理器...") self.doc_processor, self.vectorstore, self.retriever, self.doc_splits = initialize_document_processor() # 初始化评分器和路由器 print("初始化评分器和路由器...") self.graders = initialize_graders_and_router() # 初始化知识图谱 (如果启用) self.graph_retriever = None if ENABLE_GRAPHRAG: print("初始化 GraphRAG...") try: kg = initialize_knowledge_graph() # 尝试加载已有的图谱数据 try: kg.load_from_file("knowledge_graph.json") except FileNotFoundError: print(" 未找到 existing knowledge_graph.json, 将使用空图谱") self.graph_retriever = initialize_graph_retriever(kg) print("✅ GraphRAG 初始化成功") except Exception as e: print(f"⚠️ GraphRAG 初始化失败: {e}") # 初始化工作流节点 print("设置工作流节点...") # WorkflowNodes 将在 _build_workflow 中初始化 # 构建工作流 print("构建工作流图...") self.app = self._build_workflow() print("✅ 自适应RAG系统初始化完成！") def _check_ollama_service(self) -> bool: """检查 Ollama 服务是否运行""" import requests try: # 尝试连接 Ollama API response = requests.get('http://localhost:11434/api/tags', timeout=2) return response.status_code == 200 except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): return False def _build_workflow(self): """构建工作流图""" # 创建工作流节点实例，传递DocumentProcessor实例和retriever self.workflow_nodes = WorkflowNodes( doc_processor=self.doc_processor, graders=self.graders, retriever=self.retriever ) workflow = StateGraph(GraphState) # 定义节点 workflow.add_node("web_search", self.workflow_nodes.web_search) workflow.add_node("retrieve", self.workflow_nodes.retrieve) workflow.add_node("grade_documents", self.workflow_nodes.grade_documents) workflow.add_node("generate", self.workflow_nodes.generate) workflow.add_node("transform_query", self.workflow_nodes.transform_query) workflow.add_node("decompose_query", self.workflow_nodes.decompose_query) workflow.add_node("prepare_next_query", self.workflow_nodes.prepare_next_query) # 构建图 workflow.add_conditional_edges( START, self.workflow_nodes.route_question, { "web_search": "web_search", "vectorstore": "decompose_query", # 向量检索前先进行查询分解 }, ) workflow.add_edge("web_search", "generate") workflow.add_edge("decompose_query", "retrieve") workflow.add_edge("retrieve", "grade_documents") workflow.add_conditional_edges( "grade_documents", self.workflow_nodes.decide_to_generate, { "transform_query": "transform_query", "prepare_next_query": "prepare_next_query", "generate": "generate", "web_search": "web_search", # 添加 web_search 作为回退选项 }, ) workflow.add_edge("transform_query", "retrieve") workflow.add_edge("prepare_next_query", "retrieve") workflow.add_conditional_edges( "generate", self.workflow_nodes.grade_generation_v_documents_and_question, { "not supported": "transform_query", # 修复：有幻觉时重新转换查询，而不是再次生成 "useful": END, "not useful": "transform_query", }, ) # 编译（设置递归限制以防止无限循环） return workflow.compile( checkpointer=None, interrupt_before=None, interrupt_after=None, debug=False ) async def query(self, question: str, verbose: bool = True): """ 处理查询 (异步版本) Args: question (str): 用户问题 verbose (bool): 是否显示详细输出 Returns: dict: 包含最终答案和评估指标的字典 """ import asyncio print(f"\n🔍 处理问题: {question}") print("=" * 50) inputs = {"question": question, "retry_count": 0} # 初始化重试计数器 final_generation = None retrieval_metrics = None # 设置配置，增加递归限制 config = {"recursion_limit": 50} # 增加到 50，默认是 25 print("\n🤖 思考过程:") async for output in self.app.astream(inputs, config=config): for key, value in output.items(): if verbose: # 简单的节点执行提示，模拟流式感 print(f" ↳ 执行节点: {key}...", end="\r") # 异步暂停 await asyncio.sleep(0.1) print(f" ✅ 完成节点: {key} ") final_generation = value.get("generation", final_generation) # 保存检索评估指标 if "retrieval_metrics" in value: retrieval_metrics = value["retrieval_metrics"] print("\n" + "=" * 50) print("🎯 最终答案:") print("-" * 30) # 模拟流式输出效果 (打字机效果) if final_generation: import sys for char in final_generation: sys.stdout.write(char) sys.stdout.flush() # 异步暂停 await asyncio.sleep(0.01) # 控制打字速度 print() # 换行 else: print("未生成答案") print("=" * 50) # 返回包含答案和评估指标的字典 return { "answer": final_generation, "retrieval_metrics": retrieval_metrics } def interactive_mode(self): """交互模式，允许用户持续提问""" import asyncio print("\n🤖 欢迎使用自适应RAG系统!") print("💡 输入问题开始对话，输入 'quit' 或 'exit' 退出") print("-" * 50) while True: try: question = input("\n❓ 请输入您的问题: ").strip() if question.lower() in ['quit', 'exit', '退出', 'q']: print("👋 感谢使用，再见!") break if not question: print("⚠️ 请输入一个有效的问题") continue # 使用 asyncio.run 执行异步查询 result = asyncio.run(self.query(question)) # 显示检索评估摘要 if result.get("retrieval_metrics"): metrics = result["retrieval_metrics"] print("\n📊 检索评估摘要:") print(f" - 检索耗时: {metrics.get('latency', 0):.4f}秒") print(f" - 检索文档数: {metrics.get('retrieved_docs_count', 0)}") print(f" - Precision@3: {metrics.get('precision_at_3', 0):.4f}") print(f" - Recall@3: {metrics.get('recall_at_3', 0):.4f}") print(f" - MAP: {metrics.get('map_score', 0):.4f}") except KeyboardInterrupt: print("\n👋 感谢使用，再见!") break except Exception as e: print(f"❌ 发生错误: {e}") import traceback traceback.print_exc() print("请重试或输入 'quit' 退出") def main(): """主函数""" import asyncio try: # 初始化系统 rag_system: AdaptiveRAGSystem = AdaptiveRAGSystem() # 测试查询 # test_question = "AlphaCodium论文讲的是什么？" test_question = "LangGraph的作者目前在哪家公司工作？" # test_question = "解释embedding嵌入的原理，最好列举实现过程的具体步骤" # 使用 asyncio.run 执行异步查询 result = asyncio.run(rag_system.query(test_question)) # 显示测试查询的检索评估摘要 if result.get("retrieval_metrics"): metrics = result["retrieval_metrics"] print("\n📊 测试查询检索评估摘要:") print(f" - 检索耗时: {metrics.get('latency', 0):.4f}秒") print(f" - 检索文档数: {metrics.get('retrieved_docs_count', 0)}") print(f" - Precision@3: {metrics.get('precision_at_3', 0):.4f}") print(f" - Recall@3: {metrics.get('recall_at_3', 0):.4f}") print(f" - MAP: {metrics.get('map_score', 0):.4f}") # 启动交互模式 rag_system.interactive_mode() except Exception as e: print(f"❌ 系统初始化失败: {e}") import traceback traceback.print_exc() print("请检查配置和依赖是否正确安装") if __name__ == "__main__": main()