cnbugs
32bea0d4ea
扩展 Agent 命令工具，除 kubectl 外新增节点宿主机系统诊断命令支持：
- agent_tools.py: 新增 shell 命令白名单分类器 (classify_shell_command)
- 只读白名单: systemctl status/is-active、journalctl、crictl ps/logs、
docker ps/logs、df、free、ss、ip addr、ps、mount、lsmod、dmesg、
ping、nslookup、cat (限 /proc /sys /etc/kubernetes 等安全路径) 等 30+ 命令
- 写白名单: systemctl restart/stop/start/reload 等需人工审批
- 统一分发器 classify_agent_command 按命令前缀路由
- 统一执行入口 execute_command 走 subprocess_exec (无 shell 注入风险)
- ai_agent.py: 更新 AGENT_SYSTEM_PROMPT 告知 AI 两类可用工具及安全规则
- 三个调用点 (初始探测/分类/执行) 切换到统一接口
- main.py: 审批执行处改用 execute_command
- tests: 新增 19 个用例覆盖 shell 命令分类、混合执行、白名单拒绝
- 全部 44 个测试通过
- 前端/README: 文案与文档补充 shell 命令支持说明
安全策略不变: 禁止 shell 操作符/管道/重定向/多命令拼接，
非白名单命令 (rm/vi/apt 等) 直接拒绝。
432 lines
18 KiB
Python
432 lines
18 KiB
Python
"""AI Agent 分析模块 - 调用 OpenAI 兼容接口分析诊断结果 (支持流式输出)"""
|
||
import json
|
||
import re
|
||
import uuid
|
||
from collections import Counter
|
||
from typing import Optional
|
||
import httpx
|
||
from backend.config import AI_API_BASE, AI_API_KEY, AI_MODEL
|
||
from backend.agent_tools import classify_agent_command, execute_command
|
||
|
||
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。你的首要职责是主动调用命令获取现场证据，而不是只给用户一段脚本或操作建议。
|
||
|
||
可用工具：
|
||
1. kubectl 命令 - 查询集群内资源状态。例如: kubectl get pods -A, kubectl describe node <name>, kubectl logs <pod> -n <ns>
|
||
2. 系统诊断 shell 命令 - 查询节点宿主机状态。例如: systemctl status kubelet, journalctl -u kubelet -n 100 --no-pager, crictl ps, df -h, free -m, ss -tlnp, ip addr
|
||
|
||
行为要求：
|
||
1. 如果证据不足，必须返回 command 动作并逐条调查；至少检查与故障资源直接相关的状态、describe、events 或 logs。
|
||
2. 禁止在 final 中仅提供"请执行以下命令"的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
|
||
3. 每次只提出一条命令，服务端会执行并把真实输出返回给你，然后你继续下一轮判断。
|
||
|
||
安全规则：
|
||
1. 每次只返回一条命令，不得使用 shell、管道、重定向或命令连接符 (; | && > < 等)。
|
||
2. kubectl 的 get、describe、logs、top 等只读命令可以自动执行。
|
||
3. 系统诊断只读命令 (systemctl status、journalctl、crictl ps、df、free、ss、ip addr 等) 可以自动执行。
|
||
4. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准。
|
||
5. systemctl restart/stop/start 等系统修改命令必须暂停并请求人工批准。
|
||
6. 不要猜测命令输出；必须依据实际工具结果继续判断。
|
||
7. 最多执行有限轮次，优先使用 namespace 和资源名缩小范围。
|
||
|
||
每轮必须只返回一个 JSON 对象，不要使用 Markdown 代码块：
|
||
- 继续调查：{"action":"command","command":"kubectl ...","reason":"为什么执行"}
|
||
- 完成分析：{"action":"final","analysis":"中文 Markdown 分析结论","verified":true,"remaining_issues":0}
|
||
|
||
只有完成修复后的全量复检，并确认 remaining_issues 为 0 时才能返回 final。只定位根因、执行命令成功或部分资源恢复都不能结束。
|
||
"""
|
||
|
||
_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")
|
||
if action["action"] == "final":
|
||
if action.get("verified") is not True or action.get("remaining_issues") != 0:
|
||
raise ValueError("AI 最终动作必须确认 verified=true 且 remaining_issues=0")
|
||
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_command(initial_command)
|
||
yield {"type": "result", **initial_result}
|
||
messages.append({
|
||
"role": "user",
|
||
"content": "AI 执行模式已自动采集第一份现场证据：\n" + json.dumps(initial_result, ensure_ascii=False),
|
||
})
|
||
|
||
async for event in _run_agent_loop(
|
||
messages=messages,
|
||
execution_mode=execution_mode,
|
||
max_steps=max_steps,
|
||
start_step=1,
|
||
command_counts=command_counts,
|
||
report=report,
|
||
question=question,
|
||
):
|
||
yield event
|
||
|
||
|
||
async def _run_agent_loop(
|
||
messages: list[dict],
|
||
execution_mode: str,
|
||
max_steps: int,
|
||
start_step: int,
|
||
command_counts: Counter,
|
||
report: str,
|
||
question: str,
|
||
):
|
||
invalid_response_count = 0
|
||
for step in range(start_step, max_steps + 1):
|
||
yield {"type": "thinking", "step": step, "message": f"AI 正在规划第 {step} 轮诊断"}
|
||
try:
|
||
content = await _complete_chat(messages)
|
||
action = _parse_agent_action(content)
|
||
except ValueError as exc:
|
||
if "最终动作必须确认" in str(exc):
|
||
invalid_response_count = 0
|
||
messages.extend([
|
||
{"role": "assistant", "content": content},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"不能结束诊断：尚未通过修复后的全量复检确认所有问题清零。"
|
||
"请继续返回 command，检查节点、Pod、Deployment、事件以及原故障资源；"
|
||
"只有 verified=true 且 remaining_issues=0 时才可返回 final。"
|
||
),
|
||
},
|
||
])
|
||
yield {"type": "verification_required", "message": "尚未确认所有问题清零，Agent 继续执行诊断验证"}
|
||
continue
|
||
invalid_response_count += 1
|
||
if invalid_response_count >= 3:
|
||
yield {
|
||
"type": "error",
|
||
"message": "AI 连续 3 次未返回合法 JSON 动作，已停止本次诊断以避免无效循环",
|
||
}
|
||
return
|
||
messages.extend([
|
||
{"role": "assistant", "content": content},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
f"上一次响应无法解析：{exc}。请不要输出解释、脚本或多条命令，"
|
||
"严格只返回一个 JSON 对象："
|
||
'{"action":"command","command":"kubectl ...","reason":"..."} '
|
||
"或在全量复检问题清零后返回 final。"
|
||
),
|
||
},
|
||
])
|
||
yield {
|
||
"type": "response_rejected",
|
||
"message": f"AI 响应格式无效，正在自动纠正并重试（{invalid_response_count}/3）",
|
||
"reason": str(exc),
|
||
}
|
||
continue
|
||
except Exception as exc:
|
||
yield {"type": "error", "message": str(exc)}
|
||
return
|
||
|
||
if action["action"] == "final":
|
||
invalid_response_count = 0
|
||
yield {
|
||
"type": "final", "content": action["analysis"], "model": AI_MODEL,
|
||
"verified": True, "remaining_issues": 0,
|
||
}
|
||
return
|
||
|
||
command = action["command"]
|
||
invalid_response_count = 0
|
||
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_agent_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": "",
|
||
"report": report,
|
||
"question": question,
|
||
"execution_mode": execution_mode,
|
||
"messages": messages + [{"role": "assistant", "content": content}],
|
||
"next_step": step + 1,
|
||
"max_steps": max_steps,
|
||
"command_counts": dict(command_counts),
|
||
}
|
||
yield {
|
||
"type": "approval_required", "approval_id": approval_id,
|
||
"command": command, "reason": reason,
|
||
}
|
||
return
|
||
|
||
result = await execute_command(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 服务响应异常。"
|
||
),
|
||
}
|
||
|
||
|
||
async def resume_diagnostic_agent(pending: dict, result: dict):
|
||
"""审批命令执行后，把真实结果交回原 Agent 上下文并继续诊断。"""
|
||
messages = list(pending.get("messages") or [])
|
||
messages.append({
|
||
"role": "user",
|
||
"content": (
|
||
"用户已批准命令，以下是实际执行结果。请继续自主诊断并验证问题是否解决：\n"
|
||
+ json.dumps(result, ensure_ascii=False)
|
||
),
|
||
})
|
||
async for event in _run_agent_loop(
|
||
messages=messages,
|
||
execution_mode=pending.get("execution_mode", "ai"),
|
||
max_steps=int(pending.get("max_steps", 30)),
|
||
start_step=int(pending.get("next_step", 1)),
|
||
command_counts=Counter(pending.get("command_counts") or {}),
|
||
report=pending.get("report", ""),
|
||
question=pending.get("question", ""),
|
||
):
|
||
yield event
|
||
|
||
|
||
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)}"}
|