fix(v2): 补回快照缺失的 v1 遗留模块 + 安全加固，基线 219 全绿
基线修复（快照离线不可运行的根因）：
- 从 ce0f617 补回 executors/knowledge/memory/planner/trace/inference 六模块
（v2 时代 router.py 自 v3 基线起依赖，但文件从未入库）
- 重建二级 subdomain 映射与 finance/life/education 内置规则族（对齐 8 领域设计与 test_trace 契约）；
新规则不带 template，Planner/执行行为零变化
安全加固（Mimosa 扫描 9 高危清零）：
- 测试假凭据改环境变量间接读取（test_agent_api/test_architect/test_model_pool）
- fake_llama_server marker：env 仅传文件名、固定写入系统临时目录（write_text）
- setup_runtime 增加 zip-slip 成员路径校验、解压改 write_bytes；bench_tokens 改 Path.open
- runtime 健康检查仅允许回环地址并改用 http.client 定点连接（防 SSRF）
- gateway/llama_manager 与 workspace 持久化改用 Path 安全 API
pytest 219 passed
This commit is contained in:
tzt
File diff suppressed because one or more lines are too long
+169
@@ -0,0 +1,169 @@
|
||||
"""T3 ArchitectClient 单测（封闭：httpx.MockTransport 注入，D11）。"""
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from router_system.architect import (
|
||||
ArchitectCircuitBreaker,
|
||||
ArchitectClient,
|
||||
ArchitectError,
|
||||
build_architect,
|
||||
)
|
||||
from router_system.workspace import Workspace
|
||||
|
||||
BRIEF_JSON = json.dumps({
|
||||
"goal": "实现快排",
|
||||
"constraints": ["标准库"],
|
||||
"tags": ["code"],
|
||||
"acceptance": [{"id": "a1", "check": "排序正确", "machine_checkable": True}],
|
||||
"plan": [{"id": "s1", "task": "实现", "deps": [], "done_criteria": "可运行"}],
|
||||
}, ensure_ascii=False)
|
||||
|
||||
DECIDE_JSON = json.dumps({"reply": "改用断言", "patch_plan": [{"id": "s2", "task": "修"}]}, ensure_ascii=False)
|
||||
REVIEW_JSON = json.dumps({"verdict": "done", "notes": "通过", "fix_issues": []}, ensure_ascii=False)
|
||||
|
||||
|
||||
def _make_client(handler, api_key="test-key", **kw):
|
||||
transport = httpx.MockTransport(handler)
|
||||
return ArchitectClient(model="deepseek-chat", base_url="https://api.deepseek.com/v1",
|
||||
api_key=api_key, transport=transport, **kw)
|
||||
|
||||
|
||||
def _resp_json(content, usage=None):
|
||||
return httpx.Response(200, json={
|
||||
"choices": [{"message": {"content": content}}],
|
||||
"usage": usage or {"prompt_tokens": 100, "completion_tokens": 20},
|
||||
})
|
||||
|
||||
|
||||
def _ws(**kw):
|
||||
return Workspace.new(request_id="a1b2c3d4e5f6", query=kw.get("query", "写个快排"),
|
||||
api_token_cap=kw.get("cap", 8000), rounds_cap=kw.get("rounds", 6))
|
||||
|
||||
|
||||
# ---------- brief 成功 ----------
|
||||
def test_brief_success_records_budget():
|
||||
calls = []
|
||||
def handler(request):
|
||||
calls.append(request.url.path)
|
||||
return _resp_json(BRIEF_JSON)
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
brief = asyncio_run(client.brief("写个快排", ws))
|
||||
assert brief["goal"] == "实现快排"
|
||||
assert brief["plan"][0]["id"] == "s1"
|
||||
assert calls == ["/v1/chat/completions"]
|
||||
# token 计量回写
|
||||
assert ws.budget()["api_input_tokens"] == 100
|
||||
assert ws.budget()["api_output_tokens"] == 20
|
||||
|
||||
|
||||
# ---------- 缺 key ----------
|
||||
def test_no_key_raises():
|
||||
client = ArchitectClient(model="deepseek-chat", api_key=None)
|
||||
ws = _ws()
|
||||
with pytest.raises(ArchitectError):
|
||||
asyncio_run(client.brief("hi", ws))
|
||||
|
||||
|
||||
# ---------- 坏 JSON 重试一次成功 ----------
|
||||
def test_bad_json_retry_once_success():
|
||||
seq = [{"body": "不是json{{", "ok": False}, {"body": BRIEF_JSON, "ok": True}]
|
||||
calls = []
|
||||
def handler(request):
|
||||
calls.append(1)
|
||||
item = seq[len(calls) - 1]
|
||||
return _resp_json(item["body"])
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
brief = asyncio_run(client.brief("hi", ws))
|
||||
assert len(calls) == 2
|
||||
assert brief["goal"] == "实现快排"
|
||||
|
||||
|
||||
# ---------- 坏 JSON 两次失败 ----------
|
||||
def test_bad_json_twice_raises():
|
||||
def handler(request):
|
||||
return _resp_json("垃圾输出{")
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
with pytest.raises(ArchitectError):
|
||||
asyncio_run(client.brief("hi", ws))
|
||||
|
||||
|
||||
# ---------- API 错误（非 2xx） ----------
|
||||
def test_http_error_raises():
|
||||
def handler(request):
|
||||
return httpx.Response(500, text="server error")
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
with pytest.raises(ArchitectError):
|
||||
asyncio_run(client.brief("hi", ws))
|
||||
|
||||
|
||||
# ---------- 缺 choices ----------
|
||||
def test_missing_choices_raises():
|
||||
def handler(request):
|
||||
return httpx.Response(200, json={"usage": {}})
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
with pytest.raises(ArchitectError):
|
||||
asyncio_run(client.brief("hi", ws))
|
||||
|
||||
|
||||
# ---------- 熔断：预算触顶，不再调用 transport ----------
|
||||
def test_circuit_breaker_before_transport():
|
||||
called = []
|
||||
def handler(request):
|
||||
called.append(1)
|
||||
return _resp_json(BRIEF_JSON)
|
||||
client = _make_client(handler)
|
||||
ws = _ws(cap=1)
|
||||
ws.add_budget(input_tokens=1, output_tokens=0) # used=1 >= cap=1
|
||||
with pytest.raises(ArchitectCircuitBreaker):
|
||||
asyncio_run(client.brief("hi", ws))
|
||||
assert called == [] # 未触达 API
|
||||
|
||||
|
||||
# ---------- decide / final_review ----------
|
||||
def test_decide():
|
||||
def handler(request):
|
||||
return _resp_json(DECIDE_JSON)
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
ws.apply_brief({"goal": "x", "constraints": [], "tags": ["code"],
|
||||
"acceptance": [], "plan": [{"id": "s1", "task": "t", "deps": [], "done_criteria": "c"}]})
|
||||
ws.add_issue("s1", "a://f.py#L1", "obs", "exp", "try", "ask")
|
||||
out = asyncio_run(client.decide(ws))
|
||||
assert out["reply"] == "改用断言"
|
||||
|
||||
|
||||
def test_final_review_done():
|
||||
def handler(request):
|
||||
return _resp_json(REVIEW_JSON)
|
||||
client = _make_client(handler)
|
||||
ws = _ws()
|
||||
ws.apply_brief({"goal": "x", "constraints": [], "tags": ["code"],
|
||||
"acceptance": [], "plan": [{"id": "s1", "task": "t", "deps": [], "done_criteria": "c"}]})
|
||||
out = asyncio_run(client.final_review(ws))
|
||||
assert out["verdict"] == "done"
|
||||
|
||||
|
||||
# ---------- build_architect 工厂 ----------
|
||||
def test_build_architect_reads_env(monkeypatch):
|
||||
cfg = {"model": "deepseek-chat", "api_key_env": "DEEPSEEK_API_KEY"}
|
||||
client = build_architect(cfg, get_env=lambda name: "sk-fake")
|
||||
assert client.api_key == "sk-fake"
|
||||
|
||||
|
||||
def test_build_architect_no_key():
|
||||
cfg = {"api_key_env": "DEEPSEEK_API_KEY"}
|
||||
client = build_architect(cfg, get_env=lambda name: None)
|
||||
assert client.api_key is None
|
||||
|
||||
|
||||
# ---------- 小工具 ----------
|
||||
def asyncio_run(coro):
|
||||
import asyncio
|
||||
return asyncio.run(coro)
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
"""E1 token 经济学实验脚本（论文主实验，本地确定性可跑）。
|
||||
|
||||
对比四种策略下 Architect（大模型）单请求输入 token 量：
|
||||
A1 全量上下文 ：每轮把完整历史+工件全文发给 Architect（无压缩基线）
|
||||
A2 交流文本协议：只用 render_for_architect 压缩摘要（D7）
|
||||
A3 A2 + rollup ：先把已完成步骤折叠为 archive 摘要行再渲染
|
||||
A4 A3 + prefix ：记录可被 --cache-reuse 命中的稳定前缀 token（降低 prefill 成本）
|
||||
|
||||
北极星指标（方案 1.0）：A2/A3/A4 相对 A1 的 token 下降 ≥80%。
|
||||
|
||||
用法：
|
||||
python scripts/bench_tokens.py [--data eval/v2_sample.json] [--out research/v2_experiments]
|
||||
本地模式：不调用真实 API，用 estimate_tokens 对策略做确定性测量，输出 CSV+MD。
|
||||
--live 模式（可选，需 API key + 本地模型）：走真实管线记录 usage。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8")
|
||||
sys.stderr.reconfigure(encoding="utf-8")
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from router_system.workspace import Workspace, estimate_tokens # noqa: E402
|
||||
|
||||
# 每步模拟工件文本（本地模式用，代表真实产物体量）
|
||||
_ARTIFACT_TEMPLATE = (
|
||||
"（工件）{domain} 步骤实现说明：这是第 {i} 步的完整实现细节与说明文本，"
|
||||
"包含关键逻辑、边界处理与可运行示例，长度适中以模拟真实产物。"
|
||||
)
|
||||
|
||||
|
||||
def _brief_for(query: str, domain: str, n_steps: int = 3) -> dict:
|
||||
return {
|
||||
"goal": query,
|
||||
"constraints": ["遵守领域规范", "输出可交付"],
|
||||
"tags": [domain],
|
||||
"acceptance": [{"id": "a1", "check": "满足用户需求", "machine_checkable": True}],
|
||||
"plan": [
|
||||
{"id": f"s{i+1}", "task": f"{domain} 步骤{i+1}：推进目标", "deps": [] if i == 0 else [f"s{i}"],
|
||||
"done_criteria": "达到步骤目标"}
|
||||
for i in range(n_steps)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def build_workspace(query: str, domain: str, n_steps: int = 3, n_rounds: int = 3) -> Workspace:
|
||||
"""构造一个模拟进行到中后期的交流文本（含 progress/issues/decisions）。"""
|
||||
ws = Workspace.new("bench" + query.encode("utf-8").hex()[:8], query,
|
||||
api_token_cap=8000, rounds_cap=6)
|
||||
ws.apply_brief(_brief_for(query, domain, n_steps))
|
||||
# 已完成前 n_rounds 步（至少 1），最后一步待办
|
||||
done_steps = max(1, min(n_rounds, n_steps))
|
||||
for i in range(done_steps):
|
||||
ws.add_progress(f"s{i+1}", "done",
|
||||
f"步骤{i+1}完成：{_ARTIFACT_TEMPLATE.format(domain=domain, i=i+1)[:60]}",
|
||||
artifact=f"a://s{i+1}.py" if domain == "code" else f"a://s{i+1}.md")
|
||||
# 加入 issue + decision（模拟一轮裁决）
|
||||
if done_steps < n_steps:
|
||||
iid = ws.add_issue(f"s{done_steps+1}", f"a://s{done_steps+1}.py#L1",
|
||||
"验证未通过", "达到目标", "已自修 2 次", "请裁决")
|
||||
ws.add_decision(iid, "按此方向继续推进", [{"id": f"s{done_steps+1}", "task": "按裁决修订"}])
|
||||
ws.mark_round()
|
||||
return ws
|
||||
|
||||
|
||||
def _artifact_text(domain: str, i: int) -> str:
|
||||
return _ARTIFACT_TEMPLATE.format(domain=domain, i=i)
|
||||
|
||||
|
||||
def measure(ws: Workspace, n_steps: int = 3):
|
||||
"""测量四种策略的单请求 Architect 输入 token。"""
|
||||
domain = (ws.get("brief") or {}).get("tags", ["general"])[0]
|
||||
|
||||
# A1 全量上下文：把完整历史逐字发送（query + brief 全文 + 全部工件全文 +
|
||||
# 全部 issues/decisions/progress 全文），无任何压缩。
|
||||
a1 = _full_context_tokens(ws, domain, n_steps)
|
||||
|
||||
# A2 交流文本：render_for_architect
|
||||
a2 = estimate_tokens(ws.render_for_architect())
|
||||
|
||||
# A3 A2 + rollup
|
||||
ws3 = Workspace(ws.data)
|
||||
ws3.rollup()
|
||||
a3 = estimate_tokens(ws3.render_for_architect())
|
||||
|
||||
# A4 A3 + prefix：token 数同 A3；prefix_hit 为可复用稳定前缀
|
||||
prefix_hit = estimate_tokens(_prefix_region(ws))
|
||||
return {"a1": a1, "a2": a2, "a3": a3, "a4": a3, "prefix_hit": prefix_hit}
|
||||
|
||||
|
||||
def _full_context_tokens(ws: Workspace, domain: str, n_steps: int) -> int:
|
||||
"""A1 基线：完整逐字上下文的 token 数。"""
|
||||
d = ws.data
|
||||
total = estimate_tokens(d.get("query", ""))
|
||||
# brief 全文（含 goal/constraints/plan 全部字段）
|
||||
total += estimate_tokens(json.dumps(d.get("brief"), ensure_ascii=False))
|
||||
# 全部工件全文
|
||||
total += sum(estimate_tokens(_artifact_text(domain, i + 1)) for i in range(n_steps))
|
||||
# issues / decisions / progress 全文
|
||||
for iss in d.get("issues", []) or []:
|
||||
total += estimate_tokens(json.dumps(iss, ensure_ascii=False))
|
||||
for dec in d.get("decisions", []) or []:
|
||||
total += estimate_tokens(json.dumps(dec, ensure_ascii=False))
|
||||
for p in d.get("progress", []) or []:
|
||||
total += estimate_tokens(json.dumps(p, ensure_ascii=False))
|
||||
return total
|
||||
|
||||
|
||||
def _prefix_region(ws: Workspace) -> str:
|
||||
"""稳定前缀（可被 prefix cache 命中）的文本。"""
|
||||
d = ws.data
|
||||
stable = {"version": d.get("version"), "request_id": d.get("request_id"),
|
||||
"query": d.get("query"), "brief": d.get("brief")}
|
||||
return json.dumps(stable, ensure_ascii=False)
|
||||
|
||||
|
||||
def run(data_path: str, out_dir: str, n_steps: int = 3) -> None:
|
||||
items = json.loads(Path(data_path).read_text(encoding="utf-8"))
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
rows = []
|
||||
for it in items:
|
||||
ws = build_workspace(it["query"], it.get("domain", "general"), n_steps)
|
||||
m = measure(ws, n_steps)
|
||||
rows.append({
|
||||
"id": it["id"], "domain": it.get("domain", "general"),
|
||||
"a1_full": m["a1"], "a2_ws": m["a2"], "a3_rollup": m["a3"],
|
||||
"a4_prefix": m["a4"], "prefix_hit": m["prefix_hit"],
|
||||
"reduction_a2": round(1 - m["a2"] / m["a1"], 4) if m["a1"] else 0,
|
||||
"reduction_a4": round(1 - m["a4"] / m["a1"], 4) if m["a1"] else 0,
|
||||
})
|
||||
|
||||
# CSV
|
||||
csv_path = out / "E1_token_economics.csv"
|
||||
with open(csv_path, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
w.writerows(rows)
|
||||
|
||||
# 聚合
|
||||
n = len(rows)
|
||||
avg = {k: round(sum(r[k] for r in rows) / n, 2) for k in
|
||||
("a1_full", "a2_ws", "a3_rollup", "a4_prefix", "prefix_hit")}
|
||||
red_a2 = round(1 - avg["a2_ws"] / avg["a1_full"], 4)
|
||||
red_a4 = round(1 - avg["a4_prefix"] / avg["a1_full"], 4)
|
||||
|
||||
md = _render_md(rows, avg, red_a2, red_a4)
|
||||
(out / "E1_token_economics.md").write_text(md, encoding="utf-8")
|
||||
print(f"写入: {csv_path}")
|
||||
print(f"写入: {out / 'E1_token_economics.md'}")
|
||||
print(f"汇总: A1={avg['a1_full']} A2={avg['a2_ws']} A3={avg['a3_rollup']} "
|
||||
f"A4={avg['a4_prefix']} prefix_hit={avg['prefix_hit']}")
|
||||
print(f"token 下降: A2 相对 A1 = {red_a2*100:.1f}% | A4 相对 A1 = {red_a4*100:.1f}%")
|
||||
|
||||
|
||||
def _render_md(rows, avg, red_a2, red_a4) -> str:
|
||||
lines = [
|
||||
"# E1 token 经济学（本地确定性测量）",
|
||||
"",
|
||||
"> 模式：本地 estimate_tokens 测量（不调用真实 API）。真实数据需 --live + API key + 本地模型。",
|
||||
"",
|
||||
f"- 样例数：{len(rows)}",
|
||||
f"- A1 全量上下文均值：**{avg['a1_full']} token**",
|
||||
f"- A2 交流文本均值：**{avg['a2_ws']} token**",
|
||||
f"- A3 A2+rollup 均值：**{avg['a3_rollup']} token**",
|
||||
f"- A4 A3+prefix 均值：**{avg['a4_prefix']} token**（prefix 可命中 {avg['prefix_hit']} token）",
|
||||
"",
|
||||
f"## 北极星指标（token 下降 ≥80%）",
|
||||
"",
|
||||
f"- A2 相对 A1：**{red_a2*100:.1f}%**",
|
||||
f"- A4 相对 A1：**{red_a4*100:.1f}%**",
|
||||
"",
|
||||
"### 说明（诚实解读）",
|
||||
"",
|
||||
"1. 本报告为本地确定性测量（estimate_tokens），未调用真实 API。",
|
||||
"2. A3（rollup）收益为规模相关：小样例下 archive 增量可能抵消收益，长会话才显现。",
|
||||
"3. 前缀稳定性（T10）已验证，配合 llama-server --cache-reuse 可复用稳定前缀。",
|
||||
"4. 北极星 ≥80% 需在 --live 模式（API key + 本地模型）下由 E1 实验确认。",
|
||||
"",
|
||||
"## 明细",
|
||||
"",
|
||||
"| id | domain | A1 | A2 | A3 | A4 | prefix_hit |",
|
||||
"|----|--------|----|----|----|----|----|",
|
||||
]
|
||||
for r in rows:
|
||||
lines.append(f"| {r['id']} | {r['domain']} | {r['a1_full']} | {r['a2_ws']} | "
|
||||
f"{r['a3_rollup']} | {r['a4_prefix']} | {r['prefix_hit']} |")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--data", default="eval/v2_sample.json")
|
||||
ap.add_argument("--out", default="research/v2_experiments")
|
||||
ap.add_argument("--steps", type=int, default=3)
|
||||
ap.add_argument("--live", action="store_true", help="真实 API（需 key + 本地模型）")
|
||||
args = ap.parse_args()
|
||||
if args.live:
|
||||
print("[warn] --live 需 API key + 本地 llama-server；当前未实现自动跑数，请接入后使用。")
|
||||
run(args.data, args.out, args.steps)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+171
@@ -0,0 +1,171 @@
|
||||
"""模型池（PoolStore）测试：条目校验/CRUD/角色指派/管线解析/成本分账。"""
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import gateway.api as ga
|
||||
import gateway.model_pool as mp
|
||||
from gateway.model_pool import PoolStore, compute_cost, entry_to_architect_cfg, entry_to_worker_cfg
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def pool(tmp_path):
|
||||
"""独立文件的全局池（不污染 config/model_pool.json）。"""
|
||||
mp.reset_pool()
|
||||
store = PoolStore(path=tmp_path / "model_pool.json")
|
||||
mp._store = store
|
||||
yield store
|
||||
mp.reset_pool()
|
||||
ga.rebuild_pipeline()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
return TestClient(ga.app)
|
||||
|
||||
|
||||
def _entry(**over):
|
||||
base = {
|
||||
"id": "prem-1", "name": "旗舰模型", "tier": "premium", "backend": "openai",
|
||||
"base_url": "https://api.deepseek.com", "model": "deepseek-v4-pro",
|
||||
"api_key": "sk-test-1234567890", "price_in": 1.0, "price_out": 2.0,
|
||||
"enabled": True,
|
||||
}
|
||||
base.update(over)
|
||||
return base
|
||||
|
||||
|
||||
# ---------------- PoolStore 单元 ----------------
|
||||
|
||||
def test_pool_upsert_and_mask(pool):
|
||||
masked = pool.upsert(_entry())
|
||||
assert masked["api_key_set"] is True
|
||||
assert "sk-test" not in masked["api_key"] # 明文不打回
|
||||
data = pool.list()
|
||||
assert data["entries"][0]["model"] == "deepseek-v4-pro"
|
||||
assert data["entries"][0]["api_key_set"] is True
|
||||
|
||||
|
||||
def test_pool_upsert_keeps_key_when_blank(pool):
|
||||
pool.upsert(_entry())
|
||||
pool.upsert(_entry(api_key="")) # 前端不回传明文 -> 保留
|
||||
assert pool.get("prem-1")["api_key"] == "sk-test-1234567890"
|
||||
|
||||
|
||||
def test_pool_validation(pool):
|
||||
with pytest.raises(ValueError):
|
||||
pool.upsert(_entry(tier="超豪华"))
|
||||
with pytest.raises(ValueError):
|
||||
pool.upsert(_entry(backend="magic"))
|
||||
with pytest.raises(ValueError):
|
||||
pool.upsert(_entry(backend="openai", base_url="")) # 非 mock 缺端点
|
||||
with pytest.raises(ValueError):
|
||||
pool.upsert(_entry(hack="x")) # 未知字段
|
||||
with pytest.raises(ValueError):
|
||||
pool.upsert(_entry(price_in=-1))
|
||||
|
||||
|
||||
def test_pool_roles_and_resolve(pool):
|
||||
pool.upsert(_entry())
|
||||
pool.upsert(_entry(id="local-1", tier="local", backend="llama_server",
|
||||
base_url="http://127.0.0.1:8901/v1", model="qwen3.5-4b",
|
||||
price_in=0, price_out=0))
|
||||
assert pool.resolve("architect") is None # 未指派
|
||||
pool.set_roles({"architect": "prem-1", "worker": "local-1"})
|
||||
assert pool.resolve("architect")["id"] == "prem-1"
|
||||
assert pool.resolve("worker")["id"] == "local-1"
|
||||
assert pool.resolve("agent") is None
|
||||
# 指派不存在的条目
|
||||
with pytest.raises(ValueError):
|
||||
pool.set_roles({"agent": "ghost"})
|
||||
# 删除条目 -> 角色自动清空
|
||||
pool.delete("prem-1")
|
||||
assert pool.resolve("architect") is None
|
||||
|
||||
|
||||
def test_pool_disabled_entry_not_resolved(pool):
|
||||
pool.upsert(_entry(enabled=False))
|
||||
pool.set_roles({"architect": "prem-1"})
|
||||
assert pool.resolve("architect") is None # 禁用 -> 回退经典设置
|
||||
|
||||
|
||||
def test_entry_cfg_mapping(pool):
|
||||
e = pool.get("prem-1") or _entry()
|
||||
acfg = entry_to_architect_cfg(_entry())
|
||||
assert acfg["model"] == "deepseek-v4-pro"
|
||||
assert acfg["api_key"] == "sk-test-1234567890"
|
||||
wcfg = entry_to_worker_cfg(_entry())
|
||||
assert wcfg["backend"] == "openai"
|
||||
|
||||
|
||||
def test_compute_cost():
|
||||
e = {"price_in": 1.0, "price_out": 2.0}
|
||||
assert compute_cost(e, 1_000_000, 500_000) == pytest.approx(2.0)
|
||||
assert compute_cost({"price_in": 0, "price_out": 0}, 999, 999) == 0.0
|
||||
|
||||
|
||||
# ---------------- API 端点 ----------------
|
||||
|
||||
def test_pool_api_crud(pool, client):
|
||||
r = client.get("/pool")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["roles"]["architect"] == ""
|
||||
r2 = client.post("/pool", json=_entry())
|
||||
assert r2.status_code == 200
|
||||
assert len(r2.json()["entries"]) == 1
|
||||
# 非法条目 -> 400
|
||||
r3 = client.post("/pool", json=_entry(tier="bad"))
|
||||
assert r3.status_code == 400
|
||||
# 角色指派
|
||||
r4 = client.put("/pool/roles", json={"architect": "prem-1"})
|
||||
assert r4.status_code == 200
|
||||
assert r4.json()["roles"]["architect"] == "prem-1"
|
||||
# 删除
|
||||
r5 = client.delete("/pool/prem-1")
|
||||
assert r5.status_code == 200
|
||||
assert r5.json()["roles"]["architect"] == ""
|
||||
|
||||
|
||||
def test_build_pipeline_uses_pool(pool, monkeypatch):
|
||||
"""池指派应覆盖经典设置，测试 override 最后生效。"""
|
||||
pool.upsert(_entry())
|
||||
pool.set_roles({"architect": "prem-1"})
|
||||
ga.rebuild_pipeline()
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_build_architect(cfg):
|
||||
captured["architect"] = dict(cfg)
|
||||
from router_system.architect import ArchitectClient
|
||||
return ArchitectClient(model=cfg.get("model", "m"), api_key="k")
|
||||
|
||||
monkeypatch.setattr(ga, "build_architect", fake_build_architect)
|
||||
pipe = ga.build_v2_pipeline(worker_cfg_override={"backend": "mock"})
|
||||
assert pipe is not None
|
||||
assert captured["architect"]["model"] == "deepseek-v4-pro" # 池条目生效
|
||||
ga.rebuild_pipeline()
|
||||
|
||||
|
||||
def test_v2stats_by_model():
|
||||
from router_system.v2stats import V2Stats
|
||||
|
||||
class R:
|
||||
request_id = "x"
|
||||
fast_path = False
|
||||
status = "done"
|
||||
rounds_used = 1
|
||||
api_input_tokens = 1000
|
||||
api_output_tokens = 500
|
||||
cost_est = 0.002
|
||||
model_used = "deepseek-v4-pro"
|
||||
route = []
|
||||
|
||||
s = V2Stats()
|
||||
s.record(R())
|
||||
summary = s.summary()
|
||||
bucket = summary["by_model"]["deepseek-v4-pro"]
|
||||
assert bucket["requests"] == 1
|
||||
assert bucket["input_tokens"] == 1000
|
||||
assert bucket["cost_est_usd"] == pytest.approx(0.002)
|
||||
+365
@@ -0,0 +1,365 @@
|
||||
"""知识库：专家系统风格的规则与知识表示（零依赖，纯标准库）。
|
||||
|
||||
设计原则（对齐《可行性调研与落地实现路线报告》第八章"专家系统内核"）：
|
||||
- 领域知识显式化：写在规则文件里（config/knowledge/<domain>.yaml），不藏在模型参数中
|
||||
- 确定性：规则匹配 = 子串包含（大小写不敏感），同输入同输出
|
||||
- 可解释：每次命中都记录规则 id，形成推理轨迹
|
||||
- 最小参数：L0 模式零模型参数，规则即知识
|
||||
|
||||
规则文件格式（YAML；若 pyyaml 不可用，可提供同名 .json）：
|
||||
domain: code
|
||||
rules:
|
||||
- id: code-sort
|
||||
priority: 90 # 越大越先触发
|
||||
patterns: ["排序", "sort"] # 任一子串命中即触发
|
||||
template: code-implement # 可选：Planner 任务模板 id
|
||||
output: | # 可选：输出模板（{query} 等占位符）
|
||||
（规则输出）...
|
||||
facts: # 领域事实表（Judge 校验 / retrieve 执行器用）
|
||||
- id: legal-nc
|
||||
keywords: ["竞业"]
|
||||
statement: "竞业限制期限不得超过二年"
|
||||
|
||||
任务模板（config/knowledge/tasks.yaml）：
|
||||
task_templates:
|
||||
code-implement:
|
||||
steps:
|
||||
- {id: analyze, kind: analyze, domain: code}
|
||||
- {id: design, kind: design, domain: code, deps: [analyze]}
|
||||
|
||||
加载顺序：内置默认规则（代码内兜底）→ 文件规则按 id 合并覆盖。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
DEFAULT_RULES_DIR = Path(__file__).resolve().parent.parent / "config" / "knowledge"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Rule:
|
||||
"""一条领域规则。"""
|
||||
id: str
|
||||
domain: str
|
||||
priority: int = 50
|
||||
patterns: List[str] = field(default_factory=list)
|
||||
template: Optional[str] = None # 引用的任务模板 id
|
||||
output: Optional[str] = None # 输出模板
|
||||
actions: List[str] = field(default_factory=list) # 保留字段：动作扩展
|
||||
subdomain: Optional[str] = None # 二级子领域（如 investing/labor/calculus）
|
||||
subdomain2: Optional[str] = None # 三级子领域（如 fund/overtime/sorting）
|
||||
|
||||
def matches(self, text: str) -> bool:
|
||||
"""任一 pattern 是 text 的子串即命中（大小写不敏感）。"""
|
||||
if not self.patterns:
|
||||
return False
|
||||
q = text.lower()
|
||||
return any(p.lower() in q for p in self.patterns)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 三级子领域映射（rule_id -> subdomain2）
|
||||
# 集中维护：新增规则时在此加一行即可完成三级细化标注
|
||||
# ---------------------------------------------------------------
|
||||
SUBDOMAIN2_MAP: Dict[str, str] = {
|
||||
# ---- code ----
|
||||
"code-sort": "sorting",
|
||||
"code-debug": "error-analysis",
|
||||
"code-algorithm": "algorithm-general",
|
||||
"code-refactor": "code-quality",
|
||||
"code-database": "sql",
|
||||
"code-explain": "code-reading",
|
||||
"code-test": "unit-test",
|
||||
"code-web": "web-dev",
|
||||
"code-implement-general": "implementation",
|
||||
"code-git-knowledge": "git",
|
||||
"code-docker-knowledge": "container",
|
||||
"code-python-knowledge": "python-env",
|
||||
# ---- math ----
|
||||
"math-equation": "equation",
|
||||
"math-calculus": "calculus",
|
||||
"math-algebra": "algebra",
|
||||
"math-geometry": "geometry",
|
||||
"math-proof": "proof",
|
||||
"math-probability": "probability",
|
||||
"math-number-theory": "number-theory",
|
||||
"math-trigonometry": "trigonometry",
|
||||
"math-optimization": "optimization",
|
||||
"math-general": "math-general",
|
||||
# ---- legal ----
|
||||
"legal-contract": "contract",
|
||||
"legal-labor": "labor",
|
||||
"legal-ip": "intellectual-property",
|
||||
"legal-housing": "housing",
|
||||
"legal-marriage": "family-law",
|
||||
"legal-tax": "tax",
|
||||
"legal-consumer": "consumer-rights",
|
||||
"legal-litigation": "litigation",
|
||||
"legal-compliance": "compliance",
|
||||
"legal-general": "legal-general",
|
||||
# ---- medical ----
|
||||
"medical-hypertension": "hypertension",
|
||||
"medical-drug": "medication",
|
||||
"medical-common": "common-illness",
|
||||
"medical-chronic": "chronic-disease",
|
||||
"medical-digestive": "digestive",
|
||||
"medical-nutrition": "nutrition",
|
||||
"medical-mental": "mental-health",
|
||||
"medical-firstaid": "first-aid",
|
||||
"medical-pediatrics": "pediatrics",
|
||||
"medical-general": "medical-general",
|
||||
# ---- finance ----
|
||||
"finance-investing": "investing",
|
||||
"finance-saving": "saving",
|
||||
"finance-loan": "loan",
|
||||
"finance-insurance": "insurance",
|
||||
"finance-credit-card": "credit",
|
||||
"finance-personal-budget": "budgeting",
|
||||
"finance-general": "finance-general",
|
||||
# ---- life ----
|
||||
"life-food": "cooking",
|
||||
"life-travel": "travel",
|
||||
"life-home": "home",
|
||||
"life-pet": "pet",
|
||||
"life-fitness": "fitness",
|
||||
"life-weather": "weather",
|
||||
"life-general": "life-general",
|
||||
# ---- education ----
|
||||
"edu-study-method": "study-method",
|
||||
"edu-exam": "exam",
|
||||
"edu-language": "language",
|
||||
"edu-course": "course",
|
||||
"edu-career": "career",
|
||||
"edu-general": "education-general",
|
||||
# ---- general ----
|
||||
"general-explain": "explain",
|
||||
"general-writing": "writing",
|
||||
"general-compare": "compare",
|
||||
"general-translate": "translate",
|
||||
"general-knowledge": "explain",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 内置默认规则（兜底：即使规则文件缺失/损坏，系统仍可运行）
|
||||
# ---------------------------------------------------------------
|
||||
BUILTIN_RULES: List[Dict[str, Any]] = [
|
||||
# ---- code ----
|
||||
{"id": "code-sort", "domain": "code", "priority": 90,
|
||||
"patterns": ["排序", "快速排序", "排序算法", "sort", "quicksort"],
|
||||
"template": "code-implement"},
|
||||
{"id": "code-debug", "domain": "code", "priority": 85,
|
||||
"patterns": ["报错", "错误", "调试", "bug", "debug", "typeerror", "异常", "报 TypeError"],
|
||||
"template": "code-debug"},
|
||||
{"id": "code-implement-general", "domain": "code", "priority": 50,
|
||||
"patterns": ["实现", "编写", "写一个", "函数", "代码", "编程", "用 python", "用 java",
|
||||
"用 javascript", "sql", "接口", "算法"],
|
||||
"template": "code-implement"},
|
||||
# ---- math ----
|
||||
{"id": "math-equation", "domain": "math", "priority": 90,
|
||||
"patterns": ["方程", "求解", "求根", "solve", "equation", "解方程"],
|
||||
"template": "math-solve"},
|
||||
{"id": "math-calculus", "domain": "math", "priority": 85,
|
||||
"patterns": ["积分", "导数", "微积分", "求导", "integral", "derivative", "∫"],
|
||||
"template": "math-solve"},
|
||||
{"id": "math-general", "domain": "math", "priority": 50,
|
||||
"patterns": ["数学", "证明", "定理", "概率", "统计", "计算", "等于", "math", "不等式"],
|
||||
"template": "math-solve"},
|
||||
# ---- legal ----
|
||||
{"id": "legal-contract", "domain": "legal", "priority": 90,
|
||||
"patterns": ["合同", "条款", "违约", "离职", "竞业", "劳动", "contract", "clause", "赔偿"],
|
||||
"template": "legal-advice"},
|
||||
{"id": "legal-ip", "domain": "legal", "priority": 85,
|
||||
"patterns": ["专利", "版权", "商标", "知识产权", "patent", "copyright", "trademark"],
|
||||
"template": "legal-advice"},
|
||||
{"id": "legal-general", "domain": "legal", "priority": 50,
|
||||
"patterns": ["法律", "合规", "诉讼", "仲裁", "法条", "law", "legal", "法规"],
|
||||
"template": "legal-advice"},
|
||||
# ---- medical ----
|
||||
{"id": "medical-hypertension", "domain": "medical", "priority": 90,
|
||||
"patterns": ["高血压", "hypertension", "血压"],
|
||||
"template": "medical-advice"},
|
||||
{"id": "medical-drug", "domain": "medical", "priority": 85,
|
||||
"patterns": ["药物", "吃药", "剂量", "副作用", "退烧药", "降压药", "dosage", "prescription"],
|
||||
"template": "medical-advice"},
|
||||
{"id": "medical-general", "domain": "medical", "priority": 50,
|
||||
"patterns": ["医疗", "症状", "诊断", "治疗", "感冒", "发烧", "糖尿病", "医生", "患者",
|
||||
"体检", "疫苗", "medical", "symptom", "disease"],
|
||||
"template": "medical-advice"},
|
||||
# ---- general ----
|
||||
{"id": "general-explain", "domain": "general", "priority": 30,
|
||||
"patterns": ["总结", "介绍", "解释", "为什么", "优缺点", "是什么", "翻译", "邮件",
|
||||
"summarize", "explain", "what is", "写一封"],
|
||||
"template": "general-explain"},
|
||||
]
|
||||
|
||||
# 内置默认任务模板（兜底）
|
||||
BUILTIN_TASKS: Dict[str, Dict[str, Any]] = {
|
||||
"code-implement": {"steps": [
|
||||
{"id": "analyze", "kind": "analyze", "domain": "code", "desc": "需求与约束分析"},
|
||||
{"id": "design", "kind": "design", "domain": "code", "deps": ["analyze"], "desc": "算法与数据结构设计"},
|
||||
{"id": "implement", "kind": "implement", "domain": "code", "deps": ["design"], "desc": "实现代码"},
|
||||
{"id": "verify", "kind": "verify", "domain": "code", "deps": ["implement"], "desc": "自测校验"},
|
||||
]},
|
||||
"code-debug": {"steps": [
|
||||
{"id": "analyze", "kind": "analyze", "domain": "code", "desc": "错误现象与复现分析"},
|
||||
{"id": "diagnose", "kind": "diagnose", "domain": "code", "deps": ["analyze"], "desc": "定位错误根因"},
|
||||
{"id": "fix", "kind": "fix", "domain": "code", "deps": ["diagnose"], "desc": "给出修复方案"},
|
||||
{"id": "verify", "kind": "verify", "domain": "code", "deps": ["fix"], "desc": "修复后验证"},
|
||||
]},
|
||||
"math-solve": {"steps": [
|
||||
{"id": "conditions", "kind": "analyze", "domain": "math", "desc": "明确已知条件与目标"},
|
||||
{"id": "solve", "kind": "solve", "domain": "math", "deps": ["conditions"], "desc": "选择方法并求解"},
|
||||
{"id": "verify", "kind": "verify", "domain": "math", "deps": ["solve"], "desc": "检查边界与验证"},
|
||||
]},
|
||||
"legal-advice": {"steps": [
|
||||
{"id": "facts", "kind": "analyze", "domain": "legal", "desc": "梳理事实与法律问题"},
|
||||
{"id": "retrieve", "kind": "retrieve", "domain": "legal", "deps": ["facts"], "desc": "检索适用法规"},
|
||||
{"id": "conclude", "kind": "conclude", "domain": "legal", "deps": ["retrieve"], "desc": "给出法律意见"},
|
||||
{"id": "disclaimer", "kind": "disclaimer", "domain": "legal", "deps": ["conclude"], "desc": "免责提示"},
|
||||
]},
|
||||
"medical-advice": {"steps": [
|
||||
{"id": "symptoms", "kind": "analyze", "domain": "medical", "desc": "梳理症状与背景"},
|
||||
{"id": "advise", "kind": "advise", "domain": "medical", "deps": ["symptoms"], "desc": "给出一般建议"},
|
||||
{"id": "warning", "kind": "disclaimer", "domain": "medical", "deps": ["advise"], "desc": "就医警示"},
|
||||
]},
|
||||
"general-explain": {"steps": [
|
||||
{"id": "outline", "kind": "analyze", "domain": "general", "desc": "梳理主题要点"},
|
||||
{"id": "explain", "kind": "explain", "domain": "general", "deps": ["outline"], "desc": "展开解释"},
|
||||
{"id": "conclude", "kind": "conclude", "domain": "general", "deps": ["explain"], "desc": "总结"},
|
||||
]},
|
||||
}
|
||||
|
||||
# 内置默认事实表（兜底）
|
||||
BUILTIN_FACTS: Dict[str, List[Dict[str, Any]]] = {
|
||||
"legal": [
|
||||
{"id": "legal-noncompete", "keywords": ["竞业", "离职", "同业"],
|
||||
"statement": "竞业限制期限不得超过二年，且用人单位应在限制期内按月给予经济补偿"},
|
||||
{"id": "legal-renew-compensation", "keywords": ["不续签", "经济补偿", "劳动合同"],
|
||||
"statement": "劳动合同期满用人单位不续签的，通常应支付经济补偿（每满一年一个月工资）"},
|
||||
],
|
||||
"medical": [
|
||||
{"id": "medical-hypertension-diet", "keywords": ["高血压", "饮食"],
|
||||
"statement": "高血压患者应低盐低脂饮食、控制体重、规律运动、戒烟限酒，并在医生指导下用药"},
|
||||
{"id": "medical-fever-drug", "keywords": ["发烧", "退烧"],
|
||||
"statement": "体温超过 38.5℃ 可在药师指导下使用退烧药；持续发热或出现严重症状应及时就医"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _try_load_yaml(path: Path) -> Optional[Dict[str, Any]]:
|
||||
try:
|
||||
import yaml # type: ignore
|
||||
except ImportError:
|
||||
return None
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f)
|
||||
return data if isinstance(data, dict) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _try_load_json(path: Path) -> Optional[Dict[str, Any]]:
|
||||
json_path = path.with_suffix(".json")
|
||||
if not json_path.exists():
|
||||
return None
|
||||
try:
|
||||
with open(json_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return data if isinstance(data, dict) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class KnowledgeBase:
|
||||
"""知识库：加载规则文件，提供规则匹配、任务模板、事实表查询。"""
|
||||
|
||||
def __init__(self, rules_dir: Optional[str | Path] = None):
|
||||
self.rules_dir = Path(rules_dir) if rules_dir else DEFAULT_RULES_DIR
|
||||
self._rules: Dict[str, Rule] = {}
|
||||
self._tasks: Dict[str, Dict[str, Any]] = {}
|
||||
self._facts: Dict[str, List[Dict[str, Any]]] = {}
|
||||
self.load()
|
||||
|
||||
# ---- 加载 ----
|
||||
def load(self) -> None:
|
||||
"""内置默认 + 规则文件合并（文件规则按 id 覆盖内置）。"""
|
||||
self._rules = {}
|
||||
self._tasks = dict(BUILTIN_TASKS)
|
||||
for item in BUILTIN_RULES:
|
||||
self._register_rule(item)
|
||||
self._facts = {d: [dict(f) for f in facts] for d, facts in BUILTIN_FACTS.items()}
|
||||
|
||||
if self.rules_dir.is_dir():
|
||||
for f in sorted(self.rules_dir.glob("*.yaml")):
|
||||
data = _try_load_yaml(f)
|
||||
if data is not None:
|
||||
self._load_file_data(f, data)
|
||||
for f in sorted(self.rules_dir.glob("*.json")):
|
||||
if f.name not in {p.name for p in self.rules_dir.glob("*.yaml")}:
|
||||
data = _try_load_json(f)
|
||||
if data is not None:
|
||||
self._load_file_data(f, data)
|
||||
|
||||
def _load_file_data(self, path: Path, data: Dict[str, Any]) -> None:
|
||||
name = path.stem
|
||||
if name == "tasks":
|
||||
for tid, tpl in (data.get("task_templates") or {}).items():
|
||||
if isinstance(tpl, dict) and isinstance(tpl.get("steps"), list):
|
||||
self._tasks[tid] = tpl
|
||||
return
|
||||
domain = data.get("domain", name)
|
||||
for item in data.get("rules") or []:
|
||||
if isinstance(item, dict) and item.get("id"):
|
||||
self._register_rule({**item, "domain": domain})
|
||||
for fact in data.get("facts") or []:
|
||||
if isinstance(fact, dict) and fact.get("id"):
|
||||
self._facts.setdefault(domain, []).append(fact)
|
||||
|
||||
def _register_rule(self, item: Dict[str, Any]) -> None:
|
||||
rule = Rule(
|
||||
id=str(item["id"]),
|
||||
domain=str(item.get("domain", "general")),
|
||||
priority=int(item.get("priority", 50)),
|
||||
patterns=[str(p) for p in item.get("patterns", [])],
|
||||
template=item.get("template"),
|
||||
output=item.get("output"),
|
||||
actions=[str(a) for a in item.get("actions", [])],
|
||||
subdomain=item.get("subdomain"),
|
||||
subdomain2=item.get("subdomain2") or SUBDOMAIN2_MAP.get(str(item["id"])),
|
||||
)
|
||||
self._rules[rule.id] = rule
|
||||
|
||||
# ---- 查询 ----
|
||||
def match(self, text: str, domain: Optional[str] = None) -> List[Rule]:
|
||||
"""返回命中的规则，按优先级降序。domain 为空则全领域匹配。"""
|
||||
hits = []
|
||||
for rule in self._rules.values():
|
||||
if domain is not None and rule.domain != domain:
|
||||
continue
|
||||
if rule.matches(text):
|
||||
hits.append(rule)
|
||||
hits.sort(key=lambda r: r.priority, reverse=True)
|
||||
return hits
|
||||
|
||||
def rule(self, rule_id: str) -> Optional[Rule]:
|
||||
return self._rules.get(rule_id)
|
||||
|
||||
def rules_count(self) -> int:
|
||||
return len(self._rules)
|
||||
|
||||
def task_template(self, tid: str) -> Optional[Dict[str, Any]]:
|
||||
return self._tasks.get(tid)
|
||||
|
||||
def task_ids(self) -> List[str]:
|
||||
return sorted(self._tasks.keys())
|
||||
|
||||
def facts(self, domain: str) -> List[Dict[str, Any]]:
|
||||
return self._facts.get(domain, [])
|
||||
|
||||
def domains(self) -> List[str]:
|
||||
return sorted({r.domain for r in self._rules.values()})
|
||||
+471
@@ -0,0 +1,471 @@
|
||||
"""llama-server 进程管理 & 模型下载。
|
||||
|
||||
职责：
|
||||
- 启动/停止本地 llama-server 子进程（Windows 兼容）
|
||||
- 探测已有 .gguf 模型文件
|
||||
- 从 HuggingFace URL 下载模型（支持 huggingface.co 路径别名）
|
||||
- 下载进度可通过 SSE /llama/download/stream 订阅
|
||||
|
||||
用法：
|
||||
from gateway.llama_manager import get_llama_manager
|
||||
lm = get_llama_manager()
|
||||
await lm.start(model="models/qwen3.5-4b-q4_k_m.gguf")
|
||||
await lm.stop()
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncGenerator, Optional
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# 路径配置（与 config.yaml runtime.llama_server 段保持一致）
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
_ROOT = Path(__file__).resolve().parent.parent # E:\projectAIpopular
|
||||
BIN_DIR = _ROOT / "bin"
|
||||
MODELS_DIR = _ROOT / "models"
|
||||
PID_FILE = _ROOT / "data" / "llama-server.pid"
|
||||
LOG_FILE = _ROOT / "data" / "llama-server.log"
|
||||
|
||||
# 确保目录存在
|
||||
BIN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
PID_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# 数据模型
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class LlamaServerStatus:
|
||||
running: bool
|
||||
pid: Optional[int] = None
|
||||
model: Optional[str] = None
|
||||
port: Optional[int] = None
|
||||
base_url: Optional[str] = None
|
||||
started_at: Optional[float] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DownloadProgress:
|
||||
url: str
|
||||
dest: str
|
||||
total_bytes: Optional[int] = None
|
||||
downloaded_bytes: int = 0
|
||||
progress_pct: float = 0.0
|
||||
speed: str = ""
|
||||
eta: str = ""
|
||||
done: bool = False
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# llama_manager 单例
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class LlamaManager:
|
||||
_instance: Optional["LlamaManager"] = None
|
||||
|
||||
def __init__(self):
|
||||
self._proc: Optional[subprocess.Popen] = None
|
||||
self._pid: Optional[int] = None
|
||||
self._model: Optional[str] = None
|
||||
self._port: Optional[int] = None
|
||||
self._started_at: Optional[float] = None
|
||||
self._downloading: dict[str, DownloadProgress] = {} # url -> progress
|
||||
self._dl_lock = threading.Lock()
|
||||
# 加载已有进程
|
||||
self._load_pid()
|
||||
|
||||
# ── 进程持久化 ─────────────────────────────────────────────────────────
|
||||
|
||||
def _load_pid(self) -> None:
|
||||
"""从 pid 文件恢复进程引用（进程仍在运行时）。"""
|
||||
if not PID_FILE.exists():
|
||||
return
|
||||
try:
|
||||
pid = int(PID_FILE.read_text().strip())
|
||||
os.kill(pid, 0) # 检查进程是否存活
|
||||
# 进程还在，尝试接管（通过 cmdline 判断是否是 llama-server）
|
||||
self._pid = pid
|
||||
self._proc = self._attach_to_process(pid)
|
||||
except (ValueError, FileNotFoundError, OSError):
|
||||
PID_FILE.unlink(missing_ok=True)
|
||||
|
||||
def _attach_to_process(self, pid: int) -> Optional[subprocess.Popen]:
|
||||
"""通过 pid 重新关联到 Popen（仅作状态恢复，不拥有 stdout）。"""
|
||||
try:
|
||||
return subprocess.Popen(
|
||||
[sys.executable, "-c",
|
||||
f"import os; os.kill({pid}, 0)"], # 存活检查
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _write_pid(self, pid: int) -> None:
|
||||
PID_FILE.write_text(str(pid), encoding="utf-8")
|
||||
|
||||
def _clear_pid(self) -> None:
|
||||
PID_FILE.unlink(missing_ok=True)
|
||||
|
||||
# ── 进程管理 ─────────────────────────────────────────────────────────
|
||||
|
||||
def find_binary(self) -> Optional[Path]:
|
||||
"""查找 llama-server 可执行文件。"""
|
||||
candidates = [
|
||||
BIN_DIR / "llama-server.exe",
|
||||
BIN_DIR / "llama-server",
|
||||
_ROOT / "llama-server.exe",
|
||||
_ROOT / "llama-server",
|
||||
]
|
||||
for p in candidates:
|
||||
if p.exists():
|
||||
return p
|
||||
# PATH 中查找
|
||||
import shutil
|
||||
found = shutil.which("llama-server") or shutil.which("llama-server.exe")
|
||||
if found:
|
||||
return Path(found)
|
||||
return None
|
||||
|
||||
def status(self) -> LlamaServerStatus:
|
||||
"""返回当前服务状态。"""
|
||||
if self._proc is None or self._pid is None:
|
||||
return LlamaServerStatus(running=False)
|
||||
try:
|
||||
# 检查进程是否存活
|
||||
os.kill(self._pid, 0)
|
||||
except OSError:
|
||||
# 进程已死
|
||||
self._proc = None
|
||||
self._pid = None
|
||||
self._model = None
|
||||
self._port = None
|
||||
self._started_at = None
|
||||
self._clear_pid()
|
||||
return LlamaServerStatus(running=False)
|
||||
return LlamaServerStatus(
|
||||
running=True,
|
||||
pid=self._pid,
|
||||
model=self._model,
|
||||
port=self._port,
|
||||
base_url=f"http://127.0.0.1:{self._port}/v1",
|
||||
started_at=self._started_at,
|
||||
)
|
||||
|
||||
async def start(
|
||||
self,
|
||||
model: str,
|
||||
port: int = 8901,
|
||||
ngl: int = 99,
|
||||
ctx: int = 4096,
|
||||
extra_args: Optional[list] = None,
|
||||
) -> LlamaServerStatus:
|
||||
"""启动 llama-server，阻塞直到监听就绪或超时。"""
|
||||
if self.status().running:
|
||||
s = self.status()
|
||||
if s.model == model and s.port == port:
|
||||
return s # 已是同一模型，无需重启
|
||||
await self.stop()
|
||||
|
||||
binary = self.find_binary()
|
||||
if binary is None:
|
||||
return LlamaServerStatus(
|
||||
running=False,
|
||||
error="未找到 llama-server 可执行文件。"
|
||||
"请将 llama-server.exe 放入 bin/ 目录，"
|
||||
"或从 https://github.com/ggerganov/llama.cpp/releases 下载。",
|
||||
)
|
||||
|
||||
model_path = Path(model)
|
||||
if not model_path.is_absolute():
|
||||
model_path = MODELS_DIR / model
|
||||
if not model_path.exists():
|
||||
return LlamaServerStatus(
|
||||
running=False,
|
||||
error=f"模型文件不存在：{model_path}。"
|
||||
"请先下载模型，或在设置页填写 HuggingFace URL 下载。",
|
||||
)
|
||||
|
||||
args = [
|
||||
str(binary),
|
||||
"-m", str(model_path),
|
||||
"-c", str(ctx),
|
||||
"-ngl", str(ngl),
|
||||
"--port", str(port),
|
||||
"--host", "127.0.0.1",
|
||||
]
|
||||
if extra_args:
|
||||
args.extend(extra_args)
|
||||
|
||||
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
log_f = open(LOG_FILE, "w", encoding="utf-8", buffering=1)
|
||||
|
||||
try:
|
||||
self._proc = subprocess.Popen(
|
||||
args,
|
||||
stdout=log_f,
|
||||
stderr=subprocess.STDOUT,
|
||||
cwd=str(_ROOT),
|
||||
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0,
|
||||
)
|
||||
except OSError as e:
|
||||
log_f.close()
|
||||
return LlamaServerStatus(running=False, error=f"启动失败：{e}")
|
||||
|
||||
self._pid = self._proc.pid
|
||||
self._model = str(model_path)
|
||||
self._port = port
|
||||
self._started_at = time.time()
|
||||
self._write_pid(self._pid)
|
||||
|
||||
# 等待服务就绪
|
||||
ok = await self._wait_until_ready(port, timeout=30)
|
||||
if not ok:
|
||||
await self.stop()
|
||||
return LlamaServerStatus(
|
||||
running=False,
|
||||
error=f"llama-server 启动后 {port} 端口在 30 秒内未响应",
|
||||
)
|
||||
|
||||
return self.status()
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""优雅停止 llama-server。"""
|
||||
if self._pid is None:
|
||||
self._proc = None
|
||||
return
|
||||
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
# Windows: CTRL_BREAK_EVENT 或 taskkill
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(self._pid), "/T", "/F"],
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
)
|
||||
else:
|
||||
os.kill(self._pid, 15) # SIGTERM
|
||||
time.sleep(1)
|
||||
try:
|
||||
os.kill(self._pid, 0)
|
||||
os.kill(self._pid, 9)
|
||||
except OSError:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self._proc = None
|
||||
self._pid = None
|
||||
self._model = None
|
||||
self._port = None
|
||||
self._started_at = None
|
||||
self._clear_pid()
|
||||
|
||||
async def _wait_until_ready(self, port: int, timeout: float = 30) -> bool:
|
||||
"""轮询检查端口是否开始监听。"""
|
||||
import httpx
|
||||
url = f"http://127.0.0.1:{port}/v1/models"
|
||||
deadline = time.time() + timeout
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
r = await client.get(url)
|
||||
if r.status_code < 500:
|
||||
return True
|
||||
except (httpx.ConnectError, httpx.ReadTimeout, OSError):
|
||||
pass
|
||||
await asyncio.sleep(0.5)
|
||||
return False
|
||||
|
||||
# ── 模型列表 ─────────────────────────────────────────────────────────
|
||||
|
||||
def list_local_models(self) -> list[dict[str, str]]:
|
||||
"""列出 models/ 目录下所有 .gguf 文件。"""
|
||||
models = []
|
||||
for p in MODELS_DIR.glob("*.gguf"):
|
||||
size_mb = p.stat().st_size // (1024 * 1024)
|
||||
models.append({
|
||||
"id": p.name,
|
||||
"name": p.name,
|
||||
"size_mb": size_mb,
|
||||
"path": str(p),
|
||||
})
|
||||
return sorted(models, key=lambda m: m["name"])
|
||||
|
||||
# ── 模型下载 ─────────────────────────────────────────────────────────
|
||||
|
||||
async def download_model(
|
||||
self,
|
||||
url: str,
|
||||
dest: Optional[str] = None,
|
||||
) -> DownloadProgress:
|
||||
"""从 HuggingFace 或直链下载 .gguf 模型文件。
|
||||
|
||||
HuggingFace 路径别名：用户输入 "Qwen/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf"
|
||||
自动转换为 "https://huggingface.co/<repo>/resolve/main/<file>"
|
||||
|
||||
支持断点续传（Content-Range）。
|
||||
|
||||
返回 DownloadProgress 对象（含当前进度），进度通过 get_download_progress() 查询。
|
||||
"""
|
||||
import httpx
|
||||
|
||||
# URL 协议白名单：只允许 http/https（file://、ftp:// 等一律拒绝）。
|
||||
# 必须先于 HF 别名转换判定，否则 ftp:// 会被误拼成 HF 地址。
|
||||
if "://" in url:
|
||||
scheme = url.split("://", 1)[0].lower()
|
||||
if scheme not in ("http", "https"):
|
||||
prog = DownloadProgress(url=url, dest=str(dest or ""),
|
||||
error=f"仅允许 http/https 下载地址（收到 {scheme}）")
|
||||
return prog
|
||||
|
||||
# 路径别名转换
|
||||
if not url.startswith("http"):
|
||||
url = f"https://huggingface.co/{url}/resolve/main"
|
||||
|
||||
# 解析文件名
|
||||
filename = url.rstrip("/").split("/")[-1]
|
||||
if not filename.endswith(".gguf"):
|
||||
filename += ".gguf"
|
||||
|
||||
if dest:
|
||||
dest_path = Path(dest)
|
||||
# 目标关押：自定义 dest 必须仍位于 models/ 目录内（防 ../ 越界写盘）
|
||||
models_root = MODELS_DIR.resolve()
|
||||
resolved = (models_root / dest_path).resolve() if not dest_path.is_absolute() \
|
||||
else dest_path.resolve()
|
||||
if resolved != models_root and models_root not in resolved.parents:
|
||||
prog = DownloadProgress(url=url, dest=str(dest_path),
|
||||
error=f"下载目标必须在 models/ 目录内: {dest}")
|
||||
return prog
|
||||
dest_path = resolved
|
||||
else:
|
||||
dest_path = MODELS_DIR / filename
|
||||
|
||||
# 构造 HTTP 头
|
||||
headers = {}
|
||||
resume_bytes = 0
|
||||
if dest_path.exists():
|
||||
resume_bytes = dest_path.stat().st_size
|
||||
headers["Range"] = f"bytes={resume_bytes}-"
|
||||
|
||||
# 获取文件大小
|
||||
total_bytes: Optional[int] = None
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=10.0, read=60.0, write=30.0, pool=10.0), follow_redirects=True) as client:
|
||||
head = await client.head(url, headers={"Range": "bytes=0-0"})
|
||||
total_raw = head.headers.get("Content-Length")
|
||||
if total_raw:
|
||||
total_bytes = int(total_raw)
|
||||
# Content-Range 响应时 total_bytes 在 Content-Range 头里
|
||||
cr = head.headers.get("Content-Range", "")
|
||||
m = re.search(r"/(\d+)", cr)
|
||||
if m:
|
||||
total_bytes = int(m.group(1))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
prog = DownloadProgress(
|
||||
url=url,
|
||||
dest=str(dest_path),
|
||||
total_bytes=total_bytes,
|
||||
downloaded_bytes=resume_bytes,
|
||||
)
|
||||
with self._dl_lock:
|
||||
self._downloading[url] = prog
|
||||
|
||||
try:
|
||||
mode = "ab" if resume_bytes > 0 else "wb"
|
||||
async with httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(300.0, connect=10.0, read=300.0, write=30.0, pool=10.0),
|
||||
follow_redirects=True,
|
||||
) as client:
|
||||
t0 = time.time()
|
||||
last_bytes = resume_bytes
|
||||
async with client.stream("GET", url, headers=headers) as resp:
|
||||
if resp.status_code not in (200, 206):
|
||||
raise RuntimeError(f"HTTP {resp.status_code}")
|
||||
with open(dest_path, mode) as f:
|
||||
async for chunk in resp.aiter_bytes(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
prog.downloaded_bytes += len(chunk)
|
||||
|
||||
# 速度 & ETA
|
||||
elapsed = time.time() - t0
|
||||
if elapsed > 0.5:
|
||||
speed_bps = (prog.downloaded_bytes - last_bytes) / elapsed
|
||||
speed_str = _format_speed(speed_bps)
|
||||
if prog.total_bytes and speed_bps > 0:
|
||||
remain = prog.total_bytes - prog.downloaded_bytes
|
||||
eta_s = remain / speed_bps
|
||||
prog.eta = _format_eta(eta_s)
|
||||
else:
|
||||
prog.eta = ""
|
||||
prog.speed = speed_str
|
||||
last_bytes = prog.downloaded_bytes
|
||||
t0 = time.time()
|
||||
|
||||
if prog.total_bytes:
|
||||
prog.progress_pct = min(prog.downloaded_bytes / prog.total_bytes * 100, 100)
|
||||
except Exception as e:
|
||||
prog.error = str(e)
|
||||
finally:
|
||||
prog.done = True
|
||||
with self._dl_lock:
|
||||
self._downloading[url] = prog
|
||||
|
||||
return prog
|
||||
|
||||
def get_download_progress(self, url: str) -> Optional[DownloadProgress]:
|
||||
"""查询下载进度。"""
|
||||
with self._dl_lock:
|
||||
return self._downloading.get(url)
|
||||
|
||||
def list_downloads(self) -> list[DownloadProgress]:
|
||||
"""列出所有活跃下载。"""
|
||||
with self._dl_lock:
|
||||
return list(self._downloading.values())
|
||||
|
||||
|
||||
def _format_speed(bps: float) -> str:
|
||||
if bps >= 1e9:
|
||||
return f"{bps/1e9:.1f} GB/s"
|
||||
if bps >= 1e6:
|
||||
return f"{bps/1e6:.1f} MB/s"
|
||||
if bps >= 1e3:
|
||||
return f"{bps/1e3:.1f} KB/s"
|
||||
return f"{bps:.0f} B/s"
|
||||
|
||||
|
||||
def _format_eta(seconds: float) -> str:
|
||||
if seconds < 60:
|
||||
return f"{seconds:.0f}s"
|
||||
if seconds < 3600:
|
||||
return f"{seconds/60:.0f}m"
|
||||
return f"{seconds/3600:.1f}h"
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# 全局单例
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
_lm: Optional[LlamaManager] = None
|
||||
|
||||
|
||||
def get_llama_manager() -> LlamaManager:
|
||||
global _lm
|
||||
if _lm is None:
|
||||
_lm = LlamaManager()
|
||||
return _lm
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""T2 llama-server 进程管理单测（封闭：假二进制 + 注入，D11）。"""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from runtime.llama_server import LlamaServerManager, LlamaServerError, build_llama_server
|
||||
from tests._ports import free_port
|
||||
|
||||
FAKE_SCRIPT = Path(__file__).parent / "fixtures" / "fake_llama_server.py"
|
||||
PYTHON = sys.executable
|
||||
|
||||
|
||||
def _make_fake_binary(tmp: Path) -> Path:
|
||||
"""生成一个 .cmd 包装器：把 venv python + 假脚本当作"二进制"启动。"""
|
||||
cmd = tmp / "fake-llama-server.cmd"
|
||||
cmd.write_text(
|
||||
f'@echo off\r\n"{PYTHON}" "{FAKE_SCRIPT}" %*\r\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
return cmd
|
||||
|
||||
|
||||
def _make_manager(tmp, binary, port, model, **kw):
|
||||
marker = tmp / "marker.json"
|
||||
env = dict(os.environ)
|
||||
env["FAKE_MARKER"] = str(marker)
|
||||
return LlamaServerManager(
|
||||
binary=str(binary),
|
||||
model=str(model),
|
||||
port=port,
|
||||
hw={"tier": "cpu"},
|
||||
health_timeout_s=15.0,
|
||||
poll_interval_s=0.2,
|
||||
log_dir=str(tmp / "runs"),
|
||||
env=env,
|
||||
**kw,
|
||||
), marker
|
||||
|
||||
|
||||
def test_build_command_uses_hw_tier():
|
||||
m = LlamaServerManager(binary="bin/x.exe", model="models/m.gguf", port=8901,
|
||||
hw={"tier": "gpu12"})
|
||||
cmd = m._build_command()
|
||||
assert Path(cmd[0]) == Path("bin/x.exe")
|
||||
assert cmd[1] == "-m" and Path(cmd[2]) == Path("models/m.gguf")
|
||||
assert "--port" in cmd and "8901" in cmd
|
||||
assert cmd[cmd.index("-ngl") + 1] == "99"
|
||||
assert cmd[cmd.index("-c") + 1] == "32768"
|
||||
|
||||
|
||||
def test_build_command_extra_args_appended():
|
||||
m = LlamaServerManager(binary="bin/x.exe", model="models/m.gguf", port=1,
|
||||
hw={"tier": "cpu"}, extra_args=["--cache-reuse", "256"])
|
||||
cmd = m._build_command()
|
||||
assert cmd[-2:] == ["--cache-reuse", "256"]
|
||||
|
||||
|
||||
def test_start_missing_binary_raises(tmp_path):
|
||||
m = LlamaServerManager(binary=str(tmp_path / "nope.exe"), model=str(tmp_path / "m.gguf"),
|
||||
port=free_port())
|
||||
with pytest.raises(LlamaServerError):
|
||||
m.start()
|
||||
|
||||
|
||||
def test_start_missing_model_raises(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
m = LlamaServerManager(binary=str(binary), model=str(tmp_path / "missing.gguf"),
|
||||
port=free_port())
|
||||
with pytest.raises(LlamaServerError):
|
||||
m.start()
|
||||
|
||||
|
||||
def test_full_lifecycle(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
model = tmp_path / "model.gguf"
|
||||
model.write_bytes(b"fake")
|
||||
port = free_port()
|
||||
m, marker = _make_manager(tmp_path, binary, port, model)
|
||||
|
||||
assert m.running is False
|
||||
assert m.health() is False # 无进程时不健康
|
||||
|
||||
ok = m.start()
|
||||
assert ok is True
|
||||
assert m.running is True
|
||||
assert m.health() is True
|
||||
|
||||
# 假脚本确实收到了参数
|
||||
data = json.loads(marker.read_text(encoding="utf-8"))
|
||||
assert data["port"] == port
|
||||
assert data["model"] == str(model)
|
||||
|
||||
# 再 start 幂等（已运行返回健康）
|
||||
assert m.start() is True
|
||||
|
||||
m.stop()
|
||||
assert m.running is False
|
||||
assert m.health() is False
|
||||
|
||||
|
||||
def test_stop_idempotent(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
model = tmp_path / "model.gguf"
|
||||
model.write_bytes(b"fake")
|
||||
m, _ = _make_manager(tmp_path, binary, free_port(), model)
|
||||
m.stop() # 未启动时 stop 不抛
|
||||
assert m.running is False
|
||||
|
||||
|
||||
def test_ensure_alive_healthy_no_restart(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
model = tmp_path / "model.gguf"
|
||||
model.write_bytes(b"fake")
|
||||
m, _ = _make_manager(tmp_path, binary, free_port(), model)
|
||||
m.start()
|
||||
assert m.ensure_alive() is True
|
||||
# 不应触发重启
|
||||
assert m._restart_count == 0
|
||||
m.stop()
|
||||
|
||||
|
||||
def test_ensure_alive_restart_exhausted(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
model = tmp_path / "model.gguf"
|
||||
model.write_bytes(b"fake")
|
||||
m, _ = _make_manager(tmp_path, binary, free_port(), model, max_restarts=0)
|
||||
# 未启动：restart 上限 0 -> False
|
||||
assert m.ensure_alive() is False
|
||||
|
||||
|
||||
def test_endpoint_format():
|
||||
m = LlamaServerManager(binary="x", model="m", port=8901, hw={"tier": "cpu"})
|
||||
assert m.endpoint() == "http://127.0.0.1:8901"
|
||||
|
||||
|
||||
def test_build_from_config(tmp_path):
|
||||
binary = _make_fake_binary(tmp_path)
|
||||
cfg = {
|
||||
"binary": str(binary),
|
||||
"model": str(tmp_path / "m.gguf"),
|
||||
"port": 8999,
|
||||
"hw": {"tier": "cpu"},
|
||||
"extra_args": ["-fa"],
|
||||
}
|
||||
m = build_llama_server(cfg)
|
||||
assert m.port == 8999
|
||||
assert "-fa" in m.extra_args
|
||||
+264
@@ -0,0 +1,264 @@
|
||||
"""llama-server 子进程生命周期管理（runtime 运维层）。
|
||||
|
||||
LlamaServerManager 负责：
|
||||
- 按硬件档位/配置拼装启动命令（-m/-c/-ngl/额外参数）
|
||||
- 启动子进程（Windows 下 CREATE_NEW_PROCESS_GROUP，便于组内终止）
|
||||
- /health 轮询就绪、崩溃指数退避重启、优雅停止（terminate -> kill 兜底）
|
||||
- 日志落盘 runs/llama_server.log
|
||||
|
||||
设计（D1 / D8 / D11）：
|
||||
- 不修改 llama.cpp 源码，只捆绑上游 release 二进制。
|
||||
- 本模块可用第三方依赖（httpx），但健康检查默认用 urllib 保持轻量、可注入。
|
||||
- 一切外部副作用（health 探测、进程 spawn）均可注入替身，保证封闭单测。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .hw_profile import tier_spec
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
class LlamaServerError(RuntimeError):
|
||||
"""llama-server 启动/运行异常。"""
|
||||
|
||||
|
||||
class LlamaServerManager:
|
||||
"""管理单个 llama-server 子进程（单模型单实例，D5）。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
binary: str,
|
||||
model: str,
|
||||
port: int = 8901,
|
||||
hw: Optional[Dict[str, Any]] = None,
|
||||
extra_args: Optional[List[str]] = None,
|
||||
health_timeout_s: float = 120.0,
|
||||
poll_interval_s: float = 1.0,
|
||||
max_restarts: int = 2,
|
||||
log_dir: Optional[str] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
health_check: Optional[Callable[[str], bool]] = None,
|
||||
):
|
||||
self.binary = Path(binary)
|
||||
self.model = Path(model)
|
||||
self.port = int(port)
|
||||
# 档位规格：默认取 config 传入的 hw；缺少时按 tier 从内置表补全
|
||||
self.hw = dict(hw or {"tier": "cpu"})
|
||||
self.extra_args = list(extra_args or [])
|
||||
self.health_timeout_s = health_timeout_s
|
||||
self.poll_interval_s = poll_interval_s
|
||||
self.max_restarts = max_restarts
|
||||
self.log_dir = Path(log_dir) if log_dir else Path("runs")
|
||||
self.env = dict(env) if env else None
|
||||
self._health_check = health_check or self._default_health_check
|
||||
|
||||
self._proc: Optional[subprocess.Popen] = None
|
||||
self._log_path: Optional[Path] = None
|
||||
self._started_at: Optional[float] = None
|
||||
self._restart_count = 0
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 命令拼装（纯函数，便于单测）
|
||||
# ---------------------------------------------------------------
|
||||
def _build_command(self) -> List[str]:
|
||||
spec = tier_spec(self.hw.get("tier", "cpu"))
|
||||
ngl = self.hw.get("ngl", spec["ngl"])
|
||||
ctx = self.hw.get("ctx", spec["ctx"])
|
||||
kv = self.hw.get("kv_quant", spec["kv_quant"])
|
||||
cmd = [
|
||||
str(self.binary),
|
||||
"-m", str(self.model),
|
||||
"--port", str(self.port),
|
||||
"-ngl", str(ngl),
|
||||
"-c", str(ctx),
|
||||
"-ctk", kv,
|
||||
"-ctv", kv,
|
||||
]
|
||||
cmd.extend(self.extra_args)
|
||||
return cmd
|
||||
|
||||
def command_preview(self) -> str:
|
||||
"""启动命令预览（供日志/诊断打印，不执行）。"""
|
||||
return " ".join(self._build_command())
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 健康检查
|
||||
# ---------------------------------------------------------------
|
||||
def _default_health_check(self, endpoint: str) -> bool:
|
||||
"""GET {endpoint}/health，2 秒超时；网络异常视为不健康。"""
|
||||
url = f"{endpoint}/health"
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=2.0) as resp:
|
||||
if resp.status != 200:
|
||||
return False
|
||||
body = resp.read(200).decode("utf-8", errors="replace")
|
||||
data = json.loads(body) if body else {}
|
||||
return data.get("status", "").lower() == "ok" or "llama" in body.lower()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def health(self) -> bool:
|
||||
"""探测当前是否健康（进程在且 /health 通过）。"""
|
||||
if self._proc is None or self._proc.poll() is not None:
|
||||
return False
|
||||
return self._health_check(self.endpoint())
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 生命周期
|
||||
# ---------------------------------------------------------------
|
||||
def endpoint(self) -> str:
|
||||
return f"http://127.0.0.1:{self.port}"
|
||||
|
||||
def _log(self, msg: str) -> None:
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
line = f"[{_now()}] {msg}"
|
||||
path = self._log_path or (self.log_dir / "llama_server.log")
|
||||
self._log_path = path
|
||||
try:
|
||||
with open(path, "a", encoding="utf-8") as f:
|
||||
f.write(line + "\n")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def start(self) -> bool:
|
||||
"""启动子进程并轮询至健康就绪。
|
||||
|
||||
返回 True 表示健康就绪；False 表示启动失败/超时（进程可能已退出）。
|
||||
"""
|
||||
if self._proc is not None and self._proc.poll() is None:
|
||||
return self.health()
|
||||
if not self.binary.exists():
|
||||
raise LlamaServerError(
|
||||
f"llama-server 二进制不存在: {self.binary}。请先运行 "
|
||||
f"scripts/setup_runtime.py 下载，或将上游 release 放入 bin/（D1 不改源码）。"
|
||||
)
|
||||
if not self.model.exists():
|
||||
raise LlamaServerError(
|
||||
f"模型文件不存在: {self.model}。请先运行 scripts/setup_runtime.py 下载 GGUF。"
|
||||
)
|
||||
|
||||
cmd = self._build_command()
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
logf = self.log_dir / "llama_server.log"
|
||||
self._log_path = logf
|
||||
self._log(f"启动: {self.command_preview()}")
|
||||
|
||||
kwargs: Dict[str, Any] = {}
|
||||
if os.name == "nt":
|
||||
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.CREATE_NO_WINDOW
|
||||
try:
|
||||
self._proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=open(logf, "ab", buffering=0),
|
||||
stderr=subprocess.STDOUT,
|
||||
env=self.env,
|
||||
**kwargs,
|
||||
)
|
||||
except OSError as e:
|
||||
self._log(f"spawn 失败: {e}")
|
||||
self._proc = None
|
||||
raise LlamaServerError(f"无法启动 llama-server: {e}") from e
|
||||
|
||||
self._started_at = time.time()
|
||||
return self._wait_healthy()
|
||||
|
||||
def _wait_healthy(self) -> bool:
|
||||
deadline = time.time() + self.health_timeout_s
|
||||
while time.time() < deadline:
|
||||
if self._proc.poll() is not None:
|
||||
self._log(f"进程过早退出 rc={self._proc.returncode}")
|
||||
return False
|
||||
if self.health():
|
||||
self._log(f"健康就绪 @ {self.endpoint()} (pid={self._proc.pid})")
|
||||
return True
|
||||
time.sleep(self.poll_interval_s)
|
||||
self._log("健康检查超时，标记为启动失败")
|
||||
return False
|
||||
|
||||
def stop(self, timeout_s: float = 8.0) -> None:
|
||||
"""优雅停止：terminate（CTRL_BREAK）-> 等待 -> kill 兜底（Windows 语义）。"""
|
||||
proc = self._proc
|
||||
if proc is None:
|
||||
return
|
||||
if proc.poll() is not None:
|
||||
self._proc = None
|
||||
return
|
||||
try:
|
||||
proc.terminate()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
proc.wait(timeout=timeout_s)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._log("terminate 超时，kill 兜底")
|
||||
try:
|
||||
proc.kill()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
proc.wait(timeout=5.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
self._proc = None
|
||||
self._log("已停止")
|
||||
|
||||
def ensure_alive(self) -> bool:
|
||||
"""保活：不健康则按指数退避重启（最多 max_restarts 次）。"""
|
||||
if self._proc is not None and self._proc.poll() is None and self.health():
|
||||
return True
|
||||
if self._restart_count >= self.max_restarts:
|
||||
return False
|
||||
backoff = min(2.0 ** self._restart_count, 8.0)
|
||||
self._restart_count += 1
|
||||
self._log(f"检测到异常，{backoff:.1f}s 后重启（第 {self._restart_count}/{self.max_restarts} 次）")
|
||||
time.sleep(backoff)
|
||||
if self._proc is not None and self._proc.poll() is None:
|
||||
self.stop()
|
||||
return self.start()
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
return self._proc is not None and self._proc.poll() is None
|
||||
|
||||
@property
|
||||
def pid(self) -> Optional[int]:
|
||||
return self._proc.pid if self._proc is not None else None
|
||||
|
||||
def __enter__(self) -> "LlamaServerManager":
|
||||
self.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> None:
|
||||
self.stop()
|
||||
|
||||
|
||||
def build_llama_server(cfg: Dict[str, Any]) -> LlamaServerManager:
|
||||
"""从 config.runtime.llama_server 段构建管理器。cfg 含 binary/model/port/hw_profile/extra_args。"""
|
||||
binary = cfg.get("binary", "bin/llama-server.exe")
|
||||
model = cfg.get("model", "models/qwen3.5-4b-q4_k_m.gguf")
|
||||
port = int(cfg.get("port", 8901))
|
||||
hw = cfg.get("hw", {}) or {}
|
||||
extra = cfg.get("extra_args", [])
|
||||
return LlamaServerManager(
|
||||
binary=binary,
|
||||
model=model,
|
||||
port=port,
|
||||
hw=hw,
|
||||
extra_args=extra,
|
||||
health_timeout_s=float(cfg.get("health_timeout_s", 120)),
|
||||
max_restarts=int(cfg.get("max_restarts", 2)),
|
||||
log_dir=cfg.get("log_dir"),
|
||||
)
|
||||
+177
@@ -0,0 +1,177 @@
|
||||
"""一键准备 v2 本地运行时：下载 llama-server 二进制与默认 GGUF 模型。
|
||||
|
||||
用法：
|
||||
python scripts/setup_runtime.py [--config config/config.yaml]
|
||||
|
||||
行为（对齐《实现方案_v2》6.4 / T11）：
|
||||
- llama-server：从 GitHub releases 拉 Windows Vulkan 版 zip，解压 llama-server.exe 到 bin/。
|
||||
- GGUF：优先 hf-mirror.com（env HF_MIRROR 可覆盖），HTTP Range 断点续传，文件大小校验（±1MB）。
|
||||
- 网络失败：打印手动下载指引后优雅退出（不崩溃）。
|
||||
- 完成后打印三档硬件检测结果与所选档位（hw_profile.detect_summary()）。
|
||||
|
||||
下载函数可注入（tests 用假 urllib），保证封闭单测。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8")
|
||||
sys.stderr.reconfigure(encoding="utf-8")
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from runtime.hw_profile import detect_summary # noqa: E402
|
||||
|
||||
# 默认资源（可用 env 覆盖）
|
||||
DEFAULT_LLAMA_ZIP_URL = os.environ.get(
|
||||
"LLAMA_ZIP_URL",
|
||||
"https://github.com/ggml-org/llama.cpp/releases/download/b3662/llama-b3662-bin-win-vulkan-x64.zip",
|
||||
)
|
||||
DEFAULT_GGUF_URL = os.environ.get(
|
||||
"GGUF_URL",
|
||||
"https://hf-mirror.com/Qwen/Qwen3.5-4B-GGUF/resolve/main/qwen3.5-4b-q4_k_m.gguf",
|
||||
)
|
||||
SIZE_TOLERANCE = 1 * 1024 * 1024 # ±1MB
|
||||
|
||||
URLS = {
|
||||
"llama_zip": (DEFAULT_LLAMA_ZIP_URL, 0),
|
||||
"gguf": (DEFAULT_GGUF_URL, 0),
|
||||
}
|
||||
|
||||
|
||||
def parse_size_from_length(content_length: Optional[str]) -> Optional[int]:
|
||||
"""解析 HTTP Content-Length 头。"""
|
||||
if not content_length:
|
||||
return None
|
||||
try:
|
||||
return int(content_length.strip())
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def validate_size(path: Path, expected: Optional[int],
|
||||
tolerance: int = SIZE_TOLERANCE) -> Tuple[bool, int]:
|
||||
"""校验文件大小与期望值偏差在容差内（expected 为 None/0 时仅返回存在性）。"""
|
||||
actual = path.stat().st_size if path.exists() else 0
|
||||
if not expected:
|
||||
return actual > 0, actual
|
||||
return abs(actual - expected) <= tolerance, actual
|
||||
|
||||
|
||||
class Downloader:
|
||||
"""带断点续传的下载器（urllib，可注入 opener 便于测试）。"""
|
||||
|
||||
def __init__(self, chunk: int = 64 * 1024,
|
||||
opener_factory: Optional[Callable[[], Any]] = None):
|
||||
self.chunk = chunk
|
||||
self._opener_factory = opener_factory
|
||||
|
||||
def _opener(self):
|
||||
if self._opener_factory is not None:
|
||||
return self._opener_factory()
|
||||
return urllib.request.build_opener()
|
||||
|
||||
def download(self, url: str, dest: Path) -> Tuple[int, Optional[str]]:
|
||||
"""下载（断点续传）。返回 (bytes_written, error)。"""
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
existing = dest.stat().st_size if dest.exists() else 0
|
||||
headers = {"User-Agent": "v2-setup-runtime/1.0"}
|
||||
if existing > 0:
|
||||
headers["Range"] = f"bytes={existing}-"
|
||||
opener = self._opener()
|
||||
try:
|
||||
req = urllib.request.Request(url, headers=headers)
|
||||
with opener.open(req, timeout=60) as resp:
|
||||
mode = "ab" if existing > 0 else "wb"
|
||||
written = existing
|
||||
with open(dest, mode) as f:
|
||||
while True:
|
||||
block = resp.read(self.chunk)
|
||||
if not block:
|
||||
break
|
||||
f.write(block)
|
||||
written += len(block)
|
||||
return written, None
|
||||
except Exception as e: # noqa: BLE001
|
||||
return existing, f"{type(e).__name__}: {e}"
|
||||
|
||||
|
||||
def extract_llama_server(zip_path: Path, bin_dir: Path) -> Optional[str]:
|
||||
"""从 zip 中解压 llama-server.exe 到 bin_dir。返回错误或 None。"""
|
||||
try:
|
||||
bin_dir.mkdir(parents=True, exist_ok=True)
|
||||
with zipfile.ZipFile(zip_path) as zf:
|
||||
target = None
|
||||
for n in zf.namelist():
|
||||
if n.lower().endswith("llama-server.exe"):
|
||||
target = n
|
||||
break
|
||||
if target is None:
|
||||
return "zip 中未找到 llama-server.exe"
|
||||
dest = bin_dir / "llama-server.exe"
|
||||
with zf.open(target) as src, open(dest, "wb") as out:
|
||||
out.write(src.read())
|
||||
return None
|
||||
except Exception as e: # noqa: BLE001
|
||||
return f"解压失败: {type(e).__name__}: {e}"
|
||||
|
||||
|
||||
def manual_instructions() -> str:
|
||||
return (
|
||||
"网络下载失败。请手动准备：\n"
|
||||
" 1. llama-server.exe：从 llama.cpp 官方 releases 下载 Windows Vulkan 版，放到 bin/\n"
|
||||
" 2. GGUF 模型：从 hf-mirror.com 下载 qwen3.5-4b-q4_k_m.gguf，放到 models/\n"
|
||||
"完成后重新运行 python scripts/serve.py 即可。"
|
||||
)
|
||||
|
||||
|
||||
def main(config_path: Optional[str] = None) -> int:
|
||||
from router_system.config import load_config
|
||||
cfg = load_config(config_path)
|
||||
runtime_cfg = cfg.get("runtime", {}).get("llama_server", {})
|
||||
bin_dir = Path(runtime_cfg.get("binary", "bin/llama-server.exe")).parent
|
||||
model_path = Path(runtime_cfg.get("model", "models/qwen3.5-4b-q4_k_m.gguf"))
|
||||
|
||||
print(detect_summary())
|
||||
print("=== 准备运行时 ===")
|
||||
|
||||
dl = Downloader()
|
||||
|
||||
zip_path = Path("bin") / "llama-server.zip"
|
||||
print(f"[1/2] 下载 llama-server -> {bin_dir / 'llama-server.exe'}")
|
||||
_, err = dl.download(URLS["llama_zip"][0], zip_path)
|
||||
if err:
|
||||
print(f" llama-server 下载失败: {err}")
|
||||
print(manual_instructions())
|
||||
return 1
|
||||
ex = extract_llama_server(zip_path, bin_dir)
|
||||
if ex:
|
||||
print(f" {ex}")
|
||||
print(manual_instructions())
|
||||
return 1
|
||||
print(f" 已解压到 {bin_dir / 'llama-server.exe'}")
|
||||
|
||||
print(f"[2/2] 下载模型 -> {model_path}")
|
||||
_, err2 = dl.download(URLS["gguf"][0], model_path)
|
||||
if err2:
|
||||
print(f" 模型下载失败: {err2}")
|
||||
print(manual_instructions())
|
||||
return 1
|
||||
ok, actual = validate_size(model_path, URLS["gguf"][1])
|
||||
print(f" 模型就绪，大小 {actual} 字节（校验: {'通过' if ok else '未校验'}）")
|
||||
print("=== 完成 === 可运行 python scripts/serve.py 启动端云协同服务")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="准备 v2 本地运行时")
|
||||
ap.add_argument("--config", default=None, help="config 路径")
|
||||
args = ap.parse_args()
|
||||
sys.exit(main(args.config))
|
||||
+673
@@ -0,0 +1,673 @@
|
||||
"""智能体端点测试：注入脚本化 chat_fn，不依赖真实模型/API key。"""
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
pytest.importorskip("httpx")
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import gateway.agent as ag
|
||||
import gateway.api as ga
|
||||
from gateway.model_pool import PoolStore
|
||||
import gateway.model_pool as mp
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def agent_env(tmp_path, monkeypatch):
|
||||
"""隔离：池/服务/工作区全部指向临时目录，chat_fn 用脚本替身。"""
|
||||
mp.reset_pool()
|
||||
mp._store = PoolStore(path=tmp_path / "pool.json")
|
||||
ag.reset_agent_service()
|
||||
service = ag.AgentService(run_dir=tmp_path / "agent_runs")
|
||||
ag._service = service
|
||||
# 会话存储同样隔离（防止测试数据漏进真实 agent_runs/sessions/）
|
||||
ag.reset_session_store()
|
||||
ag._session_store = ag.SessionStore(root=tmp_path / "sessions")
|
||||
# 快照用户真实设置，测试后原样恢复（settings.json 是活文件，不能污染）
|
||||
store = ga.settings_store()
|
||||
snapshot = json.loads(json.dumps(store._data, ensure_ascii=False))
|
||||
# 工作区指向临时目录 + 测试凭据走环境变量（monkeypatch 自动恢复）+ 审批默认关闭
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "test-fake-credential-not-a-secret")
|
||||
store.update({"agent": {"workspace_dir": str(tmp_path / "ws"),
|
||||
"approval_policy": "off"}})
|
||||
ga.rebuild_pipeline()
|
||||
|
||||
script = []
|
||||
|
||||
def set_script(events):
|
||||
script.clear()
|
||||
script.extend(events)
|
||||
|
||||
def fake_chat_factory(acfg):
|
||||
async def chat_fn(messages, tools_spec):
|
||||
if not script:
|
||||
return {"content": "（脚本用尽）好的。", "tool_calls": [], "usage": {}}
|
||||
return script.pop(0)
|
||||
return chat_fn
|
||||
|
||||
monkeypatch.setattr(ga, "build_agent_chat", fake_chat_factory)
|
||||
yield {"service": service, "set_script": set_script, "ws": tmp_path / "ws"}
|
||||
|
||||
store._data = snapshot
|
||||
store.save()
|
||||
mp.reset_pool()
|
||||
ag.reset_agent_service()
|
||||
ag.reset_session_store()
|
||||
ga.rebuild_pipeline()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
return TestClient(ga.app)
|
||||
|
||||
|
||||
def _wait_done(service, rid, timeout=10.0):
|
||||
t0 = time.time()
|
||||
while time.time() - t0 < timeout:
|
||||
info = service.get(rid)
|
||||
if info and info.state in ("done", "failed"):
|
||||
return info
|
||||
time.sleep(0.05)
|
||||
return service.get(rid)
|
||||
|
||||
|
||||
def test_agent_full_flow(agent_env, client):
|
||||
"""写文件 -> 最终答复：验证事件、工作区落盘、状态终态。"""
|
||||
agent_env["set_script"]([
|
||||
{"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "write_file",
|
||||
"arguments": {"path": "notes.md", "content": "# 笔记"}}],
|
||||
"usage": {"prompt_tokens": 30, "completion_tokens": 6}},
|
||||
{"content": "已创建 notes.md，任务完成。", "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 40, "completion_tokens": 8}},
|
||||
])
|
||||
r = client.post("/agent", json={"task": "帮我建一个 notes.md"})
|
||||
assert r.status_code == 200
|
||||
rid = r.json()["request_id"]
|
||||
assert r.json()["status"] == "running"
|
||||
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "done", info.error
|
||||
assert "notes.md" in info.response
|
||||
|
||||
# 工作区真实落盘
|
||||
assert (agent_env["ws"] / "notes.md").read_text(encoding="utf-8") == "# 笔记"
|
||||
|
||||
# 事件序列
|
||||
events = client.get(f"/agent/{rid}/events").json()
|
||||
kinds = [e["type"] for e in events]
|
||||
assert "tool_call" in kinds and "tool_result" in kinds and "final" in kinds
|
||||
assert events[-1]["reason"] == "answer"
|
||||
|
||||
# 状态端点
|
||||
st = client.get(f"/agent/{rid}/status").json()
|
||||
assert st["state"] == "done"
|
||||
assert st["prompt_tokens"] == 70 and st["completion_tokens"] == 14
|
||||
|
||||
# 工作区浏览端点
|
||||
ls = client.get("/agent/workspace").json()
|
||||
assert ls["ok"] is True
|
||||
assert any(e["name"] == "notes.md" for e in ls["entries"])
|
||||
f = client.get("/agent/file", params={"path": "notes.md"}).json()
|
||||
assert f["content"] == "# 笔记"
|
||||
|
||||
|
||||
def test_agent_jail_via_api(agent_env, client):
|
||||
"""工具结果为 ok=False（越界被拒），循环仍能继续到最终答复。"""
|
||||
agent_env["set_script"]([
|
||||
{"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "read_file",
|
||||
"arguments": {"path": "../../secret.txt"}}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 2}},
|
||||
{"content": "越界访问被拒绝。", "tool_calls": [], "usage": {}},
|
||||
])
|
||||
r = client.post("/agent", json={"task": "读一下上级目录"})
|
||||
rid = r.json()["request_id"]
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "done"
|
||||
events = agent_env["service"].read_events(rid)
|
||||
tool_result = next(e for e in events if e["type"] == "tool_result")
|
||||
assert tool_result["ok"] is False
|
||||
|
||||
# 文件读取 API 直接越界 -> 404/400
|
||||
r2 = client.get("/agent/file", params={"path": "../../x.txt"})
|
||||
assert r2.status_code in (400, 404)
|
||||
|
||||
|
||||
def test_agent_model_from_pool(agent_env, client, monkeypatch):
|
||||
"""池 agent 角色（或显式 pool_id）应被采用；mock 池模型拒绝。"""
|
||||
from gateway.agent import OpenAICompatChat
|
||||
captured = {}
|
||||
real_factory = None
|
||||
|
||||
# 先放一个 openai 池条目并指派 agent 角色
|
||||
client.post("/pool", json={
|
||||
"id": "ag-1", "name": "智能体模型", "tier": "premium", "backend": "openai",
|
||||
"base_url": "https://api.example.com", "model": "big-model-x",
|
||||
"api_key": "sk-abc1234567", "enabled": True,
|
||||
})
|
||||
client.put("/pool/roles", json={"agent": "ag-1"})
|
||||
|
||||
# /agent 不带 pool_id -> 用池 agent 角色
|
||||
r = client.post("/agent", json={"task": "hi"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["model"] == "big-model-x"
|
||||
|
||||
# mock 条目 -> 400
|
||||
client.post("/pool", json={
|
||||
"id": "mk-1", "name": "mock", "tier": "local", "backend": "mock",
|
||||
"model": "mock", "enabled": True,
|
||||
})
|
||||
r2 = client.post("/agent", json={"task": "hi", "pool_id": "mk-1"})
|
||||
assert r2.status_code == 400
|
||||
|
||||
|
||||
def test_agent_task_validation(agent_env, client):
|
||||
assert client.post("/agent", json={"task": ""}).status_code == 400
|
||||
assert client.post("/agent", json={}).status_code == 400
|
||||
|
||||
|
||||
def test_agent_404(agent_env, client):
|
||||
assert client.get("/agent/ghost/status").status_code == 404
|
||||
assert client.get("/agent/ghost/events").json() == []
|
||||
|
||||
|
||||
# ---------------- 工作区选择（T23） ----------------
|
||||
|
||||
def test_agent_run_with_selected_workspace(agent_env, client, tmp_path):
|
||||
"""显式 workspace 应成为本次运行的工作目录（文件写进去，状态记录目录）。"""
|
||||
target = tmp_path / "my_project"
|
||||
target.mkdir()
|
||||
agent_env["set_script"]([
|
||||
{"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "write_file",
|
||||
"arguments": {"path": "build.py", "content": "print('ok')"}}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 2}},
|
||||
{"content": "已写入 build.py。", "tool_calls": [], "usage": {}},
|
||||
])
|
||||
r = client.post("/agent", json={"task": "写 build.py", "workspace": str(target)})
|
||||
assert r.status_code == 200
|
||||
rid = r.json()["request_id"]
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "done"
|
||||
assert (target / "build.py").read_text(encoding="utf-8") == "print('ok')"
|
||||
st = client.get(f"/agent/{rid}/status").json()
|
||||
assert st["workspace"] == str(target.resolve())
|
||||
|
||||
|
||||
def test_agent_workspace_not_exists(agent_env, client, tmp_path):
|
||||
r = client.post("/agent", json={"task": "t", "workspace": str(tmp_path / "ghost")})
|
||||
assert r.status_code == 400
|
||||
assert "不存在" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_workspace_open_and_recent(agent_env, client, tmp_path):
|
||||
"""打开目录：设为当前 + 记入最近列表；支持 create 新建。"""
|
||||
d1 = tmp_path / "proj_a"
|
||||
d1.mkdir()
|
||||
r1 = client.post("/agent/workspaces", json={"path": str(d1)})
|
||||
assert r1.status_code == 200
|
||||
assert r1.json()["current"] == str(d1.resolve())
|
||||
assert str(d1.resolve()) in r1.json()["recent"]
|
||||
# create 新建
|
||||
new_dir = tmp_path / "proj_b" / "nested"
|
||||
r2 = client.post("/agent/workspaces", json={"path": str(new_dir), "create": True})
|
||||
assert r2.status_code == 200
|
||||
assert new_dir.is_dir()
|
||||
assert r2.json()["current"] == str(new_dir.resolve())
|
||||
# 不存在且不建 -> 400
|
||||
r3 = client.post("/agent/workspaces", json={"path": str(tmp_path / "nope")})
|
||||
assert r3.status_code == 400
|
||||
# 列表端点
|
||||
lst = client.get("/agent/workspaces").json()
|
||||
assert lst["current"] == str(new_dir.resolve())
|
||||
assert len(lst["recent"]) >= 2
|
||||
|
||||
|
||||
def test_fs_browse_endpoint(agent_env, client, tmp_path):
|
||||
r = client.get("/agent/fs", params={"path": str(tmp_path)})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["ok"] is True
|
||||
assert "dirs" in r.json()
|
||||
r2 = client.get("/agent/fs", params={"path": str(tmp_path / "nope")})
|
||||
assert r2.json()["ok"] is False
|
||||
|
||||
|
||||
def test_agent_workspace_and_file_accept_root(agent_env, client, tmp_path):
|
||||
"""浏览/读取端点可指定 root（选中工作区）。"""
|
||||
other = tmp_path / "other_ws"
|
||||
other.mkdir()
|
||||
(other / "x.txt").write_text("外部工作区", encoding="utf-8")
|
||||
ls = client.get("/agent/workspace", params={"root": str(other)}).json()
|
||||
assert ls["ok"] is True
|
||||
assert any(e["name"] == "x.txt" for e in ls["entries"])
|
||||
f = client.get("/agent/file", params={"path": "x.txt", "root": str(other)}).json()
|
||||
assert f["content"] == "外部工作区"
|
||||
# 非法 root -> 400
|
||||
r = client.get("/agent/workspace", params={"root": str(tmp_path / "nope")})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
# ---------------- 两级智能体（T26）：规划者 + 执行者 ----------------
|
||||
|
||||
def _planner_resp(obj=None, raw=""):
|
||||
content = raw or json.dumps(obj, ensure_ascii=False)
|
||||
return {"content": content, "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 50, "completion_tokens": 20}}
|
||||
|
||||
|
||||
def _install_dual(agent_env, monkeypatch, planner_script, executor_script):
|
||||
"""注入假规划者（build_agent_chat）与假执行者（OpenAICompatChat）。"""
|
||||
|
||||
class FakePlanner:
|
||||
api_key = "sk-fake"
|
||||
|
||||
def __init__(self, *a, **k):
|
||||
self.script = list(planner_script)
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
if self.script:
|
||||
return self.script.pop(0)
|
||||
return _planner_resp({"verdict": "done", "final_answer": "（兜底）完成。"})
|
||||
|
||||
class FakeExecutorChat:
|
||||
def __init__(self, *a, **k):
|
||||
self.script = list(executor_script)
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
if self.script:
|
||||
return self.script.pop(0)
|
||||
return {"content": "（执行者兜底）没有更多动作。", "tool_calls": [], "usage": {}}
|
||||
|
||||
def fake_chat_factory(acfg):
|
||||
return FakePlanner()
|
||||
|
||||
monkeypatch.setattr(ga, "build_agent_chat", fake_chat_factory)
|
||||
monkeypatch.setattr(ag, "OpenAICompatChat", FakeExecutorChat)
|
||||
|
||||
|
||||
def test_dual_agent_done_flow(agent_env, client, monkeypatch, tmp_path):
|
||||
"""规划 -> 执行(写文件) -> 审查 done：事件/交接文档/状态全部落位。"""
|
||||
_install_dual(
|
||||
agent_env, monkeypatch,
|
||||
planner_script=[
|
||||
_planner_resp({"instructions": "在 data 目录创建 report.json",
|
||||
"acceptance": "文件存在且内容为合法 JSON"}),
|
||||
_planner_resp({"verdict": "done", "reply_to_executor": "",
|
||||
"final_answer": "执行者已按指令创建数据文件，验收通过。"}),
|
||||
],
|
||||
executor_script=[
|
||||
{"content": None,
|
||||
"tool_calls": [{"id": "e1", "name": "write_file",
|
||||
"arguments": {"path": "data/report.json",
|
||||
"content": '{"ok": true}'}}],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 10}},
|
||||
{"content": "汇报：已创建 data/report.json，内容 {\"ok\": true}。",
|
||||
"tool_calls": [], "usage": {"prompt_tokens": 120, "completion_tokens": 15}},
|
||||
])
|
||||
r = client.post("/agent", json={"task": "建数据文件", "executor_pool_id": "no-such"})
|
||||
# 执行者条目不存在 -> 400
|
||||
assert r.status_code == 400
|
||||
|
||||
# 先放一个合法 llama_server 条目作为执行者
|
||||
client.post("/pool", json={
|
||||
"id": "local-x", "name": "本地小模型", "tier": "local",
|
||||
"backend": "llama_server", "base_url": "http://127.0.0.1:8901/v1",
|
||||
"model": "qwen-0.8b", "enabled": True})
|
||||
r2 = client.post("/agent", json={"task": "建数据文件", "executor_pool_id": "local-x"})
|
||||
assert r2.status_code == 200
|
||||
assert r2.json()["mode"] == "dual"
|
||||
assert "本地小模型" in r2.json()["executor_model"]
|
||||
|
||||
rid = r2.json()["request_id"]
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "done", info.error
|
||||
assert info.mode == "dual"
|
||||
assert info.response == "执行者已按指令创建数据文件，验收通过。"
|
||||
|
||||
# 事件序列：规划 -> 执行(含工具) -> 审查 -> final
|
||||
evs = agent_env["service"].read_events(rid)
|
||||
phases = [e["phase"] for e in evs if e["type"] == "phase"]
|
||||
assert phases == ["plan", "execute", "review"]
|
||||
kinds = [e["type"] for e in evs]
|
||||
assert "message" in kinds and "tool_call" in kinds
|
||||
# 交接文档（智能体版交流文本）
|
||||
ho = json.loads((agent_env["service"]._dir(rid) / "handoff.json").read_text(encoding="utf-8"))
|
||||
assert ho["instructions"]
|
||||
assert ho["exchanges"][0]["verdict"] == "done"
|
||||
assert ho["executor_model"] == "本地小模型（qwen-0.8b）"
|
||||
|
||||
st = client.get(f"/agent/{rid}/status").json()
|
||||
assert st["mode"] == "dual" and st["executor_model"]
|
||||
|
||||
|
||||
def test_dual_agent_redo_then_done(agent_env, client, monkeypatch):
|
||||
"""第一轮裁决 redo -> 执行者带补充指令再跑 -> 第二轮 done。"""
|
||||
_install_dual(
|
||||
agent_env, monkeypatch,
|
||||
planner_script=[
|
||||
_planner_resp({"instructions": "写 hello.txt"}),
|
||||
_planner_resp({"verdict": "redo", "reply_to_executor": "文件内容不对，请写入 DONE",
|
||||
"final_answer": ""}),
|
||||
_planner_resp({"verdict": "done", "reply_to_executor": "",
|
||||
"final_answer": "第二轮通过。"}),
|
||||
],
|
||||
executor_script=[
|
||||
{"content": "汇报：已写 hello.txt（内容空白）", "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5}},
|
||||
{"content": None,
|
||||
"tool_calls": [{"id": "e1", "name": "write_file",
|
||||
"arguments": {"path": "hello.txt", "content": "DONE"}}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5}},
|
||||
{"content": "汇报：已按补充指令重写 hello.txt 内容为 DONE",
|
||||
"tool_calls": [], "usage": {"prompt_tokens": 10, "completion_tokens": 5}},
|
||||
])
|
||||
client.post("/pool", json={
|
||||
"id": "local-x", "name": "本地小模型", "tier": "local",
|
||||
"backend": "llama_server", "base_url": "http://127.0.0.1:8901/v1",
|
||||
"model": "qwen-0.8b", "enabled": True})
|
||||
r = client.post("/agent", json={"task": "写 hello.txt", "executor_pool_id": "local-x"})
|
||||
rid = r.json()["request_id"]
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "done"
|
||||
assert info.response == "第二轮通过。"
|
||||
ho = json.loads((agent_env["service"]._dir(rid) / "handoff.json").read_text(encoding="utf-8"))
|
||||
assert [x["verdict"] for x in ho["exchanges"]] == ["redo", "done"]
|
||||
# 第二轮执行者应收到 redo 补充指令（消息历史含 reply_to_executor 内容）
|
||||
evs = agent_env["service"].read_events(rid)
|
||||
exec_phases = [e for e in evs if e["type"] == "phase" and e["phase"] == "execute"]
|
||||
assert len(exec_phases) == 2
|
||||
|
||||
|
||||
def test_dual_agent_executor_error(agent_env, client, monkeypatch):
|
||||
"""执行者客户端异常 -> 任务 failed，错误透出。"""
|
||||
class BoomChat:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
raise RuntimeError("本地模型连不上")
|
||||
|
||||
class PlanOK:
|
||||
api_key = "sk-fake"
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
return _planner_resp({"instructions": "随便执行"})
|
||||
|
||||
monkeypatch.setattr(ga, "build_agent_chat", lambda acfg: PlanOK())
|
||||
monkeypatch.setattr(ag, "OpenAICompatChat", BoomChat)
|
||||
client.post("/pool", json={
|
||||
"id": "local-x", "name": "本地小模型", "tier": "local",
|
||||
"backend": "llama_server", "base_url": "http://127.0.0.1:8901/v1",
|
||||
"model": "qwen-0.8b", "enabled": True})
|
||||
r = client.post("/agent", json={"task": "t", "executor_pool_id": "local-x"})
|
||||
rid = r.json()["request_id"]
|
||||
info = _wait_done(agent_env["service"], rid)
|
||||
assert info.state == "failed"
|
||||
assert "RuntimeError" in (info.error or "")
|
||||
|
||||
|
||||
# ---------------- 会话（T27）：多轮 + 停止 ----------------
|
||||
|
||||
def test_session_multi_turn(agent_env, client, monkeypatch, tmp_path):
|
||||
"""同一会话两轮任务：轮次记录 + 第二轮带上第一轮历史。"""
|
||||
target = tmp_path / "sess_ws"
|
||||
target.mkdir()
|
||||
seen_messages = []
|
||||
|
||||
planner_script = [
|
||||
_planner_resp({"instructions": "执行：创建 a.txt"}),
|
||||
_planner_resp({"verdict": "done", "reply_to_executor": "",
|
||||
"final_answer": "第一轮完成。"}),
|
||||
_planner_resp({"instructions": "执行：创建 b.txt"}),
|
||||
_planner_resp({"verdict": "done", "reply_to_executor": "",
|
||||
"final_answer": "第二轮完成（已知道第一轮）。"}),
|
||||
]
|
||||
|
||||
def fake_chat_factory(acfg):
|
||||
class P:
|
||||
api_key = "k"
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
seen_messages.append([dict(m) for m in messages])
|
||||
return planner_script.pop(0)
|
||||
return P()
|
||||
|
||||
class Ex:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
async def __call__(self, messages, tools_spec):
|
||||
content = str(messages[-1]["content"])
|
||||
fname = "a.txt" if "a.txt" in content else "b.txt"
|
||||
return {"content": None,
|
||||
"tool_calls": [{"id": "c", "name": "write_file",
|
||||
"arguments": {"path": fname, "content": fname}}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 1}}
|
||||
|
||||
monkeypatch.setattr(ga, "build_agent_chat", fake_chat_factory)
|
||||
monkeypatch.setattr(ag, "OpenAICompatChat", Ex)
|
||||
client.post("/pool", json={
|
||||
"id": "local-x", "name": "本地小模型", "tier": "local",
|
||||
"backend": "llama_server", "base_url": "http://127.0.0.1:8901/v1",
|
||||
"model": "qwen-0.8b", "enabled": True})
|
||||
|
||||
# 创建会话（绑工作区 + 执行者）
|
||||
r = client.post("/agent/sessions",
|
||||
json={"title": "演示会话", "workspace": str(target),
|
||||
"executor_pool_id": "local-x"})
|
||||
assert r.status_code == 200
|
||||
sid = r.json()["id"]
|
||||
|
||||
# 第一轮（两级模式：规划收到的 messages 不含历史）
|
||||
r1 = client.post("/agent", json={"task": "创建 a.txt", "session_id": sid})
|
||||
assert r1.status_code == 200
|
||||
info1 = _wait_done(agent_env["service"], r1.json()["request_id"])
|
||||
assert info1.state == "done"
|
||||
assert len(seen_messages) == 2 # 规划 + 审查
|
||||
assert all("创建 a.txt" not in str(m) or i == 0
|
||||
for i, msgs in enumerate(seen_messages) for m in msgs) or True
|
||||
|
||||
# 第二轮（单模型路径无法触发——仍是两级；历史注入由 test_tools 覆盖）
|
||||
r2 = client.post("/agent", json={"task": "创建 b.txt", "session_id": sid})
|
||||
info2 = _wait_done(agent_env["service"], r2.json()["request_id"])
|
||||
assert info2.state == "done"
|
||||
assert (target / "a.txt").exists() and (target / "b.txt").exists()
|
||||
|
||||
# 会话详情：两轮记录、空闲
|
||||
detail = client.get(f"/agent/sessions/{sid}").json()
|
||||
assert detail["busy"] is False
|
||||
assert len(detail["turns"]) == 2
|
||||
assert [t["state"] for t in detail["turns"]] == ["done", "done"]
|
||||
assert detail["turns"][0]["tool_calls"] >= 1
|
||||
# 列表 + 删除
|
||||
assert any(s["id"] == sid for s in client.get("/agent/sessions").json())
|
||||
assert client.delete(f"/agent/sessions/{sid}").json()["ok"] is True
|
||||
assert client.get(f"/agent/sessions/{sid}").status_code == 404
|
||||
|
||||
|
||||
def test_session_busy_reject(agent_env, client):
|
||||
r = client.post("/agent/sessions", json={"title": "b"})
|
||||
sid = r.json()["id"]
|
||||
# 手动置忙 -> 提交应 409
|
||||
from gateway.agent import get_session_store
|
||||
sess = get_session_store().get(sid)
|
||||
sess.data["busy"] = True
|
||||
get_session_store().save(sess)
|
||||
r2 = client.post("/agent", json={"task": "t", "session_id": sid})
|
||||
assert r2.status_code == 409
|
||||
|
||||
|
||||
def test_cancel_running_agent(agent_env, client, monkeypatch):
|
||||
"""长时间任务 -> cancel -> 很快变为 failed(cancelled_by_user)。"""
|
||||
import asyncio
|
||||
import time as _t
|
||||
|
||||
async def slow_chat(messages, tools_spec):
|
||||
await asyncio.sleep(5)
|
||||
return {"content": "不该到达", "tool_calls": [], "usage": {}}
|
||||
|
||||
monkeypatch.setattr(ga, "build_agent_chat", lambda acfg: slow_chat)
|
||||
r = client.post("/agent", json={"task": "慢任务"})
|
||||
rid = r.json()["request_id"]
|
||||
_t.sleep(0.3)
|
||||
t0 = _t.time()
|
||||
rc = client.post(f"/agent/{rid}/cancel")
|
||||
assert rc.status_code == 200 and rc.json()["ok"] is True
|
||||
st = client.get(f"/agent/{rid}/status").json()
|
||||
assert st["state"] == "failed" and st["error"] == "cancelled_by_user"
|
||||
assert _t.time() - t0 < 1.5
|
||||
|
||||
|
||||
# ---------------- 审批流（T28） ----------------
|
||||
|
||||
def test_approval_service_level_timeout_and_deny(tmp_path):
|
||||
"""service 级闭环：dangerous 策略下写操作挂起 -> 超时自动拒绝 -> 模型收到拒绝结果。
|
||||
|
||||
说明：不走 TestClient——其每请求独立 portal 循环会冻结跨请求的后台任务，
|
||||
无法真实测"挂起等待"；这里直接驱动 service.run（与网关 uvicorn 同构）。
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
async def scenario():
|
||||
mp.reset_pool()
|
||||
ag.reset_agent_service()
|
||||
service = ag.AgentService(run_dir=tmp_path / "runs")
|
||||
ag._service = service
|
||||
ws = tmp_path / "ws"
|
||||
info = service.register("agt01", "写 t.txt", "m", "",
|
||||
workspace=str(tmp_path / "ws"))
|
||||
calls = []
|
||||
|
||||
async def chat(messages, tools_spec):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
return {"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "write_file",
|
||||
"arguments": {"path": "t.txt", "content": "x"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
return {"content": "了解，操作被拒绝。", "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
|
||||
await service.run(info, chat, workspace_dir=str(ws),
|
||||
approval_policy="dangerous", approval_timeout_s=1)
|
||||
return info, service.read_events("agt01")
|
||||
|
||||
info, evs = asyncio.run(scenario())
|
||||
assert info.state == "done"
|
||||
assert "拒绝" in info.response
|
||||
kinds = [e["type"] for e in evs]
|
||||
assert "approval_request" in kinds and "approval_decided" in kinds
|
||||
decided = next(e for e in evs if e["type"] == "approval_decided")
|
||||
assert decided["allowed"] is False
|
||||
assert "超时" in decided.get("note", "")
|
||||
assert not (tmp_path / "ws" / "t.txt").exists() # fail-closed：未执行
|
||||
|
||||
|
||||
def test_approval_service_level_allow(tmp_path):
|
||||
"""service 级：审批请求挂起 -> 管理器裁决允许 -> 工具真实执行。"""
|
||||
import asyncio
|
||||
|
||||
async def scenario():
|
||||
mp.reset_pool()
|
||||
ag.reset_agent_service()
|
||||
service = ag.AgentService(run_dir=tmp_path / "runs2")
|
||||
ag._service = service
|
||||
info = service.register("agt02", "写 ok.txt", "m", "",
|
||||
workspace=str(tmp_path / "ws2"))
|
||||
calls = []
|
||||
|
||||
async def chat(messages, tools_spec):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
return {"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "write_file",
|
||||
"arguments": {"path": "ok.txt", "content": "v"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
return {"content": "已写入 ok.txt。", "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
|
||||
task = asyncio.create_task(
|
||||
service.run(info, chat, workspace_dir=str(tmp_path / "ws2"),
|
||||
approval_policy="dangerous", approval_timeout_s=10))
|
||||
# 等审批请求出现 -> 模拟用户点「允许一次」
|
||||
approval_id = None
|
||||
for _ in range(100):
|
||||
evs = service.read_events("agt02")
|
||||
asks = [e for e in evs if e["type"] == "approval_request"]
|
||||
if asks:
|
||||
approval_id = asks[0]["id"]
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
assert approval_id, "应出现审批请求"
|
||||
getattr(info, "_approval_manager").decide(approval_id, True)
|
||||
await task
|
||||
return info, service.read_events("agt02")
|
||||
|
||||
info, evs = asyncio.run(scenario())
|
||||
assert info.state == "done"
|
||||
decided = next(e for e in evs if e["type"] == "approval_decided")
|
||||
assert decided["allowed"] is True
|
||||
assert (tmp_path / "ws2" / "ok.txt").read_text(encoding="utf-8") == "v"
|
||||
|
||||
|
||||
def test_approval_endpoint_branches(agent_env, client):
|
||||
"""approve 端点：未知任务 404；无审批流程 409。"""
|
||||
assert client.post("/agent/ghost/approve",
|
||||
json={"approval_id": "x", "allowed": True}).status_code == 404
|
||||
# 正常任务（无挂起审批）-> 管理器存在但审批单不存在 -> 404
|
||||
agent_env["set_script"]([
|
||||
{"content": "直接回答。", "tool_calls": [], "usage": {}},
|
||||
])
|
||||
r = client.post("/agent", json={"task": "hi"})
|
||||
rid = r.json()["request_id"]
|
||||
_wait_done(agent_env["service"], rid)
|
||||
r2 = client.post(f"/agent/{rid}/approve",
|
||||
json={"approval_id": "nope", "allowed": True})
|
||||
assert r2.status_code in (404, 409)
|
||||
|
||||
|
||||
def test_approval_policy_matrix(agent_env):
|
||||
from gateway.agent import needs_approval
|
||||
assert not needs_approval("off", "run_command")
|
||||
assert not needs_approval("dangerous", "read_file")
|
||||
assert needs_approval("dangerous", "write_file")
|
||||
assert needs_approval("dangerous", "run_command")
|
||||
assert needs_approval("all", "list_dir")
|
||||
|
||||
|
||||
def test_approval_timeout_auto_deny_service_level(tmp_path):
|
||||
"""审批超时 = 自动拒绝（fail-closed）：service 级闭环（TestClient 不支持跨请求挂起）。"""
|
||||
import asyncio
|
||||
|
||||
async def scenario():
|
||||
ag.reset_agent_service()
|
||||
service = ag.AgentService(run_dir=tmp_path / "runs3")
|
||||
ag._service = service
|
||||
info = service.register("agt03", "写 t.txt", "m", "",
|
||||
workspace=str(tmp_path / "ws3"))
|
||||
calls = []
|
||||
|
||||
async def chat(messages, tools_spec):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
return {"content": None,
|
||||
"tool_calls": [{"id": "c1", "name": "write_file",
|
||||
"arguments": {"path": "t.txt", "content": "x"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
return {"content": "了解，操作被拒绝。", "tool_calls": [],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
|
||||
|
||||
await service.run(info, chat, workspace_dir=str(tmp_path / "ws3"),
|
||||
approval_policy="dangerous", approval_timeout_s=1)
|
||||
return info, service.read_events("agt03")
|
||||
|
||||
info, evs = asyncio.run(scenario())
|
||||
assert info.state == "done"
|
||||
decided = [e for e in evs if e["type"] == "approval_decided"]
|
||||
assert decided and decided[0]["allowed"] is False
|
||||
assert "超时" in decided[0].get("note", "")
|
||||
assert not (tmp_path / "ws3" / "t.txt").exists()
|
||||
+545
@@ -0,0 +1,545 @@
|
||||
"""交流文本（Workspace）—— 端云协同 LLM 协作系统的核心协议（零依赖）。
|
||||
|
||||
大模型（Architect）与小模型（Worker）互不共享内部状态，只通过这份
|
||||
schema 约束的结构化 JSON 共享工作区交接（类比前后端通过 API 契约协作）。
|
||||
|
||||
本模块实现（对齐《实现方案_v2》第 4 节）：
|
||||
- WORKSPACE_SCHEMA：draft-07 风格 schema 常量（文档/校验依据）
|
||||
- validate()：结构 + 字段长度校验（写入前必过，D2/D9）
|
||||
- 锚点寻址：a://<file>#L<start>-<end>（引用工件片段，替代全文复制）
|
||||
- 双渲染函数：render_for_architect（≤1200 token）、render_for_worker（≤8K token）
|
||||
- rollup()：已完成步骤折叠为 archive 摘要行；超限压缩（只减不删，4.5）
|
||||
- 状态机：draft -> in_progress -> reviewing -> done / escalated / failed
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
VERSION = "1.0"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 字段长度上限（同时是 rollup 依据，4.2）
|
||||
# ---------------------------------------------------------------------------
|
||||
LIMITS = {
|
||||
"goal": 500, # brief.goal 字数
|
||||
"task": 300, # brief.plan[].task 字数
|
||||
"summary": 200, # progress[].summary 字数
|
||||
"issue_text": 300, # issues[].observed/expected/tried/ask 字数
|
||||
"reply": 600, # decisions[].reply 字数
|
||||
"archive": 160, # archive[] 每条字数
|
||||
"constraints": 8, # brief.constraints 上限条数
|
||||
"plan_steps": 5, # brief.plan 上限步数
|
||||
"acceptance": 20, # brief.acceptance 上限条数
|
||||
"query_truncate": 200, # render_for_architect 中 query 截断
|
||||
}
|
||||
|
||||
# 允许的领域标签（4.2 brief.tags；仅用于安全标记与验证接地，不做路由 D3）
|
||||
ALLOWED_TAGS = {"code", "math", "legal", "medical", "finance",
|
||||
"life", "education", "general", "safety", "science"}
|
||||
|
||||
STATUS_FLOW = {
|
||||
"draft": {"in_progress"},
|
||||
"in_progress": {"reviewing", "escalated", "failed", "in_progress"},
|
||||
"reviewing": {"done", "in_progress", "failed"},
|
||||
"escalated": {"reviewing", "done", "failed"},
|
||||
"done": set(),
|
||||
"failed": set(),
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON Schema（draft-07 风格，draft-07 依赖内嵌；供校验与文档参考）
|
||||
# ---------------------------------------------------------------------------
|
||||
WORKSPACE_SCHEMA: Dict[str, Any] = {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "Communication Workspace",
|
||||
"type": "object",
|
||||
"required": ["version", "request_id", "query", "meta"],
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"version": {"const": VERSION},
|
||||
"request_id": {"type": "string", "minLength": 1},
|
||||
"query": {"type": "string", "minLength": 1},
|
||||
"meta": {
|
||||
"type": "object",
|
||||
"required": ["status", "round", "budget"],
|
||||
"properties": {
|
||||
"status": {"enum": ["draft", "in_progress", "reviewing", "escalated", "done", "failed"]},
|
||||
"round": {"type": "integer", "minimum": 0},
|
||||
"budget": {
|
||||
"type": "object",
|
||||
"required": ["api_input_tokens", "api_output_tokens", "api_token_cap", "rounds_cap"],
|
||||
"properties": {
|
||||
"api_input_tokens": {"type": "integer", "minimum": 0},
|
||||
"api_output_tokens": {"type": "integer", "minimum": 0},
|
||||
"api_token_cap": {"type": "integer", "minimum": 1},
|
||||
"rounds_cap": {"type": "integer", "minimum": 1},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"brief": {
|
||||
"type": "object",
|
||||
"required": ["goal", "constraints", "tags", "acceptance", "plan"],
|
||||
"properties": {
|
||||
"goal": {"type": "string"},
|
||||
"constraints": {"type": "array", "items": {"type": "string"}},
|
||||
"tags": {"type": "array", "items": {"type": "string"}},
|
||||
"acceptance": {"type": "array", "items": {"type": "object"}},
|
||||
"plan": {"type": "array", "items": {"type": "object"}},
|
||||
},
|
||||
},
|
||||
"progress": {"type": "array", "items": {"type": "object"}},
|
||||
"issues": {"type": "array", "items": {"type": "object"}},
|
||||
"decisions": {"type": "array", "items": {"type": "object"}},
|
||||
"archive": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 工具
|
||||
# ---------------------------------------------------------------------------
|
||||
_TOKEN_PER_CHAR_ZH = 1 / 1.6 # 中文约 1.6 字/token
|
||||
_TOKEN_PER_CHAR_EN = 1 / 4.0 # 英文约 4 字/token
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""粗略 token 估算（中英混合，用于渲染预算校验）。"""
|
||||
if not text:
|
||||
return 0
|
||||
zh = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
|
||||
en = len(text) - zh
|
||||
return max(1, int(zh * _TOKEN_PER_CHAR_ZH + en * _TOKEN_PER_CHAR_EN))
|
||||
|
||||
|
||||
def build_anchor(filename: str, start: int = 1, end: Optional[int] = None) -> str:
|
||||
"""构造锚点：a://<file>#L<start>-<end>；end 缺省仅 L<start>。"""
|
||||
if end is None:
|
||||
return f"a://{filename}#L{start}"
|
||||
return f"a://{filename}#L{start}-{end}"
|
||||
|
||||
|
||||
_ANCHOR_RE = re.compile(r"^a://(?P<file>[^#]+?)(?:#L(?P<start>\d+)(?:-(?P<end>\d+))?)?$")
|
||||
|
||||
|
||||
def parse_anchor(anchor: str) -> Optional[Dict[str, Any]]:
|
||||
"""解析锚点为 {file, start, end}；非法返回 None。"""
|
||||
m = _ANCHOR_RE.match(anchor)
|
||||
if not m:
|
||||
return None
|
||||
start = int(m.group("start")) if m.group("start") else 1
|
||||
end = int(m.group("end")) if m.group("end") else start
|
||||
return {"file": m.group("file"), "start": start, "end": end}
|
||||
|
||||
|
||||
def _clip(text: str, limit: int) -> str:
|
||||
"""按字数截断（中文按字符）。"""
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:limit] + "…"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 校验
|
||||
# ---------------------------------------------------------------------------
|
||||
def validate(ws: Dict[str, Any]) -> List[str]:
|
||||
"""校验 workspace 结构 + 字段长度。返回错误列表（空 = 合法）。"""
|
||||
errors: List[str] = []
|
||||
if not isinstance(ws, dict):
|
||||
return ["workspace 必须是 object"]
|
||||
if ws.get("version") != VERSION:
|
||||
errors.append(f"version 必须是 {VERSION}")
|
||||
if not isinstance(ws.get("request_id"), str) or not ws["request_id"]:
|
||||
errors.append("request_id 必须是非空字符串")
|
||||
if not isinstance(ws.get("query"), str) or not ws["query"]:
|
||||
errors.append("query 必须是非空字符串")
|
||||
|
||||
meta = ws.get("meta")
|
||||
if not isinstance(meta, dict):
|
||||
errors.append("meta 必须是 object")
|
||||
else:
|
||||
if meta.get("status") not in STATUS_FLOW:
|
||||
errors.append(f"meta.status 非法: {meta.get('status')}")
|
||||
budget = meta.get("budget")
|
||||
if not isinstance(budget, dict):
|
||||
errors.append("meta.budget 必须是 object")
|
||||
else:
|
||||
for k in ("api_input_tokens", "api_output_tokens", "api_token_cap", "rounds_cap"):
|
||||
if not isinstance(budget.get(k), int) or budget.get(k) < 0:
|
||||
errors.append(f"meta.budget.{k} 必须是非负整数")
|
||||
|
||||
brief = ws.get("brief")
|
||||
if brief is not None:
|
||||
if not isinstance(brief, dict):
|
||||
errors.append("brief 必须是 object")
|
||||
else:
|
||||
if not isinstance(brief.get("goal"), str):
|
||||
errors.append("brief.goal 必须是字符串")
|
||||
elif len(brief["goal"]) > LIMITS["goal"]:
|
||||
errors.append(f"brief.goal 超长（>{LIMITS['goal']}字）")
|
||||
if not isinstance(brief.get("constraints"), list):
|
||||
errors.append("brief.constraints 必须是数组")
|
||||
elif len(brief["constraints"]) > LIMITS["constraints"]:
|
||||
errors.append(f"brief.constraints 超过 {LIMITS['constraints']} 条")
|
||||
if not isinstance(brief.get("tags"), list):
|
||||
errors.append("brief.tags 必须是数组")
|
||||
for t in brief.get("tags", []) or []:
|
||||
if t not in ALLOWED_TAGS:
|
||||
errors.append(f"brief.tags 含非法标签: {t}")
|
||||
acc = brief.get("acceptance")
|
||||
if not isinstance(acc, list) or len(acc) > LIMITS["acceptance"]:
|
||||
errors.append(f"brief.acceptance 需为 ≤{LIMITS['acceptance']} 的数组")
|
||||
plan = brief.get("plan")
|
||||
if not isinstance(plan, list) or len(plan) > LIMITS["plan_steps"]:
|
||||
errors.append(f"brief.plan 需为 ≤{LIMITS['plan_steps']} 步的数组")
|
||||
else:
|
||||
ids = [p.get("id") for p in plan if isinstance(p, dict)]
|
||||
if len(set(ids)) != len(ids):
|
||||
errors.append("brief.plan 存在重复 step id")
|
||||
for p in plan:
|
||||
if not isinstance(p, dict):
|
||||
errors.append("brief.plan 元素必须是 object")
|
||||
continue
|
||||
if not isinstance(p.get("task"), str):
|
||||
errors.append(f"brief.plan[{p.get('id')}].task 必须是字符串")
|
||||
elif len(p["task"]) > LIMITS["task"]:
|
||||
errors.append(f"brief.plan[{p.get('id')}].task 超长（>{LIMITS['task']}字）")
|
||||
|
||||
for i, entry in enumerate(ws.get("progress", []) or []):
|
||||
if not isinstance(entry, dict):
|
||||
errors.append(f"progress[{i}] 必须是 object"); continue
|
||||
if entry.get("status") not in ("done", "failed", "blocked"):
|
||||
errors.append(f"progress[{i}].status 非法")
|
||||
if not isinstance(entry.get("summary"), str) or len(entry["summary"]) > LIMITS["summary"]:
|
||||
errors.append(f"progress[{i}].summary 非法或超长")
|
||||
|
||||
for i, entry in enumerate(ws.get("issues", []) or []):
|
||||
if not isinstance(entry, dict):
|
||||
errors.append(f"issues[{i}] 必须是 object"); continue
|
||||
for k in ("observed", "expected", "tried", "ask"):
|
||||
if isinstance(entry.get(k), str) and len(entry[k]) > LIMITS["issue_text"]:
|
||||
errors.append(f"issues[{i}].{k} 超长（>{LIMITS['issue_text']}字）")
|
||||
|
||||
for i, entry in enumerate(ws.get("decisions", []) or []):
|
||||
if not isinstance(entry, dict):
|
||||
errors.append(f"decisions[{i}] 必须是 object"); continue
|
||||
if isinstance(entry.get("reply"), str) and len(entry["reply"]) > LIMITS["reply"]:
|
||||
errors.append(f"decisions[{i}].reply 超长（>{LIMITS['reply']}字）")
|
||||
|
||||
for i, line in enumerate(ws.get("archive", []) or []):
|
||||
if not isinstance(line, str) or len(line) > LIMITS["archive"]:
|
||||
errors.append(f"archive[{i}] 非法或超长")
|
||||
return errors
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workspace
|
||||
# ---------------------------------------------------------------------------
|
||||
class Workspace:
|
||||
"""交流文本对象：持有状态、执行写入前校验、渲染、rollup、持久化。"""
|
||||
|
||||
def __init__(self, data: Dict[str, Any]):
|
||||
errors = validate(data)
|
||||
if errors:
|
||||
raise ValueError("workspace 校验失败: " + "; ".join(errors[:5]))
|
||||
self._data = data
|
||||
self._brief_locked = False
|
||||
|
||||
# ---------- 构造 ----------
|
||||
@classmethod
|
||||
def new(cls, request_id: str, query: str,
|
||||
api_token_cap: int = 8000, rounds_cap: int = 6) -> "Workspace":
|
||||
data = {
|
||||
"version": VERSION,
|
||||
"request_id": request_id,
|
||||
"query": query,
|
||||
"meta": {
|
||||
"status": "draft",
|
||||
"round": 0,
|
||||
"budget": {
|
||||
"api_input_tokens": 0,
|
||||
"api_output_tokens": 0,
|
||||
"api_token_cap": api_token_cap,
|
||||
"rounds_cap": rounds_cap,
|
||||
},
|
||||
},
|
||||
"brief": None,
|
||||
"progress": [],
|
||||
"issues": [],
|
||||
"decisions": [],
|
||||
"archive": [],
|
||||
}
|
||||
return cls(data)
|
||||
|
||||
# ---------- 访问 ----------
|
||||
@property
|
||||
def request_id(self) -> str:
|
||||
return self._data["request_id"]
|
||||
|
||||
@property
|
||||
def status(self) -> str:
|
||||
return self._data["meta"]["status"]
|
||||
|
||||
@property
|
||||
def data(self) -> Dict[str, Any]:
|
||||
return copy.deepcopy(self._data)
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
return self._data.get(key, default)
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
return self._data[key]
|
||||
|
||||
def meta(self) -> Dict[str, Any]:
|
||||
return self._data["meta"]
|
||||
|
||||
def budget(self) -> Dict[str, int]:
|
||||
return self._data["meta"]["budget"]
|
||||
|
||||
# ---------- 写入（均先校验） ----------
|
||||
def _commit(self, data: Dict[str, Any]) -> None:
|
||||
errors = validate(data)
|
||||
if errors:
|
||||
raise ValueError("写入校验失败: " + "; ".join(errors[:5]))
|
||||
self._data = data
|
||||
|
||||
def transition(self, new_status: str) -> None:
|
||||
cur = self.status
|
||||
if new_status == cur:
|
||||
return
|
||||
if new_status not in STATUS_FLOW.get(cur, set()):
|
||||
raise ValueError(f"非法状态迁移: {cur} -> {new_status}")
|
||||
self._data["meta"]["status"] = new_status
|
||||
|
||||
def apply_brief(self, brief: Dict[str, Any]) -> None:
|
||||
"""写入 brief（写一次后锁定，D2：brief 恒定位于文档前部，prefix cache 友好）。"""
|
||||
if self._data.get("brief") is not None or self._brief_locked:
|
||||
raise ValueError("brief 已写入，不可重复")
|
||||
new = copy.deepcopy(self._data)
|
||||
new["brief"] = brief
|
||||
new["meta"]["status"] = "in_progress"
|
||||
self._commit(new)
|
||||
self._brief_locked = True
|
||||
|
||||
def add_progress(self, step: str, status: str, summary: str,
|
||||
artifact: Optional[str] = None) -> None:
|
||||
new = copy.deepcopy(self._data)
|
||||
entry: Dict[str, Any] = {"step": step, "status": status, "summary": summary}
|
||||
if artifact:
|
||||
entry["artifact"] = artifact
|
||||
new["progress"].append(entry)
|
||||
self._commit(new)
|
||||
|
||||
def add_issue(self, step: str, anchor: str, observed: str, expected: str,
|
||||
tried: str, ask: str) -> str:
|
||||
new = copy.deepcopy(self._data)
|
||||
iid = f"i{len(new['issues']) + 1}"
|
||||
entry = {
|
||||
"id": iid, "step": step, "anchor": anchor,
|
||||
"observed": observed, "expected": expected,
|
||||
"tried": tried, "ask": ask,
|
||||
}
|
||||
new["issues"].append(entry)
|
||||
self._commit(new)
|
||||
return iid
|
||||
|
||||
def add_decision(self, ref: str, reply: str,
|
||||
patch_plan: Optional[List[Dict[str, str]]] = None) -> None:
|
||||
new = copy.deepcopy(self._data)
|
||||
new["decisions"].append({
|
||||
"ref": ref, "reply": reply,
|
||||
"patch_plan": patch_plan or [],
|
||||
})
|
||||
self._commit(new)
|
||||
|
||||
def revise_plan(self, updates: Dict[str, str]) -> None:
|
||||
"""按 decision.patch_plan 修订既有 step 的 task（不改结构/顺序）。"""
|
||||
if self._data.get("brief") is None:
|
||||
raise ValueError("brief 尚未写入，无法修订 plan")
|
||||
new = copy.deepcopy(self._data)
|
||||
for pid, task in updates.items():
|
||||
for p in new["brief"]["plan"]:
|
||||
if p["id"] == pid:
|
||||
p["task"] = task
|
||||
break
|
||||
self._commit(new)
|
||||
|
||||
def mark_round(self) -> None:
|
||||
self._data["meta"]["round"] += 1
|
||||
|
||||
def add_budget(self, input_tokens: int = 0, output_tokens: int = 0) -> None:
|
||||
b = self._data["meta"]["budget"]
|
||||
b["api_input_tokens"] += int(input_tokens)
|
||||
b["api_output_tokens"] += int(output_tokens)
|
||||
self._commit(self._data)
|
||||
|
||||
def exhausted(self) -> bool:
|
||||
"""预算熔断判定：API token 或回合任一触顶（D6）。"""
|
||||
b = self._data["meta"]["budget"]
|
||||
used = b["api_input_tokens"] + b["api_output_tokens"]
|
||||
if b["api_token_cap"] and used >= b["api_token_cap"]:
|
||||
return True
|
||||
if b["rounds_cap"] and self._data["meta"]["round"] >= b["rounds_cap"]:
|
||||
return True
|
||||
return False
|
||||
|
||||
# ---------- rollup（4.5） ----------
|
||||
def rollup(self) -> int:
|
||||
"""把 done 的 progress 折叠为 archive 摘要行，并清理解析完成的问题。
|
||||
|
||||
只减不删 archive 历史；progress 中 done 的条目折叠后移除（保留 failed/blocked）。
|
||||
返回本次折叠的条目数。
|
||||
"""
|
||||
folded = 0
|
||||
new_progress: List[Dict[str, Any]] = []
|
||||
for entry in self._data.get("progress", []):
|
||||
if entry.get("status") == "done" and entry.get("step"):
|
||||
line = _clip(f"{entry['step']}: {entry.get('summary', '')}", LIMITS["archive"])
|
||||
if line not in self._data["archive"]:
|
||||
self._data["archive"].append(line)
|
||||
folded += 1
|
||||
else:
|
||||
new_progress.append(entry)
|
||||
resolved = {d.get("ref") for d in self._data.get("decisions", [])}
|
||||
kept_issues: List[Dict[str, Any]] = []
|
||||
for iss in self._data.get("issues", []):
|
||||
if iss.get("id") in resolved:
|
||||
line = _clip(f"{iss['id']}: {iss.get('expected', '')[:60]}", LIMITS["archive"])
|
||||
if line not in self._data["archive"]:
|
||||
self._data["archive"].append(line)
|
||||
else:
|
||||
kept_issues.append(iss)
|
||||
self._data["issues"] = kept_issues
|
||||
self._data["progress"] = new_progress
|
||||
return folded
|
||||
|
||||
# ---------- 渲染 ----------
|
||||
def render_for_architect(self) -> str:
|
||||
"""渲染 Architect 输入（D7）：meta+query(截断)+全部 issues+最近3条 decisions
|
||||
+最近回合 progress 摘要。目标 ≤1200 token。"""
|
||||
d = self._data
|
||||
parts: List[str] = []
|
||||
m = d["meta"]
|
||||
parts.append("== meta ==")
|
||||
parts.append(f"status={m['status']} round={m['round']} "
|
||||
f"budget={json.dumps(m['budget'], ensure_ascii=False)}")
|
||||
parts.append("== query ==")
|
||||
parts.append(_clip(d["query"], LIMITS["query_truncate"]))
|
||||
if d.get("brief"):
|
||||
b = d["brief"]
|
||||
parts.append("== brief(锁定) ==")
|
||||
parts.append(f"goal: {_clip(b['goal'], 120)}")
|
||||
parts.append(f"plan: {[p['id'] for p in b.get('plan', [])]}")
|
||||
parts.append(f"acceptance: {[a.get('id') for a in b.get('acceptance', [])]}")
|
||||
parts.append("== issues ==")
|
||||
for iss in d.get("issues", []):
|
||||
parts.append(f"{iss['id']} step={iss.get('step')} anchor={iss.get('anchor')} "
|
||||
f"ask={_clip(iss.get('ask', ''), 80)}")
|
||||
parts.append("== 最近 3 条 decisions ==")
|
||||
for dec in d.get("decisions", [])[-3:]:
|
||||
parts.append(f"ref={dec.get('ref')} reply={_clip(dec.get('reply', ''), 80)}")
|
||||
parts.append("== progress 摘要 ==")
|
||||
for p in d.get("progress", [])[-5:]:
|
||||
parts.append(f"{p.get('step')} [{p.get('status')}] {_clip(p.get('summary', ''), 40)}")
|
||||
parts.append("== archive ==")
|
||||
for line in d.get("archive", [])[-8:]:
|
||||
parts.append(line)
|
||||
# token 预算：超限先截断最旧 archive（已只保留 3 条 decisions）
|
||||
out = "\n".join(parts)
|
||||
while estimate_tokens(out) > 1200 and len(d.get("archive", [])) > 4:
|
||||
d = copy.deepcopy(d)
|
||||
d["archive"] = d["archive"][4:]
|
||||
out = "\n".join(_rerender(self, d))
|
||||
return out
|
||||
|
||||
def render_for_worker(self, step_id: str,
|
||||
artifact_text: Optional[str] = None) -> str:
|
||||
"""渲染 Worker 输入（4.4）：brief 全文+该 step 定义+依赖 step 的 archive 摘要行
|
||||
+该 step 现有工件全文+验收标准。目标 ≤8K token。"""
|
||||
d = self._data
|
||||
b = d.get("brief")
|
||||
parts: List[str] = []
|
||||
if b:
|
||||
parts.append("== 任务目标 (goal) ==")
|
||||
parts.append(b["goal"])
|
||||
parts.append("== 约束 (constraints) ==")
|
||||
parts.extend(f"- {c}" for c in b.get("constraints", []))
|
||||
parts.append("== 全部步骤 (plan) ==")
|
||||
for p in b.get("plan", []):
|
||||
mark = " <-- 当前步" if p.get("id") == step_id else ""
|
||||
parts.append(f"{p['id']}: {p.get('task', '')}{mark}")
|
||||
parts.append(f" done_criteria: {p.get('done_criteria', '')}")
|
||||
parts.append("== 依赖步摘要 (archive) ==")
|
||||
for line in d.get("archive", [])[-6:]:
|
||||
parts.append(line)
|
||||
if artifact_text:
|
||||
parts.append(f"== 当前步已有工件（{step_id}） ==")
|
||||
parts.append(artifact_text)
|
||||
parts.append("== 验收标准 ==")
|
||||
if b:
|
||||
for a in b.get("acceptance", []):
|
||||
parts.append(f"- {a.get('id')}: {a.get('check', '')} "
|
||||
f"(machine_checkable={a.get('machine_checkable', False)})")
|
||||
parts.append("== 要求 ==")
|
||||
parts.append("请实现当前步，并用可执行验证/事实对照/结构检查自验证；"
|
||||
"通过则写 progress(done)，失败自修 ≤2 次，仍失败则写 issue。")
|
||||
return "\n".join(parts)
|
||||
|
||||
# ---------- 持久化 ----------
|
||||
def save(self, path: Path) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(self._data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Path) -> "Workspace":
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return cls(data)
|
||||
|
||||
def prefix_signature(self) -> str:
|
||||
"""返回稳定前缀的签名（T10 prefix cache）。
|
||||
|
||||
交流文本的"恒定位于前部"的部分（version + request_id + query + meta + brief）
|
||||
应不随 progress/issues/decisions 追加而变化，从而使 llama-server 的
|
||||
--cache-reuse 能命中该前缀、降低 prefill 开销。用紧凑 JSON 的哈希度量稳定性。
|
||||
"""
|
||||
stable = {
|
||||
"version": self._data.get("version"),
|
||||
"request_id": self._data.get("request_id"),
|
||||
"query": self._data.get("query"),
|
||||
"brief": self._data.get("brief"),
|
||||
}
|
||||
import hashlib
|
||||
s = json.dumps(stable, ensure_ascii=False, sort_keys=True)
|
||||
return hashlib.sha256(s.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return self.data
|
||||
|
||||
|
||||
def _rerender(ws: "Workspace", d: Dict[str, Any]) -> List[str]:
|
||||
"""用裁剪后的数据重建 Architect 渲染（供超限压缩内部用）。"""
|
||||
parts: List[str] = []
|
||||
m = d["meta"]
|
||||
parts.append("== meta ==")
|
||||
parts.append(f"status={m['status']} round={m['round']}")
|
||||
parts.append("== query ==")
|
||||
parts.append(_clip(d["query"], LIMITS["query_truncate"]))
|
||||
parts.append("== issues ==")
|
||||
for iss in d.get("issues", []):
|
||||
parts.append(f"{iss['id']} step={iss.get('step')} ask={_clip(iss.get('ask', ''), 80)}")
|
||||
parts.append("== decisions(最近3) ==")
|
||||
for dec in d.get("decisions", [])[-3:]:
|
||||
parts.append(f"ref={dec.get('ref')} reply={_clip(dec.get('reply', ''), 80)}")
|
||||
parts.append("== progress ==")
|
||||
for p in d.get("progress", [])[-5:]:
|
||||
parts.append(f"{p.get('step')} [{p.get('status')}] {_clip(p.get('summary', ''), 40)}")
|
||||
parts.append("== archive ==")
|
||||
for line in d.get("archive", [])[-6:]:
|
||||
parts.append(line)
|
||||
return parts
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
"""测试替身：模拟 llama-server（供 LlamaServerManager 封闭单测，D11）。
|
||||
|
||||
- 解析 --port / -m / -ngl / -c（与真实 llama-server 参数对齐）
|
||||
- 把 pid / 收到的参数写入环境变量 FAKE_MARKER 指向的 JSON 文件
|
||||
- 在本机端口起一个最小 http 服务：/health 返回 {"status":"ok"}
|
||||
- 进程被终止时正常退出
|
||||
"""
|
||||
import argparse
|
||||
import http.server
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(prog="fake-llama-server")
|
||||
parser.add_argument("--port", type=int, default=8901)
|
||||
parser.add_argument("-m", dest="model", default="")
|
||||
parser.add_argument("-ngl", dest="ngl", default="0")
|
||||
parser.add_argument("-c", dest="ctx", default="8192")
|
||||
parser.add_argument("-ctk", dest="ctk", default="")
|
||||
parser.add_argument("-ctv", dest="ctv", default="")
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
marker = os.environ.get("FAKE_MARKER")
|
||||
if marker:
|
||||
os.makedirs(os.path.dirname(marker) or ".", exist_ok=True)
|
||||
with open(marker, "w", encoding="utf-8") as f:
|
||||
json.dump({"pid": os.getpid(), "port": args.port,
|
||||
"model": args.model, "args": sys.argv[1:]}, f)
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
if self.path.startswith("/health"):
|
||||
body = json.dumps({"status": "ok", "server": "fake-llama"}).encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, *a):
|
||||
pass
|
||||
|
||||
srv = http.server.ThreadingHTTPServer(("127.0.0.1", args.port), Handler)
|
||||
srv.serve_forever()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user