mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 22:51:58 +00:00
ThearchyHelios
964eeaa829
在画布上拖拽编排 agent 工作流，并实时看着它执行。 后端 FastAPI + LangGraph：把画布 JSON 编译成 StateGraph，15 种节点 （模型/Agent/多Agent、工具/沙箱、分支/循环/子图、记忆/检索/整形、 人工介入/结构校验）。选 LangGraph 是因为可视化编排真正难的是暂停、 恢复和回放，而 checkpointer + interrupt() 正好把这几件事做在了底层。 前端 React + React Flow：节点元数据驱动属性面板与连接桩，运行时 节点高亮、边流动画、卡片内流式 token 全部由同一份事件流驱动。 关键实现： - Agent 内工具审批用 @task 缓存模型调用，interrupt 重放时不重复计费 - 沙箱保持只读根文件系统，代码经 exec+base64 写入（put_archive 会被 只读容器拒绝），断网、非 root、内存/CPU/超时限额 - 检索为向量与 BM25 混合，本地哈希 embedding 零配置可用 - 分支条件走 AST 白名单求值器；HTTP 工具逐 IP 拦截内网与元数据地址 - Mock provider 让无 API Key 也能跑通完整链路 7 个内置模板全部冒烟通过，含真实模型端到端验证。
372 lines
14 KiB
Python
372 lines
14 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import contextlib
|
||
import logging
|
||
import time
|
||
from datetime import datetime, timezone
|
||
from typing import Any
|
||
|
||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||
from langgraph.types import Command
|
||
from sqlalchemy import select, update
|
||
|
||
from app.core.bus import bus
|
||
from app.core.config import settings
|
||
from app.core.events import EventType, RunEventModel
|
||
from app.db.base import SessionLocal
|
||
from app.db.models import Approval, Run, RunEvent, Workflow
|
||
from app.engine.compiler import compile_graph, initial_state
|
||
from app.engine.context import RunContext
|
||
from app.engine.schema import GraphSpec, validate_graph
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# token 事件量太大，只做实时推送不落库；回放时用节点的输出摘要重建时间线
|
||
_EPHEMERAL = {EventType.LLM_TOKEN}
|
||
|
||
|
||
class RunManager:
|
||
"""负责一次运行的完整生命周期：启动、串流、中断、恢复、取消。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._tasks: dict[str, asyncio.Task[None]] = {}
|
||
self._checkpointer: AsyncSqliteSaver | None = None
|
||
self._cm: Any = None
|
||
self._semaphore = asyncio.Semaphore(settings.max_concurrent_runs)
|
||
|
||
# ---------------- 生命周期 ----------------
|
||
|
||
async def setup(self) -> None:
|
||
settings.ensure_dirs()
|
||
self._cm = AsyncSqliteSaver.from_conn_string(str(settings.checkpoint_path))
|
||
self._checkpointer = await self._cm.__aenter__()
|
||
# 上次进程被强杀时留下的 running 记录，标成中断态而不是永远转圈
|
||
async with SessionLocal() as session:
|
||
await session.execute(
|
||
update(Run)
|
||
.where(Run.status.in_(["running", "queued"]))
|
||
.values(status="interrupted", error="服务重启，运行已挂起，可从断点恢复")
|
||
)
|
||
await session.commit()
|
||
|
||
async def shutdown(self) -> None:
|
||
for task in list(self._tasks.values()):
|
||
task.cancel()
|
||
await asyncio.gather(*self._tasks.values(), return_exceptions=True)
|
||
if self._cm is not None:
|
||
with contextlib.suppress(Exception):
|
||
await self._cm.__aexit__(None, None, None)
|
||
|
||
@property
|
||
def checkpointer(self) -> AsyncSqliteSaver:
|
||
if self._checkpointer is None:
|
||
raise RuntimeError("RunManager 还没初始化")
|
||
return self._checkpointer
|
||
|
||
# ---------------- 事件 ----------------
|
||
|
||
async def _emit(
|
||
self,
|
||
run_id: str,
|
||
event_type: EventType | str,
|
||
*,
|
||
node_id: str | None = None,
|
||
data: dict[str, Any] | None = None,
|
||
persist: bool = True,
|
||
) -> None:
|
||
event = RunEventModel(
|
||
run_id=run_id,
|
||
seq=bus.next_seq(run_id),
|
||
type=EventType(str(event_type)),
|
||
node_id=node_id,
|
||
data=data or {},
|
||
)
|
||
await bus.publish(event)
|
||
if persist and event.type not in _EPHEMERAL:
|
||
async with SessionLocal() as session:
|
||
session.add(
|
||
RunEvent(
|
||
run_id=run_id,
|
||
seq=event.seq,
|
||
type=str(event.type),
|
||
node_id=event.node_id,
|
||
ts=event.ts,
|
||
data=event.data,
|
||
)
|
||
)
|
||
await session.execute(
|
||
update(Run).where(Run.id == run_id).values(last_seq=event.seq)
|
||
)
|
||
await session.commit()
|
||
|
||
# ---------------- 启动 ----------------
|
||
|
||
async def start(
|
||
self,
|
||
*,
|
||
graph: dict[str, Any],
|
||
input_payload: dict[str, Any],
|
||
workflow_id: str | None = None,
|
||
workflow_name: str = "",
|
||
memory_scope: str = "default",
|
||
collection: str = "default",
|
||
) -> Run:
|
||
spec = GraphSpec.model_validate(graph)
|
||
report = validate_graph(spec)
|
||
if not report.ok:
|
||
raise ValueError(
|
||
"图校验未通过：" + "；".join(i.message for i in report.issues if i.level == "error")
|
||
)
|
||
|
||
async with SessionLocal() as session:
|
||
run = Run(
|
||
workflow_id=workflow_id,
|
||
workflow_name=workflow_name,
|
||
status="queued",
|
||
graph=graph,
|
||
input=input_payload,
|
||
)
|
||
run.thread_id = run.id # 一个 run 一条 checkpoint 线程
|
||
session.add(run)
|
||
await session.commit()
|
||
await session.refresh(run)
|
||
|
||
self._tasks[run.id] = asyncio.create_task(
|
||
self._drive(
|
||
run.id,
|
||
spec,
|
||
initial_state(input_payload),
|
||
workflow_id=workflow_id,
|
||
memory_scope=memory_scope,
|
||
collection=collection,
|
||
)
|
||
)
|
||
return run
|
||
|
||
async def resume(self, run_id: str, response: Any) -> Run:
|
||
"""从人工介入的断点继续。"""
|
||
async with SessionLocal() as session:
|
||
run = await session.get(Run, run_id)
|
||
if not run:
|
||
raise KeyError(f"找不到运行 {run_id}")
|
||
if run.status not in ("interrupted",):
|
||
raise ValueError(f"当前状态是 {run.status}，无法恢复")
|
||
spec = GraphSpec.model_validate(run.graph)
|
||
bus.set_seq(run_id, run.last_seq)
|
||
run.status = "running"
|
||
run.error = None
|
||
|
||
# 把这次恢复对应的待审批记录结掉，前端的待办列表才不会一直挂着
|
||
pending = (
|
||
await session.execute(
|
||
select(Approval).where(
|
||
Approval.run_id == run_id, Approval.status == "pending"
|
||
)
|
||
)
|
||
).scalars()
|
||
for approval in pending:
|
||
approval.status = "resolved"
|
||
approval.response = _safe(response) if isinstance(response, dict) else {"value": _safe(response)}
|
||
approval.resolved_at = datetime.now(timezone.utc)
|
||
|
||
await session.commit()
|
||
await session.refresh(run)
|
||
|
||
await self._emit(run_id, EventType.RUN_RESUMED, data={"response": _safe(response)})
|
||
self._tasks[run_id] = asyncio.create_task(
|
||
self._drive(run_id, spec, Command(resume=response), workflow_id=run.workflow_id)
|
||
)
|
||
return run
|
||
|
||
async def cancel(self, run_id: str) -> bool:
|
||
task = self._tasks.get(run_id)
|
||
if task and not task.done():
|
||
task.cancel()
|
||
return True
|
||
return False
|
||
|
||
def is_active(self, run_id: str) -> bool:
|
||
task = self._tasks.get(run_id)
|
||
return bool(task and not task.done())
|
||
|
||
# ---------------- 执行 ----------------
|
||
|
||
async def _drive(
|
||
self,
|
||
run_id: str,
|
||
spec: GraphSpec,
|
||
payload: Any,
|
||
*,
|
||
workflow_id: str | None = None,
|
||
memory_scope: str = "default",
|
||
collection: str = "default",
|
||
) -> None:
|
||
started = time.perf_counter()
|
||
status = "succeeded"
|
||
error: str | None = None
|
||
final_state: dict[str, Any] = {}
|
||
|
||
async with self._semaphore:
|
||
try:
|
||
async with SessionLocal() as session:
|
||
await session.execute(
|
||
update(Run)
|
||
.where(Run.id == run_id)
|
||
.values(status="running", started_at=datetime.now(timezone.utc))
|
||
)
|
||
await session.commit()
|
||
await self._emit(
|
||
run_id,
|
||
EventType.RUN_STARTED,
|
||
data={"nodes": len(spec.nodes), "resumed": isinstance(payload, Command)},
|
||
)
|
||
|
||
run_ctx = RunContext(
|
||
run_id=run_id,
|
||
thread_id=run_id,
|
||
spec=spec,
|
||
workflow_id=workflow_id,
|
||
memory_scope=memory_scope,
|
||
collection=collection,
|
||
)
|
||
app = compile_graph(spec, run_ctx).compile(checkpointer=self.checkpointer)
|
||
config = {
|
||
"configurable": {"thread_id": run_id},
|
||
"recursion_limit": settings.max_graph_steps,
|
||
}
|
||
|
||
interrupted = False
|
||
async for mode, chunk in app.astream(
|
||
payload, config=config, stream_mode=["custom", "updates"]
|
||
):
|
||
if mode == "custom":
|
||
await self._handle_custom(run_id, chunk)
|
||
elif mode == "updates":
|
||
if await self._handle_updates(run_id, chunk):
|
||
interrupted = True
|
||
|
||
snapshot = await app.aget_state(config)
|
||
final_state = dict(snapshot.values or {})
|
||
if interrupted or snapshot.interrupts:
|
||
status = "interrupted"
|
||
elif snapshot.next:
|
||
status = "interrupted"
|
||
|
||
except asyncio.CancelledError:
|
||
status = "cancelled"
|
||
error = "用户取消"
|
||
await self._emit(run_id, EventType.RUN_CANCELLED, data={})
|
||
raise
|
||
except Exception as e: # noqa: BLE001
|
||
status = "failed"
|
||
error = f"{type(e).__name__}: {e}"
|
||
logger.exception("run %s 失败", run_id)
|
||
await self._emit(run_id, EventType.RUN_FAILED, data={"error": error})
|
||
finally:
|
||
elapsed = int((time.perf_counter() - started) * 1000)
|
||
await self._finalize(run_id, status, error, final_state, elapsed)
|
||
self._tasks.pop(run_id, None)
|
||
|
||
async def _handle_custom(self, run_id: str, chunk: Any) -> None:
|
||
"""节点通过 get_stream_writer 发来的事件。"""
|
||
if not isinstance(chunk, dict) or not chunk.get("__agentlab__"):
|
||
return
|
||
await self._emit(
|
||
run_id,
|
||
chunk.get("type", EventType.LOG),
|
||
node_id=chunk.get("node_id"),
|
||
data=chunk.get("data") or {},
|
||
)
|
||
|
||
async def _handle_updates(self, run_id: str, chunk: Any) -> bool:
|
||
"""LangGraph 的状态更新。这里只关心中断信号 —— 节点级事件已经由 custom 流覆盖。"""
|
||
if not isinstance(chunk, dict):
|
||
return False
|
||
interrupts = chunk.get("__interrupt__")
|
||
if not interrupts:
|
||
return False
|
||
|
||
for item in interrupts:
|
||
value = getattr(item, "value", item)
|
||
interrupt_id = getattr(item, "id", None)
|
||
payload = value if isinstance(value, dict) else {"value": value}
|
||
async with SessionLocal() as session:
|
||
existing = (
|
||
await session.execute(
|
||
select(Approval).where(
|
||
Approval.run_id == run_id,
|
||
Approval.interrupt_id == interrupt_id,
|
||
Approval.status == "pending",
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if existing is None:
|
||
session.add(
|
||
Approval(
|
||
run_id=run_id,
|
||
node_id=str(payload.get("node_id") or ""),
|
||
interrupt_id=interrupt_id,
|
||
mode=str(payload.get("mode") or "approve"),
|
||
title=str(payload.get("title") or "需要你确认"),
|
||
payload=payload,
|
||
schema_=payload.get("schema") or {},
|
||
)
|
||
)
|
||
await session.commit()
|
||
await self._emit(
|
||
run_id,
|
||
EventType.RUN_INTERRUPTED,
|
||
node_id=payload.get("node_id"),
|
||
data={"interrupt_id": interrupt_id, "payload": _safe(payload)},
|
||
)
|
||
return True
|
||
|
||
async def _finalize(
|
||
self,
|
||
run_id: str,
|
||
status: str,
|
||
error: str | None,
|
||
state: dict[str, Any],
|
||
elapsed_ms: int,
|
||
) -> None:
|
||
usage = dict(state.get("usage") or {})
|
||
usage["duration_ms"] = elapsed_ms
|
||
output = dict(state.get("output") or {})
|
||
|
||
async with SessionLocal() as session:
|
||
run = await session.get(Run, run_id)
|
||
if run:
|
||
run.status = status
|
||
run.error = error
|
||
run.output = output
|
||
run.usage = usage
|
||
if status != "interrupted":
|
||
run.finished_at = datetime.now(timezone.utc)
|
||
await session.commit()
|
||
|
||
if status == "succeeded":
|
||
await self._emit(
|
||
run_id,
|
||
EventType.RUN_FINISHED,
|
||
data={"output": _safe(output), "usage": usage, "duration_ms": elapsed_ms},
|
||
)
|
||
if status in ("succeeded", "failed", "cancelled"):
|
||
await bus.close(run_id)
|
||
|
||
|
||
def _safe(value: Any, limit: int = 4000) -> Any:
|
||
"""把任意对象裁剪成能安全放进事件负载的形状。"""
|
||
if isinstance(value, str):
|
||
return value[:limit]
|
||
if isinstance(value, dict):
|
||
return {k: _safe(v, limit // 2) for k, v in list(value.items())[:50]}
|
||
if isinstance(value, list):
|
||
return [_safe(v, limit // 2) for v in value[:50]]
|
||
if isinstance(value, (int, float, bool)) or value is None:
|
||
return value
|
||
return str(value)[:limit]
|
||
|
||
|
||
run_manager = RunManager()
|