mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 22:31:59 +00:00
ThearchyHelios
62aac187d8
MCP / 自定义工具的审批三档（tools/trust.py） - 平台看不出这两类工具做什么，以前审批关卡一律当它们安全、照跑不误。现在默认每次调用都等审批； 工具页可以逐个改成「始终允许 · 门控把关」或「始终允许」，审批卡上的「始终允许」等于前者 - 门控把关：每次调用先由门控模型审一遍工具说明、参数和节点任务；判可疑、出错、超时、答不成格式， 一律交还给人，不会变成放行。判定包进 checkpoint，审批恢复重放时不再问一遍 - agent 节点、工具节点、协作团队三处关卡都接上；团队成员停不下来等人，门控没放行的调用不执行 - 每次运行在发起时快照档位（runs.tool_trust），恢复和接着跑沿用同一份；正式运行不看档位，一律等审批； 升级前发起的运行照旧不问——它们可能正停在某个审批上，多问一次答复就会对错号 - 接口：GET /api/tools 带 trust / trust_key，PUT /api/tools/trust；设置新增门控模型；新事件 tool.gated - 界面：工具页三选一、审批卡「始终允许」、设置页门控模型、运行记录解码 tool.gated； 顺手修了 MCP 工具详情 404、工具审批被驳回却显示成放行 智能体护栏（engine/guards.py），取代固定的「最多 12 步」 - 步数只留兜底：默认 100，设置里可改；以前画布自动写进节点的 12 按「没填」处理 - 同一个工具、同样的参数第二次不再执行，把上次的结果交还给它；连续 3 步拿不到新信息就收尾 - 每个节点的令牌预算（默认 200 万）和金额预算（默认不限），都可以在设置里改成不限 - 上下文用到窗口的 60% 先压缩早期的工具结果，到 85% 就收尾；临近任何一个上限先提醒模型一次 - 下游结构化校验失败时写明是上游哪个 agent 提前收了尾；复核信号带上收尾原因， 问数据页只在真是步数用满时才给「放宽步数重跑」 - 上限在发起时快照（runs.agent_limits）；升级前发起的运行一字不改走旧逻辑 检查脚本 - check-all 默认 4 道并行：量时间最多的两项互不重叠；并行时没过的最后单独重跑一遍，重跑才过的 单列成「偶发」。16 项全量从约 21 分钟降到约 7 分钟（加上一次重跑约 10 分钟）；--lanes 1 回到串行 顺带修的 - 设置接口里的 limits 改为只读：库里存着旧版设置页写进去的一份（max_agent_steps: 25），从不生效， 却盖住了真实上限，问数据页照着错的数许诺 - test_kb 里嵌入模型那条测试不再真连 api.openai.com（代理慢时一条要等近 4 分钟） - check-stream 读滚动容器前先等它挂上，并行跑时不再偶发 验证（沙箱）：tsc 干净，生产构建通过，pytest 1655 passed（两遍），check-all 16 项检查、3562 项全过； 新增后端测试 53 条，关键逻辑做了 22 处变异验证、全部抓到。
428 lines
18 KiB
Python
428 lines
18 KiB
Python
"""agent 护栏：取代以前固定的「最多 12 步」。
|
||
|
||
用户的原话是「12 太少了，撞上就判定运行失败」。真实运行里（719253eb、6310f291）agent
|
||
用满 12 步收了尾，结论缺数据，下游的结构化校验缺字段，整个运行失败，报错里看不出根源。
|
||
|
||
现在：
|
||
- 步数只剩一个很高的兜底（默认 100，设置里可改）；以前画布写死的 12 按「没填」处理；
|
||
- 同一个调用第二次不再执行，把上次结果交还；连续几步拿不到新信息就收尾；
|
||
- 令牌 / 金额预算，可以设成不限；
|
||
- 上下文接近窗口先压缩早期工具结果，再不够就收尾；
|
||
- 临近上限提醒模型一次；
|
||
- 下游校验失败时写明是上游哪个 agent 提前收了尾；
|
||
- 升级前发起的运行（没有快照）一字不改走旧逻辑。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
import pytest
|
||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
|
||
from langchain_core.outputs import ChatGenerationChunk
|
||
from sqlalchemy import select
|
||
|
||
from app.db.base import SessionLocal
|
||
from app.db.models import CustomTool, Run, RunEvent, Setting
|
||
from app.engine import runner as runner_mod
|
||
from app.engine.guards import node_limits, run_limits
|
||
from app.engine.runner import run_manager
|
||
from app.tools import custom as custom_tools
|
||
|
||
FINAL = "结论：都查完了。"
|
||
SETTLED = "收尾：按已经查到的部分。"
|
||
TOOL = "shop_lookup"
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
async def engine_up():
|
||
await run_manager.setup()
|
||
yield
|
||
await run_manager.shutdown()
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
async def clean_run_settings():
|
||
async def drop():
|
||
async with SessionLocal() as session:
|
||
row = await session.get(Setting, "run")
|
||
if row:
|
||
await session.delete(row)
|
||
await session.commit()
|
||
await drop()
|
||
yield
|
||
await drop()
|
||
|
||
|
||
@pytest.fixture
|
||
async def lookup(monkeypatch):
|
||
"""自定义工具：按参数返回一段文字，长度可调。真正的请求换成记账。"""
|
||
calls: list[dict] = []
|
||
size = {"chars": 20}
|
||
|
||
async def fake_http(row, args):
|
||
calls.append(dict(args))
|
||
return f"门店 {args.get('q')}：" + "数" * size["chars"]
|
||
|
||
monkeypatch.setattr(custom_tools, "_run_http", fake_http)
|
||
async with SessionLocal() as session:
|
||
session.add(CustomTool(name=TOOL, kind="http", description="查门店",
|
||
parameters={"type": "object", "properties": {"q": {"type": "string"}},
|
||
"required": ["q"]},
|
||
config={"url": "https://example.com/shop"}))
|
||
await session.commit()
|
||
yield calls, size
|
||
async with SessionLocal() as session:
|
||
for row in (await session.execute(select(CustomTool).where(CustomTool.name == TOOL))).scalars():
|
||
await session.delete(row)
|
||
await session.commit()
|
||
|
||
|
||
def graph(**config):
|
||
cfg = {"prompt": "查门店", "tools": [TOOL], "approval": "never", **config}
|
||
return {
|
||
"nodes": [
|
||
{"id": "in", "type": "input", "data": {"config": {}}},
|
||
{"id": "work", "type": "agent", "data": {"label": "查数", "config": cfg}},
|
||
{"id": "out", "type": "output", "data": {"config": {}}},
|
||
],
|
||
"edges": [{"source": "in", "target": "work"}, {"source": "work", "target": "out"}],
|
||
}
|
||
|
||
|
||
def script(monkeypatch, plan, *, usage=None):
|
||
"""plan(第几次调用, messages) → 查询参数（字符串）或 None（给结论）。
|
||
|
||
流式和非流式都换掉：mock 走流式时不报用量，预算的测试要自己带上（usage）。
|
||
seen 记下每次调用模型时看到的 messages。
|
||
"""
|
||
from app.providers import mock_model
|
||
|
||
seen: list[list] = []
|
||
|
||
def decide(self, messages):
|
||
if not self.tools:
|
||
return AIMessage(content=SETTLED)
|
||
seen.append(list(messages))
|
||
q = plan(len(seen) - 1, messages)
|
||
if q is None:
|
||
return AIMessage(content=FINAL)
|
||
return AIMessage(content="再查一个。", tool_calls=[
|
||
{"name": TOOL, "args": {"q": q}, "id": f"call_{len(seen)}"}])
|
||
|
||
async def astream(self, messages, stop=None, run_manager=None, **kwargs):
|
||
reply = decide(self, messages)
|
||
extra = {"usage_metadata": usage(messages, reply)} if usage else {}
|
||
yield ChatGenerationChunk(message=AIMessageChunk(
|
||
content=reply.content, tool_calls=reply.tool_calls, **extra))
|
||
|
||
monkeypatch.setattr(mock_model.MockChatModel, "_decide", decide)
|
||
monkeypatch.setattr(mock_model.MockChatModel, "_astream", astream)
|
||
return seen
|
||
|
||
|
||
async def run(g, *, settings_run: dict | None = None):
|
||
if settings_run is not None:
|
||
async with SessionLocal() as session:
|
||
session.add(Setting(key="run", value=settings_run))
|
||
await session.commit()
|
||
started = await run_manager.start(graph=g, input_payload={"question": "有几家门店"})
|
||
for _ in range(400):
|
||
await asyncio.sleep(0.05)
|
||
async with SessionLocal() as session:
|
||
row = await session.get(Run, started.id)
|
||
if row.status in ("succeeded", "failed", "cancelled", "interrupted"):
|
||
break
|
||
else:
|
||
pytest.fail("运行没有在 20 秒内结束")
|
||
async with SessionLocal() as session:
|
||
events = list((await session.execute(
|
||
select(RunEvent).where(RunEvent.run_id == started.id).order_by(RunEvent.seq))).scalars())
|
||
return row, events
|
||
|
||
|
||
def logs(events, code):
|
||
return [e.data for e in events if e.type == "log" and e.data.get("code") == code]
|
||
|
||
|
||
def node_out(row):
|
||
return row.output
|
||
|
||
|
||
# ---------------------------------------------------------------- 上限怎么定
|
||
|
||
|
||
def test_run_limits_from_settings():
|
||
assert run_limits({}, 100) == {"max_steps": 100, "budget_tokens": 2_000_000, "budget_usd": None}
|
||
# 写 null / 0 / 空都是不限；存坏了的步数按缺省
|
||
assert run_limits({"agent_budget_tokens": None, "agent_budget_usd": 0}, 100)["budget_tokens"] is None
|
||
assert run_limits({"agent_budget_tokens": ""}, 100)["budget_tokens"] is None
|
||
assert run_limits({"agent_max_steps": "abc"}, 100)["max_steps"] == 100
|
||
# 环境变量的硬顶压过设置
|
||
assert run_limits({"agent_max_steps": 500}, 100)["max_steps"] == 100
|
||
assert run_limits({"agent_budget_usd": 1.5}, 100)["budget_usd"] == 1.5
|
||
|
||
|
||
def test_node_overrides_and_the_legacy_twelve():
|
||
snap = {"max_steps": 100, "budget_tokens": 2_000_000, "budget_usd": None}
|
||
# 画布以前自动写进去的 12：当作没填
|
||
assert node_limits({"max_steps": 12}, snap, 100).max_steps == 100
|
||
assert node_limits({"max_steps": "12"}, snap, 100).max_steps == 100
|
||
assert node_limits({"max_steps": 20}, snap, 100).max_steps == 20
|
||
assert node_limits({"max_steps": 300}, snap, 100).max_steps == 100
|
||
got = node_limits({"budget_tokens": 5000, "budget_usd": 0.5}, snap, 100)
|
||
assert (got.budget_tokens, got.budget_usd) == (5000, 0.5)
|
||
assert node_limits({}, {**snap, "budget_tokens": None}, 100).budget_tokens is None
|
||
|
||
|
||
async def test_settings_expose_the_defaults():
|
||
from httpx import ASGITransport, AsyncClient
|
||
|
||
from app.main import app
|
||
|
||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
|
||
got = (await c.get("/api/settings")).json()["run"]
|
||
assert (got["agent_max_steps"], got["agent_budget_tokens"], got["agent_budget_usd"]) == (100, 2_000_000, None)
|
||
|
||
|
||
# ---------------------------------------------------------------- 步数
|
||
|
||
|
||
async def test_twelve_steps_no_longer_cut_it_short(lookup, monkeypatch):
|
||
"""查 15 个门店再下结论：以前第 12 步就被掐断，现在跑完。节点上写着 12 也一样。"""
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: f"店{i}" if i < 15 else None)
|
||
row, events = await run(graph(max_steps=12))
|
||
assert row.status == "succeeded", row.error
|
||
assert len(calls) == 15
|
||
assert not logs(events, "step_limit_settled")
|
||
assert "limited" not in str(row.output)
|
||
|
||
|
||
async def test_an_explicit_step_limit_still_settles(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: f"店{i}")
|
||
row, events = await run(graph(max_steps=4))
|
||
assert row.status == "succeeded", row.error
|
||
assert len(calls) == 4
|
||
[settled] = logs(events, "step_limit_settled")
|
||
assert settled["reason"] == "steps" and "4 步" in settled["message"]
|
||
|
||
|
||
async def test_the_settings_default_applies(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: f"店{i}")
|
||
row, events = await run(graph(), settings_run={"agent_max_steps": 6})
|
||
assert len(calls) == 6
|
||
assert logs(events, "step_limit_settled")[0]["reason"] == "steps"
|
||
|
||
|
||
async def test_run_started_records_the_limits(lookup, monkeypatch):
|
||
script(monkeypatch, lambda i, _: None)
|
||
row, events = await run(graph(), settings_run={"agent_budget_tokens": None, "agent_max_steps": 30})
|
||
started = next(e for e in events if e.type == "run.started")
|
||
assert started.data["agent_limits"] == {"max_steps": 30, "budget_tokens": None, "budget_usd": None}
|
||
|
||
|
||
# ---------------------------------------------------------------- 重复与停滞
|
||
|
||
|
||
async def test_a_repeated_call_is_answered_from_the_first_result(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
seen = script(monkeypatch, lambda i, _: ["店A", "店A", None][i])
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert calls == [{"q": "店A"}], "同样的调用又执行了一次"
|
||
[repeat] = logs(events, "tool_repeat")
|
||
assert TOOL in repeat["message"]
|
||
reused = [m for m in seen[2] if isinstance(m, ToolMessage)][-1].content
|
||
assert "已经执行过" in reused and "门店 店A" in reused
|
||
|
||
|
||
async def test_going_in_circles_settles(lookup, monkeypatch):
|
||
"""一直用同样的参数查：第一次真查，之后每步都拿不到新信息，连续 3 步就收尾。"""
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: "店A")
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert calls == [{"q": "店A"}]
|
||
assert len(logs(events, "tool_repeat")) == 3, "没进展的 3 步之后就该收尾，不能一路耗到步数兜底"
|
||
[settled] = logs(events, "step_limit_settled")
|
||
assert settled["reason"] == "stall"
|
||
assert SETTLED in str(row.output)
|
||
|
||
|
||
async def test_new_information_resets_the_stall_count(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
# A、A、A（两步没进展）、B（有进展，清零）、B、B（又两步）、结论
|
||
order = ["店A", "店A", "店A", "店B", "店B", "店B", None]
|
||
script(monkeypatch, lambda i, _: order[i])
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert calls == [{"q": "店A"}, {"q": "店B"}]
|
||
assert not logs(events, "step_limit_settled")
|
||
|
||
|
||
# ---------------------------------------------------------------- 预算
|
||
|
||
|
||
def per_call(tokens):
|
||
return lambda messages, reply: {"input_tokens": tokens, "output_tokens": 0, "total_tokens": tokens}
|
||
|
||
|
||
async def test_token_budget_settles(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: f"店{i}", usage=per_call(1000))
|
||
row, events = await run(graph(budget_tokens=3500))
|
||
assert row.status == "succeeded", row.error
|
||
# 第 4 次调用之后累计 4000，超过 3500：那一步的工具照常跑完，然后收尾
|
||
assert len(calls) == 4
|
||
[settled] = logs(events, "step_limit_settled")
|
||
assert settled["reason"] == "budget_tokens" and "3,500" in settled["message"]
|
||
|
||
|
||
async def test_budget_can_be_unlimited(lookup, monkeypatch):
|
||
calls, _ = lookup
|
||
# 用量记在输出上：输入一大就先撞上下文水位了，这里只测预算
|
||
script(monkeypatch, lambda i, _: f"店{i}" if i < 8 else None,
|
||
usage=lambda m, r: {"input_tokens": 10, "output_tokens": 1_000_000, "total_tokens": 1_000_010})
|
||
row, events = await run(graph(), settings_run={"agent_budget_tokens": None})
|
||
assert row.status == "succeeded", row.error
|
||
assert len(calls) == 8 and not logs(events, "step_limit_settled")
|
||
|
||
|
||
async def test_usd_budget_settles_when_prices_are_known(lookup, monkeypatch):
|
||
from app.providers import catalog
|
||
|
||
monkeypatch.setattr(catalog, "estimate_cost", lambda model, i, o: i / 1000 * 0.01)
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: f"店{i}", usage=per_call(1000))
|
||
row, events = await run(graph(budget_usd=0.025))
|
||
assert len(calls) == 3
|
||
assert logs(events, "step_limit_settled")[0]["reason"] == "budget_usd"
|
||
|
||
|
||
async def test_reminder_before_the_budget_runs_out(lookup, monkeypatch):
|
||
seen = script(monkeypatch, lambda i, _: f"店{i}" if i < 5 else None, usage=per_call(1000))
|
||
row, _ = await run(graph(budget_tokens=5000))
|
||
notes = [m.content for call in seen for m in call if isinstance(m, HumanMessage) and "系统提醒" in m.content]
|
||
assert notes and "令牌预算已用 80%" in notes[0]
|
||
assert len(set(notes)) == 1, "同一种提醒只该说一次"
|
||
|
||
|
||
async def test_reminder_before_the_steps_run_out(lookup, monkeypatch):
|
||
seen = script(monkeypatch, lambda i, _: f"店{i}")
|
||
await run(graph(max_steps=8))
|
||
# 第 5 步（下标 4）之后还剩 3 步：第 6 次调用时看得到提醒
|
||
assert any("还剩 3 步" in m.content for m in seen[5] if isinstance(m, HumanMessage))
|
||
assert not any("系统提醒" in m.content for m in seen[4] if isinstance(m, HumanMessage))
|
||
|
||
|
||
# ---------------------------------------------------------------- 上下文
|
||
|
||
|
||
async def test_old_tool_results_are_compressed_near_the_window(lookup, monkeypatch):
|
||
from app.providers import catalog
|
||
|
||
calls, size = lookup
|
||
size["chars"] = 3000
|
||
# 没报用量时按字数的一半估：每条结果约 1500 token，第 6 次调用前后到窗口的 60%
|
||
monkeypatch.setattr(catalog, "find_context", lambda model: 12_000)
|
||
seen = script(monkeypatch, lambda i, _: f"店{i}" if i < 7 else None)
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert logs(events, "context_compressed")
|
||
last = [m.content for m in seen[-1] if isinstance(m, ToolMessage)]
|
||
assert any("早期的工具结果已压缩" in c for c in last[:-3])
|
||
assert all("早期的工具结果已压缩" not in c for c in last[-3:]), "最近 3 条不该压缩"
|
||
|
||
|
||
async def test_too_close_to_the_window_settles(lookup, monkeypatch):
|
||
from app.providers import catalog
|
||
|
||
calls, size = lookup
|
||
size["chars"] = 3000
|
||
monkeypatch.setattr(catalog, "find_context", lambda model: 1_000)
|
||
script(monkeypatch, lambda i, _: f"店{i}")
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert logs(events, "step_limit_settled")[0]["reason"] == "context"
|
||
assert len(calls) <= 2
|
||
|
||
|
||
# ---------------------------------------------------------------- 下游校验与老运行
|
||
|
||
|
||
async def test_validation_failure_names_the_agent_that_was_cut_short(lookup, monkeypatch):
|
||
script(monkeypatch, lambda i, _: f"店{i}")
|
||
g = graph(max_steps=3)
|
||
g["nodes"][2] = {"id": "check", "type": "validate", "data": {"label": "校验", "config": {
|
||
"source": "{{ nodes.work.text }}", "repair_with_llm": False,
|
||
"schema": {"type": "object", "required": ["stores"], "properties": {"stores": {"type": "array"}}}}}}
|
||
g["edges"][1] = {"source": "work", "target": "check"}
|
||
row, _ = await run(g)
|
||
assert row.status == "failed"
|
||
assert "上游「查数」用满了步数" in (row.error or ""), row.error
|
||
|
||
|
||
async def test_runs_from_before_the_upgrade_keep_the_old_rules(lookup, monkeypatch):
|
||
"""没有快照的运行：默认 12 步，重复调用照样执行——它们可能正停在某个审批上，
|
||
多一次复用、少一次调用，重放时 checkpoint 就对错号了。"""
|
||
async def no_snapshot(session):
|
||
return None
|
||
|
||
monkeypatch.setattr(runner_mod, "_agent_limits", no_snapshot)
|
||
calls, _ = lookup
|
||
script(monkeypatch, lambda i, _: "店A")
|
||
row, events = await run(graph())
|
||
assert row.status == "succeeded", row.error
|
||
assert len(calls) == 12
|
||
[settled] = logs(events, "step_limit_settled")
|
||
assert "reason" not in settled and "用满了 12 步" in settled["message"]
|
||
assert not logs(events, "tool_repeat")
|
||
|
||
|
||
def test_review_signal_carries_the_reason():
|
||
"""问数据页只在步数用满时给「放宽步数重跑」：复核信号要带上收尾原因。"""
|
||
from app.engine.review import scan
|
||
|
||
def log(code, reason=None):
|
||
data = {"level": "warn", "code": code, "message": f"{code} {reason}"}
|
||
if reason:
|
||
data["reason"] = reason
|
||
return {"type": "log", "node_id": "work", "data": data}
|
||
|
||
got = {s.detail: s.as_dict() for s in scan([log("step_limit_settled", "stall"),
|
||
log("step_limit_settled", "steps"),
|
||
log("step_limit_settled")], {"a": "答案"})}
|
||
assert got["step_limit_settled stall"]["reason"] == "stall"
|
||
assert got["step_limit_settled steps"]["reason"] == "steps"
|
||
assert "reason" not in got["step_limit_settled None"], "老运行的信号不该多出字段"
|
||
|
||
|
||
async def test_settings_limits_reflect_the_real_config():
|
||
"""库里存着旧版设置页写进去的 limits（见过 max_agent_steps: 25）：不生效，也不能盖住真实上限。"""
|
||
from httpx import ASGITransport, AsyncClient
|
||
|
||
from app.core.config import settings as app_settings
|
||
from app.main import app
|
||
|
||
async with SessionLocal() as session:
|
||
row = await session.get(Setting, "limits")
|
||
if row:
|
||
row.value = {"max_agent_steps": 25}
|
||
else:
|
||
session.add(Setting(key="limits", value={"max_agent_steps": 25}))
|
||
await session.commit()
|
||
try:
|
||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
|
||
assert (await c.get("/api/settings")).json()["limits"]["max_agent_steps"] == app_settings.max_agent_steps
|
||
await c.put("/api/settings", json={"values": {"limits": {"max_agent_steps": 7}}})
|
||
async with SessionLocal() as session:
|
||
assert (await session.get(Setting, "limits")).value == {"max_agent_steps": 25}, "limits 不该被写"
|
||
finally:
|
||
async with SessionLocal() as session:
|
||
row = await session.get(Setting, "limits")
|
||
if row:
|
||
await session.delete(row)
|
||
await session.commit()
|