213 lines
9.2 KiB
Python
213 lines
9.2 KiB
Python
"""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())
|