cnbugs
323 lines
13 KiB
Python
323 lines
13 KiB
Python
"""AI Agent 分析模块 - 调用 OpenAI 兼容接口分析诊断结果 (支持流式输出)"""
|
|
import json
|
|
import re
|
|
import uuid
|
|
from collections import Counter
|
|
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. 如果证据不足，必须返回 command 动作并逐条调查；至少检查与故障资源直接相关的状态、describe、events 或 logs。
|
|
2. 禁止在 final 中仅提供“请执行以下命令”的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
|
|
3. 每次只提出一条 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:
|
|
text = content.strip()
|
|
candidates = [text]
|
|
fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
|
|
if fenced:
|
|
candidates.insert(0, fenced.group(1))
|
|
first_brace = text.find("{")
|
|
last_brace = text.rfind("}")
|
|
if first_brace >= 0 and last_brace > first_brace:
|
|
candidates.append(text[first_brace:last_brace + 1])
|
|
|
|
action = None
|
|
for candidate in candidates:
|
|
try:
|
|
action = json.loads(candidate)
|
|
break
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if action is None:
|
|
raise ValueError("AI 未返回合法 JSON 动作")
|
|
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 = 30,
|
|
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},
|
|
]
|
|
|
|
command_counts = Counter()
|
|
|
|
if execution_mode == "ai":
|
|
initial_command = "kubectl get nodes -o wide"
|
|
command_counts[initial_command] += 1
|
|
yield {
|
|
"type": "command", "step": 0, "command": initial_command,
|
|
"reason": "建立集群节点状态基线，确保 AI 执行模式实际采集现场证据",
|
|
"mode": "read",
|
|
}
|
|
initial_result = await execute_kubectl(initial_command)
|
|
yield {"type": "result", **initial_result}
|
|
messages.append({
|
|
"role": "user",
|
|
"content": "AI 执行模式已自动采集第一份现场证据：\n" + json.dumps(initial_result, ensure_ascii=False),
|
|
})
|
|
|
|
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 需要更多集群证据")
|
|
command_counts[command] += 1
|
|
if command_counts[command] > 2:
|
|
messages.extend([
|
|
{"role": "assistant", "content": content},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"命令 {command!r} 已重复执行两次，不能再次执行。"
|
|
"请基于已有输出选择不同的只读诊断命令，或在证据充分时返回 final。"
|
|
),
|
|
},
|
|
])
|
|
yield {
|
|
"type": "command_rejected", "command": command,
|
|
"reason": "相同命令已执行两次，请更换诊断方向",
|
|
}
|
|
continue
|
|
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,
|
|
"run_id": "",
|
|
}
|
|
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": "error",
|
|
"message": (
|
|
f"自主诊断已达到安全上限（{max_steps} 轮），但 Agent 仍未形成有效结论。"
|
|
"系统已停止继续执行以避免无限循环，请检查是否重复执行相同命令或 AI 服务响应异常。"
|
|
),
|
|
}
|
|
|
|
|
|
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)}"}
|