200 lines
7.1 KiB
Python
200 lines
7.1 KiB
Python
"""AI 控制台 API：模型/Key 前端配置 + Agent 流式执行（SSE 事件流）。
|
||
|
||
流程：
|
||
POST /api/agent/run {prompt} → 启动 Agent 线程，返回 run_id
|
||
GET /api/agent/stream?run_id= → SSE 事件流（EventSource 订阅）：
|
||
event: delta {text, kind: content|reasoning} 流式文本增量
|
||
event: step {tool, args, image?} 工具调用完成（MCP 步骤）
|
||
event: done {answer} 完成
|
||
event: error {message} 失败
|
||
GET/POST /api/agent/config → 配置读写（key 打码回显）
|
||
|
||
单实例：同时只允许一个 Agent 运行。
|
||
"""
|
||
import asyncio
|
||
import base64
|
||
import io
|
||
import json
|
||
import os
|
||
import queue
|
||
import sys
|
||
import threading
|
||
import uuid
|
||
|
||
from flask import Blueprint, Response, jsonify, request
|
||
|
||
from core.logger import get_logger
|
||
from core.models import db
|
||
from web.auth import admin_required
|
||
|
||
_log = get_logger("web.agent")
|
||
bp = Blueprint("agent", __name__)
|
||
|
||
# ---------- 配置键（app_meta） ----------
|
||
_CFG_KEYS = {"api_base": "agent_api_base",
|
||
"model": "agent_model",
|
||
"api_key": "agent_api_key",
|
||
"default_serial": "agent_default_serial"}
|
||
|
||
# ---------- 运行状态（单实例 + 事件队列） ----------
|
||
_run = {"id": None, "state": "idle", "prompt": "", "serial": "",
|
||
"answer": "", "error": ""}
|
||
_queues = {} # run_id -> queue.Queue（SSE 消费者读取）
|
||
_lock = threading.Lock()
|
||
|
||
|
||
def _meta_get(key):
|
||
return db.session.execute(
|
||
db.text("SELECT value FROM app_meta WHERE key=:k"), {"k": key}).scalar() or ""
|
||
|
||
|
||
def _meta_put(key, value):
|
||
db.session.execute(
|
||
db.text("INSERT OR REPLACE INTO app_meta(key,value) VALUES(:k,:v)"),
|
||
{"k": key, "v": str(value)})
|
||
|
||
|
||
def _read_cfg():
|
||
return {k: _meta_get(v) for k, v in _CFG_KEYS.items()}
|
||
|
||
|
||
# ================== 配置 ==================
|
||
@bp.route("/api/agent/config", methods=["GET"])
|
||
@admin_required
|
||
def agent_config_get():
|
||
"""读 Agent 配置（key 打码返回）。"""
|
||
cfg = _read_cfg()
|
||
if cfg["api_key"]:
|
||
k = cfg["api_key"]
|
||
cfg["api_key_masked"] = k[:6] + "***" + k[-4:]
|
||
return jsonify({"ok": True, **cfg})
|
||
|
||
|
||
@bp.route("/api/agent/config", methods=["POST"])
|
||
@admin_required
|
||
def agent_config_save():
|
||
"""保存 Agent 配置：{api_base?, model?, api_key?, default_serial?} 部分更新。"""
|
||
data = request.json or {}
|
||
for key, meta_key in _CFG_KEYS.items():
|
||
if key in data and data[key] is not None:
|
||
_meta_put(meta_key, str(data[key]).strip())
|
||
db.session.commit()
|
||
return jsonify({"ok": True, "msg": "已保存"})
|
||
|
||
|
||
# ================== 运行 ==================
|
||
@bp.route("/api/agent/run", methods=["POST"])
|
||
@admin_required
|
||
def agent_run():
|
||
"""启动 Agent：{prompt, serial?}。运行中返回 409。"""
|
||
data = request.json or {}
|
||
prompt = (data.get("prompt") or "").strip()
|
||
if not prompt:
|
||
return jsonify({"ok": False, "error": "请输入指令"}), 400
|
||
cfg = _read_cfg()
|
||
if not cfg.get("api_key"):
|
||
return jsonify({"ok": False, "error": "请先在配置区填写 API Key"}), 400
|
||
if not cfg.get("model"):
|
||
return jsonify({"ok": False, "error": "请先填写模型名"}), 400
|
||
with _lock:
|
||
if _run["state"] == "running":
|
||
return jsonify({"ok": False, "error": "已有 Agent 运行中，请等待完成"}), 409
|
||
run_id = uuid.uuid4().hex[:8]
|
||
_run.update(id=run_id, state="running", prompt=prompt,
|
||
serial=(data.get("serial") or "").strip(),
|
||
answer="", error="")
|
||
_queues[run_id] = queue.Queue()
|
||
_log.info(f"Agent 启动: {prompt[:60]}")
|
||
threading.Thread(target=_agent_thread, args=(run_id, prompt, cfg),
|
||
daemon=True).start()
|
||
return jsonify({"ok": True, "run_id": run_id})
|
||
|
||
|
||
@bp.route("/api/agent/stream")
|
||
@admin_required
|
||
def agent_stream():
|
||
"""SSE 事件流（EventSource）：delta/step/done/error。"""
|
||
run_id = request.args.get("run_id", "")
|
||
with _lock:
|
||
if run_id != _run["id"]:
|
||
return jsonify({"ok": False, "error": "run_id 不存在"}), 404
|
||
q = _queues.get(run_id)
|
||
|
||
def gen():
|
||
while True:
|
||
try:
|
||
evt = q.get(timeout=15)
|
||
except queue.Empty:
|
||
yield ": keepalive\n\n" # 心跳防超时
|
||
continue
|
||
if evt is None:
|
||
break
|
||
kind, payload = evt
|
||
yield f"event: {kind}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
||
if kind in ("done", "error"):
|
||
break
|
||
|
||
return Response(gen(), mimetype="text/event-stream",
|
||
headers={"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no"})
|
||
|
||
|
||
def _shrink_image(b64, width=220, quality=50):
|
||
"""截图降采样（SSE step 事件用，控制传输体积）。失败原样返回。"""
|
||
try:
|
||
from PIL import Image
|
||
img = Image.open(io.BytesIO(base64.b64decode(b64)))
|
||
if img.width > width:
|
||
img = img.resize((width, int(img.height * width / img.width)))
|
||
buf = io.BytesIO()
|
||
img.convert("RGB").save(buf, "JPEG", quality=quality)
|
||
return base64.b64encode(buf.getvalue()).decode()
|
||
except Exception:
|
||
return b64
|
||
|
||
|
||
def _agent_thread(run_id, prompt, cfg):
|
||
"""后台线程：Agent 流式执行，事件推入队列供 SSE 消费。"""
|
||
q = _queues.get(run_id)
|
||
try:
|
||
# 平台进程 cwd=/app（含 mcp_agent 包），进程内 import
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
from mcp_agent.agent import Agent
|
||
|
||
def on_delta(text, kind):
|
||
q.put(("delta", {"text": text, "kind": kind}))
|
||
|
||
def on_tool(step):
|
||
rec = {"tool": step.get("tool"),
|
||
"args": str(step.get("args"))[:200]}
|
||
if step.get("image"):
|
||
rec["image"] = _shrink_image(step["image"])
|
||
q.put(("step", rec))
|
||
|
||
agent = Agent()
|
||
agent.s.api_base = cfg.get("api_base") or agent.s.api_base
|
||
agent.s.model = cfg.get("model") or agent.s.model
|
||
agent.s.api_key = cfg.get("api_key") or agent.s.api_key
|
||
agent.s.default_serial = cfg.get("default_serial") or agent.s.default_serial
|
||
|
||
async def _execute():
|
||
await agent._load_tools()
|
||
return await agent.run_stream(prompt, cfg.get("default_serial") or "",
|
||
on_delta=on_delta, on_tool=on_tool)
|
||
|
||
# 整体超时保护：卡死时结束，释放单实例
|
||
answer = asyncio.run(asyncio.wait_for(_execute(), timeout=600))
|
||
with _lock:
|
||
_run["state"] = "done"
|
||
_run["answer"] = answer
|
||
q.put(("done", {"answer": answer}))
|
||
except Exception as e:
|
||
_log.warning(f"Agent 运行异常: {e}")
|
||
with _lock:
|
||
_run["state"] = "error"
|
||
_run["error"] = f"{type(e).__name__}: {str(e)[:200]}"
|
||
q.put(("error", {"message": str(e)[:200]}))
|
||
finally:
|
||
q.put(None) # 关闭 SSE
|
||
_queues.pop(run_id, None)
|