from typing import List, Dict, Any, TypedDict 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.prompts import STRUCTURED_TOOL_DEMAND_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]] knowledgeType: str 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.task_dao = QueryTaskDAO() # 创建状态图 self.graph = self._create_graph() def _create_graph(self) -> StateGraph: """创建LangGraph状态图""" workflow = StateGraph(AgentState) # 添加节点（仅保留 生成 与 保存） workflow.add_node("generate_initial_queries", self._generate_initial_queries) workflow.add_node("save_queries", self._save_queries) # 设置入口点 workflow.set_entry_point("generate_initial_queries") # 添加边 workflow.add_edge("generate_initial_queries", "save_queries") workflow.add_edge("save_queries", END) return workflow.compile() def _generate_initial_queries(self, state: AgentState) -> AgentState: """生成 refined_queries（从结构化JSON中聚合三类关键词）""" question = state["question"] # 使用新的结构化系统提示 prompt = ChatPromptTemplate.from_messages([ SystemMessage(content=STRUCTURED_TOOL_DEMAND_PROMPT), HumanMessage(content=question) ]) try: response = self.llm.invoke(prompt.format_messages()) text = (response.content or "").strip() # 解析严格的JSON数组；若失败，尝试从文本中提取 try: data = json.loads(text) except Exception: data = self._extract_json_array_from_text(text) logger.info(f"需求分析结果: {data}") aggregated: List[str] = [] for item in data: ek = (item or {}).get("expanded_keywords", {}) g = ek.get("general_discovery_queries", []) or [] t = ek.get("themed_function_queries", []) or [] h = ek.get("how_to_use_queries", []) or [] for q in [*g, *t, *h]: q_str = str(q).strip() if q_str: aggregated.append(q_str) # 去重，保持顺序 seen = set() deduped: List[str] = [] for q in aggregated: if q not in seen: seen.add(q) deduped.append(q) state["initial_queries"] = deduped state["refined_queries"] = deduped except Exception as e: logger.warning(f"结构化需求解析失败，降级为原始问题: {e}") state["initial_queries"] = [question] state["refined_queries"] = [question] return state # 删除 refine/validate/classify 节点 def _save_queries(self, state: AgentState) -> AgentState: """保存查询词到外部接口节点""" refined_queries = state.get("refined_queries", []) question = state.get("question", "") knowledge_type = state.get("knowledgeType", "") or "内容知识" if not refined_queries: logger.warning("没有查询词需要保存") return state # 合并 knowledgeType 与每个查询词，形成提交数据 result_items: List[Dict[str, str]] = [ {"query": q, "knowledgeType": knowledge_type} for q in refined_queries ] state["result_queries"] = result_items try: url = "http://aigc-testapi.cybertogether.net/aigc/agent/knowledgeWorkflow/addQuery" headers = {"Content-Type": "application/json"} with httpx.Client() as client: data_content = result_items logger.info(f"查询词保存数据: {data_content}") resp1 = client.post(url, headers=headers, json=data_content, timeout=30) resp1.raise_for_status() logger.info(f"查询词保存结果: {resp1.text}") 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, knowledge_type: str = "") -> List[str]: """ 生成查询词的主入口 Args: question: 用户问题 task_id: 任务ID Returns: 生成的查询词列表 """ initial_state = { "question": question, "task_id": task_id, "initial_queries": [], "refined_queries": [], "result_queries": [], "knowledgeType": knowledge_type or "内容知识" } 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]