"""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 集群诊断报告，找出根本原因并给出修复建议。 分析要求： 1. **问题定位**：明确指出每个问题的根本原因（Root Cause），不要只描述表面现象 2. **影响评估**：说明每个问题对业务的潜在影响 3. **修复方案**：给出具体的修复命令或操作步骤（kubectl 命令、YAML 修改等） 4. **优先级排序**：按紧急程度排序，先处理影响最大的问题 5. **预防建议**：给出避免类似问题再次发生的建议 输出格式：使用清晰的 Markdown，包含标题、列表、代码块。语言使用中文。""" _HEADERS = { "Authorization": f"Bearer {AI_API_KEY}", "Content-Type": "application/json", } 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 服务无响应时前端无限卡死 timeout = httpx.Timeout(90, connect=10) async with httpx.AsyncClient(timeout=timeout) as client: async with client.stream( "POST", f"{AI_API_BASE}/chat/completions", headers=_HEADERS, json={ "model": AI_MODEL, "messages": messages, "temperature": 0.3, "max_tokens": 4096, "stream": True, }, ) as resp: # 非 2xx 直接抛异常 (由上层转成 error 事件) resp.raise_for_status() # 有些服务认证失败时返回 200 + JSON 错误体 (非 SSE)，需要识别 ctype = resp.headers.get("content-type", "") if "text/event-stream" not in ctype and "application/json" in ctype: body = await resp.aread() try: err = json.loads(body.decode()) raise httpx.HTTPStatusError( f"AI 服务返回错误: {err.get('error', body.decode()[:300])}", request=resp.request, response=resp, ) except (json.JSONDecodeError, UnicodeDecodeError): raise httpx.HTTPStatusError( f"AI 服务返回非流式响应: {body.decode()[:300]}", request=resp.request, response=resp, ) async for line in resp.aiter_lines(): if not line.startswith("data: "): continue data_str = line[6:].strip() if data_str == "[DONE]": break try: chunk = json.loads(data_str) delta = chunk["choices"][0].get("delta", {}) content = delta.get("content", "") if content: yield content except (json.JSONDecodeError, KeyError, IndexError): continue async def analyze_with_ai_stream(report: str, question: str = ""): """流式 AI 分析诊断报告 - yield 文本 chunk""" user_msg = f"以下是 K8S 集群诊断报告，请进行详细分析：\n\n{report}" if question: user_msg += f"\n\n用户额外问题：{question}" messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_msg}, ] async for chunk in _stream_chat(messages): yield chunk async def chat_with_ai_stream(messages: list[dict], context: str = ""): """流式多轮对话 - yield 文本 chunk""" msgs = [{"role": "system", "content": SYSTEM_PROMPT}] if context: msgs.append({"role": "system", "content": f"当前诊断上下文：\n{context}"}) msgs.extend(messages) async for chunk in _stream_chat(msgs): yield chunk async def analyze_with_ai(report: str, question: str = "") -> dict: """非流式 AI 分析 (兼容旧接口)""" try: content = "" async for chunk in analyze_with_ai_stream(report, question): content += chunk return {"success": True, "analysis": content, "model": AI_MODEL} except httpx.HTTPStatusError as e: return {"success": False, "analysis": f"AI 接口返回错误: {e.response.status_code} - {e.response.text[:500]}", "model": AI_MODEL} except httpx.ConnectError: return {"success": False, "analysis": f"无法连接 AI 服务 ({AI_API_BASE})，请检查 .env 中的 AI_API_BASE 配置", "model": AI_MODEL} except Exception as e: return {"success": False, "analysis": f"AI 分析失败: {str(e)}", "model": AI_MODEL} async def chat_with_ai(messages: list[dict], context: str = "") -> dict: """非流式多轮对话 (兼容旧接口)""" try: content = "" async for chunk in chat_with_ai_stream(messages, context): content += chunk return {"success": True, "reply": content} except Exception as e: return {"success": False, "reply": f"AI 对话失败: {str(e)}"}