knowledge-tools/src/agent/query_agent.py

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]