fengqun/examples/export_model_agnet_io_report.py

from __future__ import annotations

from pathlib import Path
import re
import sys


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from swarm_minimal.azure_store import PostgresRedisBlobSwarmStore
from swarm_minimal.config import SwarmConfig
from swarm_minimal.local_env import load_env_file


RUN_IDS = [
    "78f189ccd1924ed0a4fb0a0a447ad449",
    "623b5f6e5cc24cc7967fd9577f9c224b",
    "42688e7b6245466dab8398dfe4790456",
]

OUTPUT_PATH = ROOT / "MODEL_AGNET_IO_REPORT.zh-CN.md"


def main() -> None:
    load_env_file(ROOT / ".env")
    store = PostgresRedisBlobSwarmStore(SwarmConfig.from_env())
    try:
        store.ensure_schema()
        report = build_report(store)
        OUTPUT_PATH.write_text(report, encoding="utf-8")
        print(str(OUTPUT_PATH))
    finally:
        store.close()


def build_report(store: PostgresRedisBlobSwarmStore) -> str:
    sections = [
        "# 模型 / Agnet 输入输出报告",
        "",
        "这份报告直接从 Azure PostgreSQL 的 `swarm_convergence`、`swarm_tasks` 和 `swarm_observations` 读取历史 live run。",
        "报告只展示任务输入、模型提示词模板、Agnet 输出和评分，不展示 `.env` 或任何密钥。",
        "",
    ]
    for run_id in RUN_IDS:
        convergence = fetch_convergence(store, run_id)
        if convergence is None:
            sections.extend([f"## Run `{run_id}`", "", "未找到该 run。", ""])
            continue
        task_ids = [item["task_id"] for item in convergence["observations"]]
        tasks = fetch_tasks(store, task_ids)
        sections.extend(render_run(convergence, tasks))
    return "\n".join(sections).rstrip() + "\n"


def render_run(convergence: dict[str, object], tasks: dict[str, dict[str, object]]) -> list[str]:
    run_id = str(convergence["run_id"])
    goal = str(convergence["goal"])
    system_prompt = system_prompt_for_goal(goal)
    lines = [
        f"## Run `{run_id}`",
        "",
        f"- 任务目标：{goal}",
        f"- 完成任务数：{convergence['completed_tasks']}",
        f"- 收敛分数：{convergence['accepted_score']}",
        f"- Blob artifact：`{convergence['artifact_path']}`",
        "",
        "### 模型系统提示词",
        "",
        "```text",
        system_prompt,
        "```",
        "",
    ]

    for index, observation in enumerate(convergence["observations"], start=1):
        task_id = observation["task_id"]
        task = tasks[task_id]
        output = str(task.get("output") or "")
        model = infer_model(task, output)
        lines.extend(
            [
                f"### Agnet 调用 {index}: `{task['kind']}`",
                "",
                f"- Agnet：`{task['claimed_by']}`",
                f"- 模型：`{model}`",
                f"- 状态：`{task['status']}`",
                f"- 分数：`{task['score']}`",
                f"- 观测信号：`{observation['signal']}`",
                "",
                "#### 给模型的 user prompt 结构",
                "",
                "```text",
                user_prompt_shape_for_goal(goal),
                "```",
                "",
                "#### 本次任务输入 task.input",
                "",
                "```text",
                str(task["input"]).strip(),
                "```",
                "",
                "#### Agnet / 模型实际输出 task.output",
                "",
                "```text",
                output.strip(),
                "```",
                "",
            ]
        )
    return lines


def system_prompt_for_goal(goal: str) -> str:
    if goal.startswith("full live test"):
        return (
            "You are a minimal Agnet worker inside a swarm. "
            "Return a concise result that can be scored and converged."
        )
    if goal.startswith("真实全面场景"):
        return (
            "You are a senior coding/algorithm agent in a multi-task swarm. "
            "Return a concrete engineering answer for the assigned subtask. "
            "Do not include secrets."
        )
    if goal.startswith("连续性长推理场景"):
        return (
            "You are one stage in a continuous long-reasoning swarm. "
            "Carry forward prior conclusions, expose risks, and hand off a concise next-state. "
            "Do not reveal secrets."
        )
    return "<unknown system prompt>"


def user_prompt_shape_for_goal(goal: str) -> str:
    if goal.startswith("full live test"):
        return "\n".join(
            [
                "Task kind: <task.kind>",
                "Task input: <task.input>",
                "Known shared state keys: <sorted(shared_state.keys())>",
            ]
        )
    if goal.startswith("真实全面场景"):
        return "\n".join(
            [
                "Task kind: <task.kind>",
                "Task input:",
                "<task.input>",
                "",
                "Shared state keys: <sorted(shared_state.keys())>",
            ]
        )
    if goal.startswith("连续性长推理场景"):
        return "\n".join(
            [
                "Previous marker: <previous_marker>",
                "Previous summary:",
                "<previous_step_summary>",
                "",
                "Task kind: <task.kind>",
                "Task input:",
                "<task.input>",
            ]
        )
    return "<unknown user prompt shape>"


def infer_model(task: dict[str, object], output: str) -> str:
    for pattern in [r"used_model=([^;\\n]+)", r"primary_model=([^;\\n]+)", r"model=([^;\\n]+)"]:
        match = re.search(pattern, output)
        if match:
            return match.group(1).strip()
    match = re.search(r"model=([^;\\n]+)", str(task.get("input") or ""))
    if match:
        return match.group(1).strip()
    return "<unknown>"


def fetch_convergence(store: PostgresRedisBlobSwarmStore, run_id: str) -> dict[str, object] | None:
    def operation(cur):
        cur.execute(
            """
            select run_id, goal, accepted_score, completed_tasks, observations, artifact_path, created_at
            from swarm_convergence
            where run_id = %s
            """,
            (run_id,),
        )
        row = cur.fetchone()
        if row is None:
            return None
        return {
            "run_id": row[0],
            "goal": row[1],
            "accepted_score": row[2],
            "completed_tasks": row[3],
            "observations": row[4],
            "artifact_path": row[5],
            "created_at": row[6],
        }

    return store._run_pg(operation)


def fetch_tasks(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> dict[str, dict[str, object]]:
    def operation(cur):
        cur.execute(
            """
            select id, kind, input, status, claimed_by, output, score
            from swarm_tasks
            where id = any(%s)
            """,
            (task_ids,),
        )
        return {
            row[0]: {
                "id": row[0],
                "kind": row[1],
                "input": row[2],
                "status": row[3],
                "claimed_by": row[4],
                "output": row[5],
                "score": row[6],
            }
            for row in cur.fetchall()
        }

    return store._run_pg(operation)


if __name__ == "__main__":
    main()