from typing import List, Dict, Any, TypedDict, Annotated from langgraph.graph import StateGraph, END from langchain_google_genai import ChatGoogleGenerativeAI from langchain.prompts import ChatPromptTemplate from langchain.schema import HumanMessage, SystemMessage import httpx import json from ..tools.query_tool import SuggestQueryTool from ..tools.prompts import QUERY_GENERATION_PROMPT, QUERY_REFINEMENT_PROMPT from ..database.models import QueryTaskDAO, QueryTaskStatus, logger class AgentState(TypedDict): """Agent状态定义""" question: str task_id: int initial_queries: List[str] refined_queries: List[str] result_queries: List[Dict[str, str]] context: str iteration_count: int class QueryGenerationAgent: """查询词生成Agent""" def __init__(self, gemini_api_key: str, model_name: str = "gemini-1.5-pro"): """ 初始化Agent Args: gemini_api_key: Gemini API密钥 model_name: 使用的模型名称 """ self.llm = ChatGoogleGenerativeAI( google_api_key=gemini_api_key, model=model_name, temperature=0.7 ) self.query_tool = SuggestQueryTool() self.task_dao = QueryTaskDAO() # 创建状态图 self.graph = self._create_graph() def _create_graph(self) -> StateGraph: """创建LangGraph状态图""" workflow = StateGraph(AgentState) # 添加节点 workflow.add_node("analyze_question", self._analyze_question) workflow.add_node("generate_initial_queries", self._generate_initial_queries) workflow.add_node("refine_queries", self._refine_queries) workflow.add_node("validate_queries", self._validate_queries) workflow.add_node("classify_queries", self._classify_queries) workflow.add_node("save_queries", self._save_queries) # 设置入口点 workflow.set_entry_point("analyze_question") # 添加边 workflow.add_edge("analyze_question", "generate_initial_queries") workflow.add_edge("generate_initial_queries", "refine_queries") workflow.add_edge("refine_queries", "validate_queries") workflow.add_edge("validate_queries", "classify_queries") workflow.add_edge("classify_queries", "save_queries") workflow.add_edge("save_queries", END) return workflow.compile() def _analyze_question(self, state: AgentState) -> AgentState: """分析问题节点""" question = state["question"] # 分析问题的复杂度和类型 analysis_prompt = ChatPromptTemplate.from_messages([ SystemMessage(content="你是一个问题分析专家。请分析用户问题的类型和复杂度。"), HumanMessage(content=f"请分析这个问题：{question}\n\n分析要点：1.问题类型 2.复杂度 3.关键词 4.需要的查询角度") ]) try: response = self.llm.invoke(analysis_prompt.format_messages()) logger.info(f"问题分析结果: {response.content}") context = response.content except Exception as e: context = f"问题分析失败: {str(e)}" state["context"] = context state["iteration_count"] = 0 return state def _generate_initial_queries(self, state: AgentState) -> AgentState: """生成初始查询词节点""" question = state["question"] task_id = state["task_id"] context = state.get("context", "") # 使用工具生成查询词 try: initial_queries = self.query_tool._run(question, context, task_id) except Exception as e: # 如果工具失败，使用LLM生成 prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=QUERY_GENERATION_PROMPT), HumanMessage(content=question) ]) try: response = self.llm.invoke(prompt.format_messages()) queries_text = response.content initial_queries = [q.strip() for q in queries_text.split('\n') if q.strip()] except Exception: initial_queries = [question] # 降级处理 state["initial_queries"] = initial_queries return state def _refine_queries(self, state: AgentState) -> AgentState: """优化查询词节点""" question = state["question"] initial_queries = state["initial_queries"] if not initial_queries: state["refined_queries"] = [question] return state # 使用LLM优化查询词 queries_text = '\n'.join(initial_queries) prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=QUERY_REFINEMENT_PROMPT), HumanMessage(content=f"问题：{question}\n查询词：{queries_text}") ]) try: response = self.llm.invoke(prompt.format_messages()) logger.info(f"查询词优化结果: {response.content}") refined_text = response.content refined_queries = [q.strip() for q in refined_text.split('\n') if q.strip()] except Exception as e: # 如果优化失败，使用原始查询词 refined_queries = initial_queries state["refined_queries"] = refined_queries return state def _validate_queries(self, state: AgentState) -> AgentState: """验证查询词节点""" refined_queries = state["refined_queries"] # 基本验证：去重、过滤空字符串、限制长度 validated_queries = [] seen = set() for query in refined_queries: if query and len(query.strip()) > 0 and len(query.strip()) < 100: if query.strip() not in seen: validated_queries.append(query.strip()) seen.add(query.strip()) # 限制最终数量 if len(validated_queries) > 10: validated_queries = validated_queries[:10] # 确保至少有一个查询词 if not validated_queries: validated_queries = [state["question"]] logger.info(f"查询词验证结果: {validated_queries}") state["refined_queries"] = validated_queries return state def _classify_queries(self, state: AgentState) -> AgentState: """推测每个查询词的知识类型并写入result_queries""" refined_queries = state.get("refined_queries", []) # 使用大模型进行分类 result_items: List[Dict[str, str]] = self._classify_with_llm(refined_queries) state["result_queries"] = result_items return state def _save_queries(self, state: AgentState) -> AgentState: """保存查询词到外部接口节点""" refined_queries = state["refined_queries"] question = state["question"] if not refined_queries: logger.warning("没有查询词需要保存") return state # 调用外部接口保存查询词（按类型分组） try: url = "http://aigc-testapi.cybertogether.net/aigc/agent/knowledgeWorkflow/addQuery" headers = {"Content-Type": "application/json"} # 仅使用前一步的分类结果，不做即时分类 result_items: List[Dict[str, str]] = state.get("result_queries", []) if not result_items: logger.warning("缺少分类结果result_queries，跳过外部提交") return state if result_items: with httpx.Client() as client: data_content = {"queryWords": result_items} data_content = json.dumps(data_content) data_content = data_content.encode('utf-8') logger.info(f"查询词保存数据: {data_content}") resp1 = client.post(url, headers=headers, data=data_content, timeout=30) logger.info(f"查询词保存结果: {resp1.json()}") resp1.raise_for_status() logger.info(f"查询词保存成功: question={question}，query数量={len(result_items)}") except httpx.HTTPError as e: logger.error(f"保存查询词时发生HTTP错误: {str(e)}") except Exception as e: logger.error(f"保存查询词时发生错误: {str(e)}") return state def _infer_knowledge_type(self, query: str) -> str: """根据查询词简单推断知识类型（内容知识/工具知识）""" tool_keywords = [ "安装", "配置", "使用", "教程", "API", "SDK", "命令", "指令", "版本", "错误", "异常", "调试", "部署", "集成", "调用", "参数", "示例", "代码", "CLI", "tool", "library", "framework" ] lower_q = query.lower() for kw in tool_keywords: if kw.lower() in lower_q: return "工具知识" return "内容知识" def _classify_with_llm(self, queries: List[str]) -> List[Dict[str, str]]: """调用LLM将查询词分类为 内容知识 / 工具知识。 返回形如 [{"query": q, "knowledgeType": "内容知识"|"工具知识"}, ...] 若解析失败，降级为将所有查询标记为 内容知识（不使用关键词启发）。 """ if not queries: return [] instruction = ( "你是一名分类助手。请将下面的查询词逐一分类为‘内容知识’或‘工具知识’。\n" "请只返回严格的JSON数组，每个元素为对象：{\"query\": 原始查询词, \"knowledgeType\": \"内容知识\" 或 \"工具知识\"}。\n" "不要输出任何解释或多余文本。" ) payload = "\n".join(queries) prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=instruction), HumanMessage(content=f"查询词列表（每行一个）：\n{payload}") ]) try: response = self.llm.invoke(prompt.format_messages()) text = (response.content or "").strip() logger.info(f"LLM分类结果: {text}") # 尝试解析为JSON数组；若失败，尝试从代码块或文本中提取 try: data = json.loads(text) except Exception: data = self._extract_json_array_from_text(text) result: List[Dict[str, str]] = [] for item in data: q = str(item.get("query", "")).strip() kt = str(item.get("knowledgeType", "")).strip() if q and kt in ("内容知识", "工具知识"): result.append({"query": q, "knowledgeType": kt}) # 保证顺序与输入一致，且都包含 if len(result) != len(queries): # 尝试基于输入进行对齐 mapped = {it["query"]: it["knowledgeType"] for it in result} aligned: List[Dict[str, str]] = [] for q in queries: kt = mapped.get(q, "内容知识") aligned.append({"query": q, "knowledgeType": kt}) return aligned return result except Exception as e: # 降级：全部标注为内容知识（不做关键词匹配） logger.warning(f"LLM分类失败，使用降级策略: {e}") return [{"query": q, "knowledgeType": "内容知识"} for q in queries] def _extract_json_array_from_text(self, text: str) -> List[Dict[str, Any]]: """尽力从模型输出（可能包含```json代码块或多余文本）中提取JSON数组。""" s = (text or "").strip() # 去除三引号包裹的代码块 if s.startswith("```"): # 去掉第一行的 ``` 或 ```json first_newline = s.find('\n') if first_newline != -1: s = s[first_newline + 1:] if s.endswith("```"): s = s[:-3] s = s.strip() # 在文本中查找首个JSON数组 import re match = re.search(r"\[[\s\S]*\]", s) if not match: raise ValueError("未找到JSON数组片段") json_str = match.group(0) data = json.loads(json_str) if not isinstance(data, list): raise ValueError("提取内容不是JSON数组") return data async def generate_queries(self, question: str, task_id: int = 0) -> List[str]: """ 生成查询词的主入口 Args: question: 用户问题 task_id: 任务ID Returns: 生成的查询词列表 """ initial_state = { "question": question, "task_id": task_id, "initial_queries": [], "refined_queries": [], "result_queries": [], "context": "", "iteration_count": 0 } try: result = await self.graph.ainvoke(initial_state) return result["result_queries"] except Exception as e: # 更新任务状态为失败 if task_id > 0: self.task_dao.update_task_status(task_id, QueryTaskStatus.FAILED) # 降级处理：返回原始问题 return [question]