feat: 增加 AI Agent 自主诊断执行模式

2026-07-25 12:21:54 +08:00
parent 26ae846501
commit 4f3a7833a0
8 changed files with 659 additions and 1 deletions
@@ -1,7 +1,9 @@
 """AI Agent 分析模块 - 调用 OpenAI 兼容接口分析诊断结果 (支持流式输出)"""
 import json
+import uuid
 import httpx
 from backend.config import AI_API_BASE, AI_API_KEY, AI_MODEL
+from backend.agent_tools import classify_kubectl_command, execute_kubectl

 SYSTEM_PROMPT = """你是一位资深的 Kubernetes 运维专家和 SRE 工程师。你的任务是分析 K8S 集群诊断报告，找出根本原因并给出修复建议。

@@ -20,6 +22,143 @@ _HEADERS = {
 }


+AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你可以根据诊断报告决定执行 kubectl 命令获取更多证据，并在证据充分后给出结论。
+
+安全规则：
+1. 只允许提出单条 kubectl 命令，不得使用 shell、管道、重定向或命令连接符。
+2. get、describe、logs、top 等只读命令可以自动执行。
+3. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准。
+4. 不要猜测命令输出；必须依据实际工具结果继续判断。
+5. 最多执行有限轮次，优先使用 namespace 和资源名缩小范围。
+
+每轮必须只返回一个 JSON 对象，不要使用 Markdown 代码块：
+- 继续调查：{"action":"command","command":"kubectl ...","reason":"为什么执行"}
+- 完成分析：{"action":"final","analysis":"中文 Markdown 分析结论"}
+"""
+
+_PENDING_APPROVALS: dict[str, dict] = {}
+
+
+def _parse_agent_action(content: str) -> dict:
+    try:
+        action = json.loads(content.strip())
+    except json.JSONDecodeError as exc:
+        raise ValueError("AI 未返回合法 JSON 动作") from exc
+    if not isinstance(action, dict) or action.get("action") not in {"command", "final"}:
+        raise ValueError("AI 返回了不支持的动作")
+    if action["action"] == "command" and not action.get("command"):
+        raise ValueError("AI 命令动作缺少 command")
+    if action["action"] == "final" and not action.get("analysis"):
+        raise ValueError("AI 最终动作缺少 analysis")
+    return action
+
+
+async def _complete_chat(messages: list[dict]) -> str:
+    timeout = httpx.Timeout(90, connect=10)
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        response = await client.post(
+            f"{AI_API_BASE}/chat/completions",
+            headers=_HEADERS,
+            json={
+                "model": AI_MODEL,
+                "messages": messages,
+                "temperature": 0.1,
+                "max_tokens": 4096,
+                "stream": False,
+            },
+        )
+        response.raise_for_status()
+        payload = response.json()
+        return payload["choices"][0]["message"]["content"]
+
+
+async def run_diagnostic_agent(
+    report: str,
+    question: str = "",
+    max_steps: int = 6,
+    execution_mode: str = "ai",
+):
+    """自主诊断循环：AI 模式执行只读命令，手动模式等待用户执行。"""
+    if execution_mode not in {"ai", "manual"}:
+        yield {"type": "error", "message": "不支持的命令执行模式，仅支持 ai 或 manual"}
+        return
+
+    user_content = f"当前 K8S 诊断报告：\n\n{report}"
+    if question:
+        user_content += f"\n\n用户关注点：{question}"
+    messages = [
+        {"role": "system", "content": AGENT_SYSTEM_PROMPT},
+        {"role": "user", "content": user_content},
+    ]
+
+    for step in range(1, max_steps + 1):
+        yield {"type": "thinking", "step": step, "message": f"AI 正在规划第 {step} 轮诊断"}
+        try:
+            content = await _complete_chat(messages)
+            action = _parse_agent_action(content)
+        except Exception as exc:
+            yield {"type": "error", "message": str(exc)}
+            return
+
+        if action["action"] == "final":
+            yield {"type": "final", "content": action["analysis"], "model": AI_MODEL}
+            return
+
+        command = action["command"]
+        reason = action.get("reason", "AI 需要更多集群证据")
+        try:
+            decision = classify_kubectl_command(command)
+        except ValueError as exc:
+            messages.extend([
+                {"role": "assistant", "content": content},
+                {"role": "user", "content": f"命令被安全策略拒绝：{exc}。请改用单条合法 kubectl 命令。"},
+            ])
+            yield {"type": "command_rejected", "command": command, "reason": str(exc)}
+            continue
+
+        yield {
+            "type": "command", "step": step, "command": command,
+            "reason": reason, "mode": decision.mode,
+        }
+        if execution_mode == "manual":
+            yield {
+                "type": "manual_command", "command_id": uuid.uuid4().hex,
+                "command": command, "reason": reason, "mode": decision.mode,
+                "message": "请在终端手动执行命令，并根据结果继续排查",
+            }
+            return
+
+        if decision.requires_approval:
+            approval_id = uuid.uuid4().hex
+            _PENDING_APPROVALS[approval_id] = {
+                "command": command,
+                "reason": reason,
+                "created_by": AI_MODEL,
+            }
+            yield {
+                "type": "approval_required", "approval_id": approval_id,
+                "command": command, "reason": reason,
+            }
+            return
+
+        result = await execute_kubectl(command)
+        yield {"type": "result", **result}
+        messages.extend([
+            {"role": "assistant", "content": content},
+            {"role": "user", "content": "命令实际执行结果：\n" + json.dumps(result, ensure_ascii=False)},
+        ])
+
+    yield {
+        "type": "final",
+        "content": f"已达到自主诊断最大轮次（{max_steps}），请根据上方命令输出继续人工排查。",
+        "model": AI_MODEL,
+    }
+
+
+def take_pending_approval(approval_id: str) -> dict | None:
+    return _PENDING_APPROVALS.pop(approval_id, None)
+
+
 async def _stream_chat(messages: list[dict]):
    """通用流式请求 - yield 文本 chunk"""
    # connect=10s 快速失败; read=90s 防止 AI 服务无响应时前端无限卡死