#!/usr/bin/env python3 # -*- coding: utf-8 -*- """作业 7 · 步骤 4：挂 3 个自愈工具 + 提示词。""" from __future__ import annotations import json import re import sys from pathlib import Path import _bootstrap # noqa: F401 import _agent_cfg as AC import _config as C from dify_api import Dify, pp PREFERRED = ( "detect_vllm_anomaly", "restart_vllm_container", "send_recovery_notification", ) SYSTEM_PROMPT = """你是 vLLM 故障自愈运维 Agent。当用户报告 vLLM OOM、崩溃、不可用或要求自愈时，必须严格按顺序执行： 1. detect_vllm_anomaly — 检测是否 OOM 或容器异常 2. 若 anomaly=true 且 recommendation=restart_vllm_container，调用 restart_vllm_container 重启容器 3. 重启后再次 detect_vllm_anomaly 确认恢复 4. send_recovery_notification — 向钉钉发送恢复通知（message 写结论，detail 写证据与重启结果） 禁止跳过检测直接重启。禁止编造未发生的重启。 通知示例 message：「vLLM OOM 已自愈，容器已重启恢复服务」。""" def find_ids(d: Dify) -> tuple[str, str, str, str]: agent_id = "" for item in (d.get("/agent?page=1&limit=50").get("json") or {}).get("data") or []: if item.get("name") == C.HOMEWORK7_APP_NAME: agent_id = item.get("agent_id") or item.get("id") or "" break if not agent_id: raise RuntimeError("找不到 Agent: %s" % C.HOMEWORK7_APP_NAME) sql = ( "SELECT a.app_id, a.id, a.active_config_snapshot_id " "FROM agents a WHERE a.id='%s' LIMIT 1;" % agent_id ) out, err = d.sh( 'docker exec %s psql -U postgres -d dify -t -A -F "|" -c "%s"' % (C.PG_CONTAINER, sql) ) line = (out or "").strip().splitlines() if not line: raise RuntimeError(err) app_id, agent_id, snap_id = line[0].split("|") draft_id = ensure_draft(d, agent_id, snap_id) return app_id, agent_id, snap_id, draft_id def ensure_draft(d: Dify, agent_id: str, snap_id: str) -> str: out, _ = d.sh( 'docker exec %s psql -U postgres -d dify -t -A -c ' '"SELECT id FROM agent_config_drafts WHERE agent_id=\'%s\' LIMIT 1;"' % (C.PG_CONTAINER, agent_id) ) draft_id = (out or "").strip().splitlines() if draft_id and draft_id[0].strip(): return draft_id[0].strip() ins = ( "INSERT INTO agent_config_drafts (" "id, tenant_id, agent_id, draft_type, draft_owner_key, account_id, " "base_snapshot_id, config_snapshot, created_by, updated_by, created_at, updated_at" ") SELECT gen_random_uuid(), s.tenant_id, s.agent_id, 'draft', '', NULL, " "s.id, s.config_snapshot, s.created_by, s.created_by, NOW(), NOW() " "FROM agent_config_snapshots s WHERE s.id='%s' RETURNING id;" % snap_id ) remote = "/tmp/hw7_ins_draft.sql" with d.sftp.open(remote, "w") as f: f.write(ins) out, err = d.sh( "docker exec -i %s psql -U postgres -d dify -t -A < %s" % (C.PG_CONTAINER, remote) ) draft_id = (out or "").strip().splitlines() draft_id = draft_id[0].strip() if draft_id else "" if not draft_id: raise RuntimeError("无法创建草稿: %s %s" % (out, err)) return draft_id def resolve_tool_entries(d: Dify) -> list[dict]: out, err = d.sh( "docker exec %s psql -U postgres -d dify -t -A -c " "\"SELECT id FROM tool_api_providers WHERE name='%s' LIMIT 1;\"" % (C.PG_CONTAINER, C.SELF_HEAL_PROVIDER) ) provider_uuid = (out or "").strip() r = d.get( "/workspaces/current/tool-provider/api/tools?provider=%s" % C.SELF_HEAL_PROVIDER ) data = r.get("json") or [] by_name = {} for tool in data: name = tool.get("name") or tool.get("tool_name") or "" by_name[name] = { "enabled": True, "provider_type": "api", "provider_id": provider_uuid, "tool_name": name, "credential_type": "unauthorized", "credential_ref": None, "runtime_parameters": {}, } return [by_name[n] for n in PREFERRED if n in by_name] def patch_cfg(cfg: dict, tool_entries: list[dict]) -> dict: prompt = cfg.setdefault("prompt", {}) if prompt is None: cfg["prompt"] = {} prompt = cfg["prompt"] prompt["system_prompt"] = SYSTEM_PROMPT tools = cfg.get("tools") if not isinstance(tools, dict): tools = {} cfg["tools"] = tools tools["dify_tools"] = tool_entries tools["cli_tools"] = tools.get("cli_tools") or [] model = cfg.get("model") if not isinstance(model, dict): model = {} cfg["model"] = model model["model"] = C.QWEN_MODEL model["plugin_id"] = "langgenius/openai" model["model_provider"] = C.OPENAI_PROVIDER model["credential_ref"] = None ms = model.get("model_settings") if not isinstance(ms, dict): ms = {} model["model_settings"] = ms ms["max_tokens"] = ms.get("max_tokens") or 4096 return cfg def write_ids(app_id: str, agent_id: str, snap_id: str, draft_id: str) -> None: mapping = { "HOMEWORK7_APP_ID": app_id, "HOMEWORK7_AGENT_ID": agent_id, "HOMEWORK7_SNAP_ID": snap_id, "HOMEWORK7_DRAFT_ID": draft_id, } for path in ( Path(__file__).resolve().parent / "_config.py", Path(__file__).resolve().parents[1] / "_config.py", ): if not path.is_file(): continue text = path.read_text(encoding="utf-8") for key, val in mapping.items(): text, n = re.subn(r'%s = "[^"]*"' % key, '%s = "%s"' % (key, val), text, count=1) if n != 1: raise RuntimeError("update failed %s in %s" % (key, path)) path.write_text(text, encoding="utf-8") def main() -> int: d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD) try: d.login() app_id, agent_id, snap_id, draft_id = find_ids(d) tool_entries = resolve_tool_entries(d) print("tool_entries", json.dumps(tool_entries, ensure_ascii=False)) snap = AC.load_cfg(d, "agent_config_snapshots", snap_id) draft = AC.load_cfg(d, "agent_config_drafts", draft_id) AC.save_cfg(d, "agent_config_snapshots", snap_id, patch_cfg(snap, tool_entries), "hw7s") AC.save_cfg(d, "agent_config_drafts", draft_id, patch_cfg(draft, tool_entries), "hw7d") write_ids(app_id, agent_id, snap_id, draft_id) print("[OK] 已挂工具，勿点发布") return 0 finally: d.close() if __name__ == "__main__": sys.exit(main())