""" 本地 AI Agent 模块 通过 OpenAI 兼容的 HTTP API 跟大模型对话，并把工具调用能力 （执行 SSH 命令、查询系统状态、文件操作）暴露给模型。 无需本地 GPU，配置 API Key + endpoint 即可使用。 """ import json import re import time from typing import Optional, List, Dict, Any, Callable import requests # AI Agent 可调用的工具定义（OpenAI function calling 格式） TOOL_DEFINITIONS = [ { "type": "function", "function": { "name": "exec_ssh_command", "description": "在当前 SSH 连接的远程主机上执行一条 shell 命令，返回标准输出、错误和退出码。用于排查问题、查看文件、启停服务、修改配置等所有需要 shell 的场景。", "parameters": { "type": "object", "properties": { "command": {"type": "string", "description": "要执行的完整 shell 命令字符串"}, "timeout": {"type": "integer", "description": "超时秒数，默认 30，最大 300", "default": 30}, }, "required": ["command"], }, }, }, { "type": "function", "function": { "name": "get_system_metrics", "description": "获取远程主机的实时系统指标：CPU 使用率/核心数、内存、磁盘使用率、各网卡收发字节、负载。", "parameters": {"type": "object", "properties": {}}, }, }, { "type": "function", "function": { "name": "list_remote_files", "description": "列出远程主机上指定目录的文件和子目录。", "parameters": { "type": "object", "properties": { "path": {"type": "string", "description": "要列出的远程绝对路径", "default": "/"}, }, "required": ["path"], }, }, }, { "type": "function", "function": { "name": "read_remote_file", "description": "读取远程主机上一个文本文件的内容（最大 8000 字节；超过会截断）。", "parameters": { "type": "object", "properties": { "path": {"type": "string", "description": "远程文件的绝对路径"}, }, "required": ["path"], }, }, }, { "type": "function", "function": { "name": "upload_local_file", "description": "把本地文件上传到远程主机的指定路径。", "parameters": { "type": "object", "properties": { "local_path": {"type": "string", "description": "本地文件绝对路径"}, "remote_path": {"type": "string", "description": "远程目标绝对路径"}, }, "required": ["local_path", "remote_path"], }, }, }, ] class AIAgent: """OpenAI 兼容协议的对话 + 工具调用 Agent""" def __init__(self, api_key: str = "", base_url: str = "https://api.openai.com/v1", model: str = "gpt-4o-mini", system_prompt: str = ""): self.api_key = api_key self.base_url = base_url.rstrip("/") self.model = model self.system_prompt = system_prompt or self._default_system_prompt() self.history: List[Dict[str, Any]] = [] self.max_steps = 8 # 单轮最大工具调用次数，避免死循环 @staticmethod def _default_system_prompt() -> str: return ( "你是一个专业的服务器运维 AI Agent，可以通过工具调用操作当前 SSH 会话里的远程主机。" "优先使用工具获取真实数据再回答，不要凭空猜测。" "对破坏性操作（rm -rf、kill -9、systemctl stop、重启等）务必先确认。" "回复简洁，用中文，给出可执行的结论和命令。" ) def update_config(self, api_key: str = "", base_url: str = "", model: str = "", system_prompt: str = ""): """热更新配置""" if api_key: self.api_key = api_key if base_url: self.base_url = base_url.rstrip("/") if model: self.model = model if system_prompt: self.system_prompt = system_prompt def clear_history(self): self.history = [] def chat(self, user_message: str, tool_executor: Callable[[str, dict], str], on_step: Optional[Callable[[str, str], None]] = None) -> str: """ 发起一轮对话。tool_executor(tool_name, arguments) -> str（工具执行的纯文本结果）。 on_step(role, content) 在每个思考/工具步骤触发，用于把过程实时渲染到 UI。 返回最终助手回复文本。 """ if not self.api_key: raise RuntimeError("未配置 API Key，请先在「AI 设置」中填写") self.history.append({"role": "user", "content": user_message}) for step in range(self.max_steps): try: resp = self._call_llm() except Exception as e: msg = f"[AI 调用失败] {e}" if on_step: on_step("error", msg) return msg msg = resp.choices[0].message tool_calls = getattr(msg, "tool_calls", None) or [] content = (msg.content or "").strip() if not tool_calls: # 没有工具调用 -> 终态 self.history.append({"role": "assistant", "content": content}) if on_step and content: on_step("assistant", content) return content # 工具调用阶段 self.history.append({ "role": "assistant", "content": content, "tool_calls": [ { "id": tc.id, "type": "function", "function": { "name": tc.function.name, "arguments": tc.function.arguments, }, } for tc in tool_calls ], }) if on_step and content: on_step("assistant_thinking", content) for tc in tool_calls: fn_name = tc.function.name try: args = json.loads(tc.function.arguments) if tc.function.arguments else {} except json.JSONDecodeError: args = {} if on_step: on_step("tool_call", f"调用工具: {fn_name}({json.dumps(args, ensure_ascii=False)})") result = tool_executor(fn_name, args) # 截断过长的结果，避免撑爆上下文 result_str = result if len(result) < 6000 else result[:6000] + "\n...(已截断)" self.history.append({ "role": "tool", "tool_call_id": tc.id, "name": fn_name, "content": result_str, }) if on_step: on_step("tool_result", f"[{fn_name}] -> {result_str[:400]}") return "已达最大工具调用步数，可能未完成任务。可继续提问或追加要求。" def _call_llm(self): """发起一次 LLM 调用""" if not self.api_key: raise RuntimeError("未配置 API Key") headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } payload = { "model": self.model, "messages": [{"role": "system", "content": self.system_prompt}] + self.history, "tools": TOOL_DEFINITIONS, "tool_choice": "auto", "temperature": 0.2, } r = requests.post( f"{self.base_url}/chat/completions", headers=headers, json=payload, timeout=60, ) if r.status_code != 200: raise RuntimeError(f"HTTP {r.status_code}: {r.text[:300]}") # 用一个轻量对象包装，调用方用 .choices[0].message.tool_calls / .content return _Resp(r.json()) class _Msg: def __init__(self, d): self.content = d.get("content") or "" self.tool_calls = None tcs = d.get("tool_calls") if tcs: self.tool_calls = [] for tc in tcs: fn = tc.get("function", {}) self.tool_calls.append(_TC(tc.get("id", ""), fn.get("name", ""), fn.get("arguments", ""))) class _TC: def __init__(self, _id, name, arguments): self.id = _id self.function = _FN(name, arguments) class _FN: def __init__(self, name, arguments): self.name = name self.arguments = arguments class _Choice: def __init__(self, d): self.message = _Msg(d.get("message", {})) class _Resp: def __init__(self, j): self.choices = [_Choice(c) for c in j.get("choices", [])]