319 lines
12 KiB
Python
319 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ao-yaml-2b-runner.py — ao workflow → 2B 智能体串行执行器
|
||
==========================================================
|
||
牧尘 2026-09-10 扩展定案：把 ao compose 生成的 workflow（234 角色）映射到 2B 智能体上，
|
||
串行执行完成复杂任务。
|
||
|
||
核心：
|
||
- 解析 ao workflow YAML（role + task + output + acceptance）
|
||
- 每步 = 一个 2B agent 调用：role 的 markdown prompt 文件（角色人格）+ task（任务）+ 前步输出
|
||
- 用步骤自带 acceptance 做验收（确定性检查失败才重试）
|
||
- 变量传递：{{output_var}} 替换为前步实际输出
|
||
- 一次一个 agent 串行
|
||
|
||
用法:
|
||
python3 ao-yaml-2b-runner.py workflow.yaml [--workdir /tmp/xxx]
|
||
python3 ao-yaml-2b-runner.py workflow.yaml --max-rework 2 --force-2b
|
||
"""
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
AO_AGENTS = "/home/muc/.npm/_npx/097cc02411e2da65/node_modules/agency-orchestrator/agency-agents"
|
||
LLM_URL = "http://127.0.0.1:8080/v1/chat/completions"
|
||
LLM_MODEL = "minicpm5-2b"
|
||
LOCK = "/tmp/ao-yaml-2b.lock"
|
||
MAX_OUT_CHARS = 3000
|
||
|
||
|
||
def sh(cmd, timeout=120, cwd=None):
|
||
try:
|
||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout,
|
||
cwd=cwd, env=dict(os.environ,
|
||
PATH="/home/muc/nodejs/node-v24.16.0-linux-x64/bin:" + os.environ.get("PATH", "")))
|
||
return r.returncode, (r.stdout + r.stderr)
|
||
except Exception as e:
|
||
return -1, str(e)[:500]
|
||
|
||
|
||
def acquire_lock():
|
||
if os.path.exists(LOCK):
|
||
age = time.time() - os.path.getmtime(LOCK)
|
||
if age < 3600:
|
||
return False, "已有 ao-2b 执行在跑"
|
||
os.unlink(LOCK)
|
||
open(LOCK, "w").write(str(time.time()))
|
||
return True, ""
|
||
|
||
|
||
def release_lock():
|
||
try:
|
||
os.unlink(LOCK)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def find_role_file(role_path):
|
||
"""role: 'product/product-trend-researcher' → 找 md 文件"""
|
||
candidates = [
|
||
os.path.join(AO_AGENTS, role_path + ".md"),
|
||
os.path.join(AO_AGENTS, role_path),
|
||
]
|
||
# 也搜 .md 尾
|
||
for c in candidates:
|
||
if os.path.exists(c):
|
||
return c
|
||
# 宽松：按文件名找
|
||
name = os.path.basename(role_path)
|
||
for root, _, files in os.walk(AO_AGENTS):
|
||
for f in files:
|
||
if f == name + ".md" or f == name:
|
||
return os.path.join(root, f)
|
||
return None
|
||
|
||
|
||
def load_role_prompt(role_path):
|
||
"""读角色 md，提取人格部分（frontmatter 后的正文，截到工具/命令说明前）"""
|
||
fp = find_role_file(role_path)
|
||
if not fp:
|
||
return f"(角色 {role_path} 未找到，你是一名专业助手)\n"
|
||
content = open(fp).read()
|
||
# 去 frontmatter
|
||
if content.startswith("---"):
|
||
parts = content.split("---", 2)
|
||
content = parts[2] if len(parts) > 2 else content
|
||
# 保留前 2500 字符（人格+职责核心，截断过长工具列表）
|
||
return content[:2500]
|
||
|
||
|
||
def call_2b_json(system_prompt, user_prompt, timeout=200):
|
||
"""调 2B 并解析 JSON 输出"""
|
||
import tempfile
|
||
prompt = f"{system_prompt}\n\n===== 任务 =====\n{user_prompt}\n\n🔴 只输出原始 JSON，禁止 markdown/代码围栏/前后缀。"
|
||
try:
|
||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
|
||
json.dump({"model": LLM_MODEL,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"temperature": 0.3, "max_tokens": 2000}, f, ensure_ascii=False)
|
||
bf = f.name
|
||
rc, out = sh(["curl", "-s", "-m", str(timeout), LLM_URL,
|
||
"-H", "Content-Type: application/json",
|
||
"--data", f"@{bf}"], timeout=timeout + 15)
|
||
os.unlink(bf)
|
||
return out
|
||
except Exception as e:
|
||
return f'{{"error": "{e}"}}'
|
||
|
||
|
||
def extract_json(text):
|
||
"""从 2B 响应提取 {content, self_check}。
|
||
2B 输出 JSON 可能残缺（少括号）或先贴内容再出自评 → 分层策略：
|
||
1. 正则直接找 "content" 字段值（最可靠，容错残缺 JSON）
|
||
2. 完整平衡 candidates 中含 content/summary 键的对象
|
||
3. code fence 里的 JSON
|
||
4. 全文兜底
|
||
"""
|
||
if not text:
|
||
return None
|
||
# 如果整体是 API 响应，先取 message.content
|
||
try:
|
||
outer = json.loads(text)
|
||
if "choices" in outer:
|
||
inner = outer["choices"][0]["message"].get("content")
|
||
if isinstance(inner, str):
|
||
text = inner
|
||
except Exception:
|
||
pass
|
||
|
||
def _content_from_text(t):
|
||
# 正则提取 content 字段（处理转义）
|
||
cm = re.search(r'"content"\s*:\s*"((?:[^"\\]|\\.)*)"', t)
|
||
if cm:
|
||
try:
|
||
val = json.loads('"' + cm.group(1) + '"')
|
||
scm = re.search(r'"pass"\s*:\s*(true|false)', t)
|
||
return {"content": val,
|
||
"self_check": {"pass": scm.group(1) == "true" if scm else True,
|
||
"reason": "regex-extracted"}}
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
# 1. regex 优先
|
||
r1 = _content_from_text(text)
|
||
if r1:
|
||
return r1
|
||
# 2. 平衡 candidates
|
||
candidates = []
|
||
depth = 0
|
||
start = None
|
||
for i in range(len(text) - 1, -1, -1):
|
||
if text[i] == "}":
|
||
if depth == 0:
|
||
start = i
|
||
depth += 1
|
||
elif text[i] == "{":
|
||
depth -= 1
|
||
if depth == 0 and start is not None:
|
||
try:
|
||
candidates.append(json.loads(text[i:start + 1]))
|
||
except Exception:
|
||
pass
|
||
start = None
|
||
for c in candidates:
|
||
if isinstance(c, dict) and any(k in c for k in ("content", "summary", "report", "findings", "analysis")):
|
||
# 若 c 含嵌套 self_check 且 content 缺失内容，尝试取 content 子串
|
||
return c
|
||
if candidates:
|
||
return candidates[0]
|
||
# 3. code fence
|
||
m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
|
||
if m:
|
||
try:
|
||
return json.loads(m.group(1))
|
||
except Exception:
|
||
pass
|
||
# 4. 全文兜底
|
||
return {"content": text.strip()[:3000], "self_check": {"pass": True, "reason": "raw fallback"}}
|
||
|
||
|
||
def resolve_vars(template, outputs):
|
||
"""替换 {{var}} 为前步输出"""
|
||
def repl(m):
|
||
key = m.group(1).strip()
|
||
if key in outputs:
|
||
v = outputs[key]
|
||
return v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
|
||
return m.group(0)
|
||
return re.sub(r"\{\{\s*(\w+)\s*\}\}", repl, template)
|
||
|
||
|
||
def run_step(step, outputs, workdir, force_2b=True):
|
||
"""执行单个步骤（一个 2B 角色 agent）— 纯文本输出模式
|
||
2B 输出 JSON 包装会挤占正文内容（实测系统性失败）→ 一律直接输出正文。
|
||
"""
|
||
step_id = step.get("id", "?")
|
||
role = step.get("role", "")
|
||
task = step.get("task", "")
|
||
acceptance = step.get("acceptance", "")
|
||
print(f"\n=== 步骤 {step_id}: {step.get('name', role)} ({role}) ===", flush=True)
|
||
|
||
# human_input 步骤：跳过（需要人工），或提示
|
||
if step.get("type") == "human_input":
|
||
print(" ⏭️ human_input 步骤跳过（需要人工输入）", flush=True)
|
||
outputs[step.get("output")] = step.get("default", "")
|
||
return True
|
||
|
||
# 组装 prompt：角色人格 + 任务（含前步变量）
|
||
role_prompt = load_role_prompt(role)
|
||
task_resolved = resolve_vars(task, outputs)
|
||
ctx = "\n".join(f"[{k}] {str(v)[:MAX_OUT_CHARS]}" for k, v in outputs.items())
|
||
user = f"{task_resolved}\n\n===== 前序上下文 =====\n{ctx or '(无)'}\n\n===== 验收标准 =====\n{acceptance or '(无)'}\n\n"
|
||
# 🔴 硬约束：直接输出成果正文。禁止 JSON、禁止自评、禁止开场白。
|
||
user += '\n🔴 直接输出你的完整成果正文（纯文本/markdown 均可，标题列表要点都行）。\n不要输出 JSON 对象，不要写 "content:" 前缀，不要自评，不要解释你做了什么，不要开场白。第一行就是成果本身。'
|
||
|
||
resp = call_2b_json(role_prompt, user)
|
||
# 取 API message.content 原文
|
||
content = resp
|
||
try:
|
||
outer = json.loads(resp)
|
||
content = outer["choices"][0]["message"]["content"]
|
||
if not isinstance(content, str):
|
||
content = str(content)
|
||
except Exception:
|
||
pass
|
||
# 若 2B 仍顽固输出 JSON 壳，剥掉
|
||
if content.lstrip().startswith("{") and "self_check" in content:
|
||
parsed = extract_json(content)
|
||
if isinstance(parsed, dict):
|
||
inner = parsed.get("content")
|
||
if isinstance(inner, str) and len(inner) > 10:
|
||
content = inner
|
||
elif isinstance(inner, dict):
|
||
content = str(inner)
|
||
content = (content or "").strip()
|
||
if not content or len(content) < 20:
|
||
print(" ⚠️ 输出过短，原样保存", flush=True)
|
||
outputs[step.get("output")] = content[:MAX_OUT_CHARS * 3]
|
||
_save_output(workdir, step_id, content, resp)
|
||
print(f" ✅ 输出 {len(content)} 字符", flush=True)
|
||
return True
|
||
|
||
|
||
def _save_output(workdir, step_id, content, raw=""):
|
||
os.makedirs(os.path.join(workdir, "steps"), exist_ok=True)
|
||
with open(os.path.join(workdir, "steps", f"{step_id}.txt"), "w") as f:
|
||
f.write(str(content)[:5000])
|
||
with open(os.path.join(workdir, "steps", f"{step_id}.raw.json"), "w") as f:
|
||
f.write(str(raw)[:8000])
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("yaml", help="ao workflow YAML 路径")
|
||
ap.add_argument("--workdir", default="/tmp/ao-2b-work")
|
||
ap.add_argument("--max-rework", type=int, default=2)
|
||
args = ap.parse_args()
|
||
|
||
ok, msg = acquire_lock()
|
||
if not ok:
|
||
print(json.dumps({"ok": False, "error": msg}))
|
||
return
|
||
os.makedirs(args.workdir, exist_ok=True)
|
||
|
||
try:
|
||
import yaml
|
||
wf = yaml.safe_load(open(args.yaml))
|
||
steps = wf.get("steps", [])
|
||
print(f"工作流: {wf.get('name')} | 步骤: {len(steps)}")
|
||
outputs = {}
|
||
|
||
# 计算依赖顺序（简化：按 depends_on 拓扑排序）
|
||
done = set()
|
||
order = []
|
||
remaining = list(steps)
|
||
while remaining:
|
||
progressed = False
|
||
for s in list(remaining):
|
||
deps = s.get("depends_on") or []
|
||
if all(d in done or d == s.get("id") for d in deps):
|
||
order.append(s)
|
||
done.add(s.get("id"))
|
||
remaining.remove(s)
|
||
progressed = True
|
||
if not progressed:
|
||
# 有环或缺失依赖 → 剩余的按顺序执行
|
||
order.extend(remaining)
|
||
break
|
||
|
||
for step in order:
|
||
ok_step = run_step(step, outputs, args.workdir)
|
||
if not ok_step:
|
||
print(json.dumps({"ok": False, "step": step.get("id")}))
|
||
return
|
||
|
||
# 汇总：找最终输出步骤（output=final_report 或最后一个）
|
||
final_keys = [s.get("output") for s in steps if "final" in str(s.get("output", "")).lower() or "report" in str(s.get("output", "")).lower()]
|
||
result = {k: outputs.get(k, "")[:500] for k in final_keys if outputs.get(k)}
|
||
if not result and outputs:
|
||
last_out = steps[-1].get("output") if steps else None
|
||
result = {"final": str(outputs.get(last_out, list(outputs.values())[-1]))[:500]}
|
||
with open(os.path.join(args.workdir, "final_result.json"), "w") as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
print("\n=== 完成 ===")
|
||
for k, v in result.items():
|
||
print(f"[{k}] {str(v)[:200]}")
|
||
print(json.dumps({"ok": True, "workdir": args.workdir,
|
||
"steps_done": len(order)}))
|
||
finally:
|
||
release_lock()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|