Your Name
00d95c4c5e
现在 AI Agent 支持在单条命令中使用: - 管道 (|): crictl ps -a | grep nginx | head -5 - 重定向 (> / >>): journalctl -u kubelet > /tmp/log - 分号 (;): systemctl status kubelet; df -h - 逻辑连接符 (&& / ||): df -h && free -m 安全策略: - FORBIDDEN_TOKENS 清空，所有 shell 操作符均放行 - 新增 _split_into_commands() 按操作符分割命令段 - 新增 _command_has_dangerous_cmds() 扫描所有段中的武器级命令 - 永远拒绝: sh/bash/dash/zsh/dd/format/.../parted 等 - 检查第一条命令的白名单，cat 不安全路径检查 - 只读白名单新增 echo printf cd pwd export test [
516 lines
21 KiB
Python
516 lines
21 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, top, du -sh, curl, tcpdump
|
||
3. containerd 命令 - 查询容器运行时状态。例如: ctr containers list, ctr images list, ctr tasks ls, nerdctl ps -a, nerdctl logs <container>, nerdctl stats
|
||
4. etcd 诊断命令 - 检查 etcd 集群健康与数据。例如: etcdctl endpoint status, etcdctl endpoint health, etcdctl member list, etcdctl alarm list, etcdctl get /kubernetes --prefix
|
||
5. Linux 基础调优命令 - 例如: iostat -xz 1 3, vmstat 1 3, mpstat -P ALL 1 3, pidstat -d 1 3, lsof -i :<port>, ethtool <iface>, tcpdump -i any -c 10 -nn
|
||
|
||
行为要求：
|
||
1. 如果证据不足，必须返回 command 动作并逐条调查；至少检查与故障资源直接相关的状态、describe、events 或 logs。
|
||
2. 禁止在 final 中仅提供"请执行以下命令"的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
|
||
3. 每次只提出一条命令，服务端会执行并把真实输出返回给你，然后你继续下一轮判断。
|
||
|
||
安全规则：
|
||
1. 可以使用管道 (|)、重定向 (>/>>)、逻辑连接符 (&&/||) 或分号 (;) 组合多条命令进行精确排查。不得使用反引号 (`` ` ``) 执行命令替换。
|
||
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. 只要问题没有完全解决（verified=true 且 remaining_issues=0），就要持续执行命令诊断，没有轮数上限。
|
||
|
||
每轮必须只返回一个 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",
|
||
run_id: str = "",
|
||
history_store=None,
|
||
):
|
||
"""自主诊断循环：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,
|
||
run_id=run_id,
|
||
history_store=history_store,
|
||
):
|
||
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,
|
||
run_id: str = "",
|
||
history_store=None,
|
||
):
|
||
invalid_response_count = 0
|
||
step = start_step
|
||
while True:
|
||
if max_steps > 0 and step > max_steps:
|
||
break
|
||
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": "请在终端手动执行命令，并根据结果继续排查",
|
||
}
|
||
# 手动模式也保存上下文，支持断点续传
|
||
self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
|
||
max_steps, report, question, execution_mode, content)
|
||
return
|
||
|
||
if decision.requires_approval:
|
||
approval_id = uuid.uuid4().hex
|
||
_PENDING_APPROVALS[approval_id] = {
|
||
"command": command,
|
||
"reason": reason,
|
||
"created_by": AI_MODEL,
|
||
"run_id": run_id,
|
||
"report": report,
|
||
"question": question,
|
||
"execution_mode": execution_mode,
|
||
"history_store": history_store,
|
||
"messages": messages + [{"role": "assistant", "content": content}],
|
||
"next_step": step + 1,
|
||
"max_steps": max_steps,
|
||
"command_counts": dict(command_counts),
|
||
}
|
||
# 保存上下文到 DB，避免浏览器关闭/刷新丢失进度
|
||
self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
|
||
max_steps, report, question, execution_mode, content)
|
||
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)},
|
||
])
|
||
|
||
# 每次执行后保存上下文，支持断点续传
|
||
self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
|
||
max_steps, report, question, execution_mode)
|
||
|
||
step += 1
|
||
|
||
if max_steps > 0:
|
||
yield {
|
||
"type": "error",
|
||
"message": (
|
||
f"自主诊断已达到安全上限（{max_steps} 轮），但 Agent 仍未形成有效结论。"
|
||
"系统已停止继续执行以避免无限循环，请检查是否重复执行相同命令或 AI 服务响应异常。"
|
||
),
|
||
}
|
||
# max_steps=0 表示不限轮数，不输出错误 — Agent 会持续执行直到完成或遇到致命错误
|
||
|
||
|
||
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)
|
||
),
|
||
})
|
||
run_id = pending.get("run_id", "")
|
||
history_store = pending.get("history_store")
|
||
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", ""),
|
||
run_id=run_id,
|
||
history_store=history_store,
|
||
):
|
||
yield event
|
||
|
||
|
||
def take_pending_approval(approval_id: str) -> dict | None:
|
||
return _PENDING_APPROVALS.pop(approval_id, None)
|
||
|
||
|
||
def self_save_context(
|
||
run_id: str, history_store,
|
||
messages: list, command_counts: Counter,
|
||
step: int, next_step: int,
|
||
max_steps: int, report: str, question: str,
|
||
execution_mode: str,
|
||
last_response_content: str = "",
|
||
):
|
||
"""保存 Agent 执行上下文到 DB，支持断点续传。"""
|
||
if not run_id or not history_store:
|
||
return
|
||
context = {
|
||
"messages": messages,
|
||
"command_counts": dict(command_counts),
|
||
"step": step,
|
||
"next_step": next_step,
|
||
"max_steps": max_steps,
|
||
"report": report,
|
||
"question": question,
|
||
"execution_mode": execution_mode,
|
||
"last_response": last_response_content,
|
||
}
|
||
history_store.update_context(run_id, context)
|
||
|
||
|
||
async def resume_from_context(
|
||
ctx: dict,
|
||
run_id: str = "",
|
||
history_store=None,
|
||
):
|
||
"""从 DB 保存的上下文恢复 Agent 诊断，不重跑历史命令，直接进入下一轮。"""
|
||
messages = list(ctx.get("messages") or [])
|
||
command_counts = Counter(ctx.get("command_counts") or {})
|
||
next_step = int(ctx.get("next_step", 1))
|
||
max_steps = int(ctx.get("max_steps", 0))
|
||
report = ctx.get("report", "")
|
||
question = ctx.get("question", "")
|
||
execution_mode = ctx.get("execution_mode", "ai")
|
||
|
||
async for event in _run_agent_loop(
|
||
messages=messages,
|
||
execution_mode=execution_mode,
|
||
max_steps=max_steps,
|
||
start_step=next_step,
|
||
command_counts=command_counts,
|
||
report=report,
|
||
question=question,
|
||
run_id=run_id,
|
||
history_store=history_store,
|
||
):
|
||
yield event
|
||
|
||
|
||
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)}"}
|