# agent/citu_agent.py from typing import Dict, Any, Literal from langgraph.graph import StateGraph, END from langchain.agents import AgentExecutor, create_openai_tools_agent from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.messages import SystemMessage, HumanMessage from agent.state import AgentState from agent.classifier import QuestionClassifier from agent.tools import TOOLS, generate_sql, execute_sql, generate_summary, general_chat from agent.utils import get_compatible_llm class CituLangGraphAgent: """Citu LangGraph智能助手主类 - 使用@tool装饰器 + Agent工具调用""" def __init__(self): # 加载配置 try: from agent.config import get_current_config, get_nested_config self.config = get_current_config() print("[CITU_AGENT] 加载Agent配置完成") except ImportError: self.config = {} print("[CITU_AGENT] 配置文件不可用，使用默认配置") self.classifier = QuestionClassifier() self.tools = TOOLS self.llm = get_compatible_llm() # 预创建Agent实例以提升性能 enable_reuse = self.config.get("performance", {}).get("enable_agent_reuse", True) if enable_reuse: print("[CITU_AGENT] 预创建Agent实例中...") self._database_executor = self._create_database_agent() self._chat_executor = self._create_chat_agent() print("[CITU_AGENT] Agent实例预创建完成") else: self._database_executor = None self._chat_executor = None print("[CITU_AGENT] Agent实例重用已禁用，将在运行时创建") self.workflow = self._create_workflow() print("[CITU_AGENT] LangGraph Agent with Tools初始化完成") def _create_workflow(self) -> StateGraph: """创建LangGraph工作流""" workflow = StateGraph(AgentState) # 添加节点 workflow.add_node("classify_question", self._classify_question_node) workflow.add_node("agent_chat", self._agent_chat_node) workflow.add_node("agent_database", self._agent_database_node) workflow.add_node("format_response", self._format_response_node) # 设置入口点 workflow.set_entry_point("classify_question") # 添加条件边：分类后的路由 # 完全信任QuestionClassifier的决策，不再进行二次判断 workflow.add_conditional_edges( "classify_question", self._route_after_classification, { "DATABASE": "agent_database", "CHAT": "agent_chat" # CHAT分支处理所有非DATABASE的情况（包括UNCERTAIN） } ) # 添加边 workflow.add_edge("agent_chat", "format_response") workflow.add_edge("agent_database", "format_response") workflow.add_edge("format_response", END) return workflow.compile() def _classify_question_node(self, state: AgentState) -> AgentState: """问题分类节点""" try: print(f"[CLASSIFY_NODE] 开始分类问题: {state['question']}") classification_result = self.classifier.classify(state["question"]) # 更新状态 state["question_type"] = classification_result.question_type state["classification_confidence"] = classification_result.confidence state["classification_reason"] = classification_result.reason state["classification_method"] = classification_result.method state["current_step"] = "classified" state["execution_path"].append("classify") print(f"[CLASSIFY_NODE] 分类结果: {classification_result.question_type}, 置信度: {classification_result.confidence}") return state except Exception as e: print(f"[ERROR] 问题分类异常: {str(e)}") state["error"] = f"问题分类失败: {str(e)}" state["error_code"] = 500 state["execution_path"].append("classify_error") return state def _create_database_agent(self): """创建数据库专用Agent（预创建）""" from agent.config import get_nested_config # 获取配置 max_iterations = get_nested_config(self.config, "database_agent.max_iterations", 5) enable_verbose = get_nested_config(self.config, "database_agent.enable_verbose", True) early_stopping_method = get_nested_config(self.config, "database_agent.early_stopping_method", "generate") database_prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=""" 你是一个专业的数据库查询助手。你的任务是帮助用户查询数据库并生成报告。 工具使用流程： 1. 首先使用 generate_sql 工具将用户问题转换为SQL 2. 然后使用 execute_sql 工具执行SQL查询 3. 最后使用 generate_summary 工具为结果生成自然语言摘要 如果任何步骤失败，请提供清晰的错误信息并建议解决方案。 """), MessagesPlaceholder(variable_name="chat_history", optional=True), HumanMessage(content="{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ]) database_tools = [generate_sql, execute_sql, generate_summary] agent = create_openai_tools_agent(self.llm, database_tools, database_prompt) return AgentExecutor( agent=agent, tools=database_tools, verbose=enable_verbose, handle_parsing_errors=True, max_iterations=max_iterations, early_stopping_method=early_stopping_method ) def _agent_database_node(self, state: AgentState) -> AgentState: """数据库Agent节点 - 使用预创建或动态创建的Agent""" try: print(f"[DATABASE_AGENT] 开始处理数据库查询: {state['question']}") # 使用预创建的Agent或动态创建 if self._database_executor is not None: executor = self._database_executor print(f"[DATABASE_AGENT] 使用预创建的Agent实例") else: executor = self._create_database_agent() print(f"[DATABASE_AGENT] 动态创建Agent实例") # 执行Agent result = executor.invoke({ "input": state["question"] }) # 解析Agent执行结果 self._parse_database_agent_result(state, result) state["current_step"] = "database_completed" state["execution_path"].append("agent_database") print(f"[DATABASE_AGENT] 数据库查询完成") return state except Exception as e: print(f"[ERROR] 数据库Agent异常: {str(e)}") state["error"] = f"数据库查询失败: {str(e)}" state["error_code"] = 500 state["current_step"] = "database_error" state["execution_path"].append("agent_database_error") return state def _create_chat_agent(self): """创建聊天专用Agent（预创建）""" from agent.config import get_nested_config # 获取配置 max_iterations = get_nested_config(self.config, "chat_agent.max_iterations", 3) enable_verbose = get_nested_config(self.config, "chat_agent.enable_verbose", True) chat_prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=""" 你是Citu智能数据问答平台的友好助手。 使用 general_chat 工具来处理用户的一般性问题、概念解释、操作指导等。 特别注意： - 如果用户的问题可能涉及数据查询，建议他们尝试数据库查询功能 - 如果问题不够明确，主动询问更多细节以便更好地帮助用户 - 对于模糊的问题，可以提供多种可能的解决方案 - 当遇到不确定的问题时，通过友好的对话来澄清用户意图 """), MessagesPlaceholder(variable_name="chat_history", optional=True), HumanMessage(content="{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ]) chat_tools = [general_chat] agent = create_openai_tools_agent(self.llm, chat_tools, chat_prompt) return AgentExecutor( agent=agent, tools=chat_tools, verbose=enable_verbose, handle_parsing_errors=True, max_iterations=max_iterations ) def _agent_chat_node(self, state: AgentState) -> AgentState: """聊天Agent节点 - 使用预创建或动态创建的Agent""" try: print(f"[CHAT_AGENT] 开始处理聊天: {state['question']}") # 使用预创建的Agent或动态创建 if self._chat_executor is not None: executor = self._chat_executor print(f"[CHAT_AGENT] 使用预创建的Agent实例") else: executor = self._create_chat_agent() print(f"[CHAT_AGENT] 动态创建Agent实例") # 构建上下文 enable_context_injection = self.config.get("chat_agent", {}).get("enable_context_injection", True) context = None if enable_context_injection and state.get("classification_reason"): context = f"分类原因: {state['classification_reason']}" # 执行Agent input_text = state["question"] if context: input_text = f"{state['question']}\n\n上下文: {context}" result = executor.invoke({ "input": input_text }) # 提取聊天响应 state["chat_response"] = result.get("output", "") state["current_step"] = "chat_completed" state["execution_path"].append("agent_chat") print(f"[CHAT_AGENT] 聊天处理完成") return state except Exception as e: print(f"[ERROR] 聊天Agent异常: {str(e)}") state["chat_response"] = "抱歉，我暂时无法处理您的问题。请稍后再试，或者尝试询问数据相关的问题。" state["current_step"] = "chat_error" state["execution_path"].append("agent_chat_error") return state def _format_response_node(self, state: AgentState) -> AgentState: """格式化最终响应节点""" try: print(f"[FORMAT_NODE] 开始格式化响应，问题类型: {state['question_type']}") state["current_step"] = "completed" state["execution_path"].append("format_response") # 根据问题类型和执行状态格式化响应 if state.get("error"): # 有错误的情况 state["final_response"] = { "success": False, "error": state["error"], "error_code": state.get("error_code", 500), "question_type": state["question_type"], "execution_path": state["execution_path"], "classification_info": { "confidence": state.get("classification_confidence", 0), "reason": state.get("classification_reason", ""), "method": state.get("classification_method", "") } } elif state["question_type"] == "DATABASE": # 数据库查询类型 if state.get("data_result") and state.get("summary"): # 完整的数据库查询流程 state["final_response"] = { "success": True, "response": state["summary"], "type": "DATABASE", "sql": state.get("sql"), "data_result": state["data_result"], "summary": state["summary"], "execution_path": state["execution_path"], "classification_info": { "confidence": state["classification_confidence"], "reason": state["classification_reason"], "method": state["classification_method"] } } else: # 数据库查询失败，但有部分结果 state["final_response"] = { "success": False, "error": state.get("error", "数据库查询未完成"), "type": "DATABASE", "sql": state.get("sql"), "execution_path": state["execution_path"] } else: # 聊天类型 state["final_response"] = { "success": True, "response": state.get("chat_response", ""), "type": "CHAT", "execution_path": state["execution_path"], "classification_info": { "confidence": state["classification_confidence"], "reason": state["classification_reason"], "method": state["classification_method"] } } print(f"[FORMAT_NODE] 响应格式化完成") return state except Exception as e: print(f"[ERROR] 响应格式化异常: {str(e)}") state["final_response"] = { "success": False, "error": f"响应格式化异常: {str(e)}", "error_code": 500, "execution_path": state["execution_path"] } return state def _route_after_classification(self, state: AgentState) -> Literal["DATABASE", "CHAT"]: """ 分类后的路由决策 完全信任QuestionClassifier的决策： - DATABASE类型 → 数据库Agent - CHAT和UNCERTAIN类型 → 聊天Agent 这样避免了双重决策的冲突，所有分类逻辑都集中在QuestionClassifier中 """ question_type = state["question_type"] confidence = state["classification_confidence"] print(f"[ROUTE] 分类路由: {question_type}, 置信度: {confidence} (完全信任分类器决策)") if question_type == "DATABASE": return "DATABASE" else: # 将 "CHAT" 和 "UNCERTAIN" 类型都路由到聊天流程 # 聊天Agent可以处理不确定的情况，并在必要时引导用户提供更多信息 return "CHAT" def _parse_database_agent_result(self, state: AgentState, agent_result: Dict[str, Any]): """解析数据库Agent的执行结果""" try: output = agent_result.get("output", "") intermediate_steps = agent_result.get("intermediate_steps", []) # 从intermediate_steps中提取工具调用结果 for step in intermediate_steps: if len(step) >= 2: action, observation = step[0], step[1] if hasattr(action, 'tool') and hasattr(action, 'tool_input'): tool_name = action.tool tool_result = observation # 解析工具结果 if tool_name == "generate_sql" and isinstance(tool_result, dict): if tool_result.get("success"): state["sql"] = tool_result.get("sql") else: state["error"] = tool_result.get("error") elif tool_name == "execute_sql" and isinstance(tool_result, dict): if tool_result.get("success"): state["data_result"] = tool_result.get("data_result") else: state["error"] = tool_result.get("error") elif tool_name == "generate_summary" and isinstance(tool_result, dict): if tool_result.get("success"): state["summary"] = tool_result.get("summary") # 如果没有从工具结果中获取到摘要，使用Agent的最终输出 if not state.get("summary") and output: state["summary"] = output except Exception as e: print(f"[WARNING] 解析数据库Agent结果失败: {str(e)}") # 使用Agent的输出作为摘要 state["summary"] = agent_result.get("output", "查询处理完成") def process_question(self, question: str, session_id: str = None) -> Dict[str, Any]: """ 统一的问题处理入口 Args: question: 用户问题 session_id: 会话ID Returns: Dict包含完整的处理结果 """ try: print(f"[CITU_AGENT] 开始处理问题: {question}") # 初始化状态 initial_state = self._create_initial_state(question, session_id) # 执行工作流 final_state = self.workflow.invoke( initial_state, config={ "configurable": {"session_id": session_id} } if session_id else None ) # 提取最终结果 result = final_state["final_response"] print(f"[CITU_AGENT] 问题处理完成: {result.get('success', False)}") return result except Exception as e: print(f"[ERROR] Agent执行异常: {str(e)}") return { "success": False, "error": f"Agent系统异常: {str(e)}", "error_code": 500, "execution_path": ["error"] } def _create_initial_state(self, question: str, session_id: str = None) -> AgentState: """创建初始状态""" return AgentState( # 输入信息 question=question, session_id=session_id, # 分类结果 question_type="", classification_confidence=0.0, classification_reason="", classification_method="", # 数据库查询流程状态 sql=None, sql_generation_attempts=0, data_result=None, summary=None, # 聊天响应 chat_response=None, # 最终输出 final_response={}, # 错误处理 error=None, error_code=None, # 流程控制 current_step="start", execution_path=[], retry_count=0, max_retries=2, # 调试信息 debug_info={} ) def health_check(self) -> Dict[str, Any]: """健康检查""" try: # 从配置获取健康检查参数 from agent.config import get_nested_config test_question = get_nested_config(self.config, "health_check.test_question", "你好") enable_full_test = get_nested_config(self.config, "health_check.enable_full_test", True) if enable_full_test: # 完整流程测试 test_result = self.process_question(test_question, "health_check") return { "status": "healthy" if test_result.get("success") else "degraded", "test_result": test_result.get("success", False), "workflow_compiled": self.workflow is not None, "tools_count": len(self.tools), "agent_reuse_enabled": self._database_executor is not None and self._chat_executor is not None, "message": "Agent健康检查完成" } else: # 简单检查 return { "status": "healthy", "test_result": True, "workflow_compiled": self.workflow is not None, "tools_count": len(self.tools), "agent_reuse_enabled": self._database_executor is not None and self._chat_executor is not None, "message": "Agent简单健康检查完成" } except Exception as e: return { "status": "unhealthy", "error": str(e), "workflow_compiled": self.workflow is not None, "tools_count": len(self.tools) if hasattr(self, 'tools') else 0, "agent_reuse_enabled": False, "message": "Agent健康检查失败" }