fengqun/examples/run_real_code_task_acceptance.py

from pathlib import Path
from uuid import uuid4
import json
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.core import Agent, SwarmCoordinator, Task
from swarm_minimal.local_env import load_project_env
from swarm_minimal.newapi_agnet import (
    NewApiAgnet,
    NewApiChannelConfig,
    discover_newapi_models,
    select_distinct_models,
)


CODE_TASK = {
    "title": "真实代码场景：给 swarm-minimal 做可交付的长任务改造方案",
    "description": (
        "基于当前 swarm-minimal Python 原型，设计一个代码级补丁方案："
        "让 NewAPI 三模型长任务验收更适合持续运行。方案必须覆盖模型自动发现、"
        "长任务超时配置、失败模型重试或降级、Azure 持久化验收，以及测试命令。"
        "输出控制在 700 字以内，并必须逐字包含指定测试命令。"
    ),
    "target_files": [
        "swarm_minimal/newapi_agnet.py",
        "swarm_minimal/azure_store.py",
        "examples/run_long_task_acceptance.py",
        "tests/test_newapi_agnet.py",
        "tests/test_minimal_swarm.py",
        "README.md",
    ],
}


ACCEPTANCE_CRITERIA = [
    "3 个 Agnet 必须来自模型发现结果，且模型互不相同。",
    "3 个模型任务都必须完成，并写入 PostgreSQL task pool。",
    "每个任务必须在 PostgreSQL 和 Redis pheromone score 中都有正分。",
    "共享状态必须收敛为 converged。",
    "收敛结果必须写入 PostgreSQL，并存在 Blob artifact。",
    "Redis Stream 必须新增事件，证明任务事件流被写入。",
    "被接受的模型输出必须引用至少 3 个真实目标文件。",
    "被接受的模型输出必须包含可执行测试命令。",
    "被接受的模型输出必须明确不能写死 NEWAPI_MODEL，要使用模型发现。",
    "被接受的模型输出不能建议 NATS 或 Cosmos 作为 MVP 必需依赖。",
]


def main() -> None:
    load_project_env(ROOT)
    azure_config = SwarmConfig.from_env()
    newapi_config = NewApiChannelConfig.from_env()
    if newapi_config.timeout_seconds < 120:
        newapi_config = NewApiChannelConfig(
            base_url=newapi_config.base_url,
            api_key=newapi_config.api_key,
            model=newapi_config.model,
            timeout_seconds=120,
        )

    store = PostgresRedisBlobSwarmStore(azure_config)
    try:
        store.ensure_schema()
        models = select_distinct_models(discover_newapi_models(newapi_config), count=3)
        run_id = uuid4().hex
        stream_before = store.redis.xlen("swarm:events")
        code_context = build_code_context()
        prompt = build_task_prompt(code_context)

        store.shared_state[f"run:{run_id}:goal"] = CODE_TASK["title"]
        store.shared_state[f"run:{run_id}:status"] = "running"

        agents: list[Agent] = []
        task_ids: list[str] = []
        for index, model in enumerate(models):
            capability = f"real_code_model_{index + 1}"
            task = Task(kind=capability, input=prompt)
            store.add_task(task)
            task_ids.append(task.id)
            agents.append(build_code_task_agent(newapi_config, model, index + 1, capability))

        result = SwarmCoordinator(store=store, agents=agents).run_until_converged(run_id)
        report = collect_report(
            store=store,
            run_id=run_id,
            result=result,
            selected_models=models,
            task_ids=task_ids,
            stream_before=stream_before,
        )
        print(json.dumps(report, ensure_ascii=False, indent=2))
        if report["summary"]["status"] != "PASS":
            raise SystemExit(1)
    finally:
        store.close()


def build_code_context() -> str:
    return """
--- swarm_minimal/newapi_agnet.py ---
Key surfaces:
- NewApiChannelConfig.from_env reads NEWAPI_BASE_URL / NEWAPI_API_KEY and optional NEWAPI_MODEL.
- discover_newapi_models(config) tries /v1/models, /models, /model.
- select_distinct_models(models, count=3) chooses unique model ids.
- build_model_test_agnets(config, models) creates one Agent per model.
- NewApiAgnet.chat posts to /v1/chat/completions and requires config.model.

--- swarm_minimal/azure_store.py ---
Key surfaces:
- PostgresRedisBlobSwarmStore persists task pool, pheromone scores, shared state, outbox, convergence.
- add_task writes swarm_tasks and swarm_pheromones, then emits task.created.
- complete_task writes task output, PostgreSQL score, Redis sorted-set score, observation, task.done.
- converge writes swarm_convergence and Blob artifact.

--- examples/run_long_task_acceptance.py ---
Key surfaces:
- Loads .env.
- Uses SwarmConfig and NewApiChannelConfig.
- Discovers models and selects 3 distinct models.
- Runs long task with Azure-backed store.
- Verifies model discovery, task pool, scores, shared state, convergence, Redis stream, no fixed model.

--- tests/test_newapi_agnet.py ---
Key surfaces:
- Mock HTTP client validates model discovery and chat payloads.
- Tests three NewAPI Agnets use three different models.

--- tests/test_minimal_swarm.py ---
Key surfaces:
- Validates minimal swarm resources and env parsing/redaction.
""".strip()


def build_task_prompt(code_context: str) -> str:
    return (
        f"{CODE_TASK['title']}\n\n"
        f"任务说明：{CODE_TASK['description']}\n\n"
        "目标文件：\n"
        + "\n".join(f"- {item}" for item in CODE_TASK["target_files"])
        + "\n\n验收标准：\n"
        + "\n".join(f"- {item}" for item in ACCEPTANCE_CRITERIA)
        + "\n\n当前代码上下文：\n"
        + code_context
        + "\n\n输出要求：\n"
        "- 用中文输出。\n"
        "- 控制在 700 字以内。\n"
        "- 给出补丁计划，必须引用具体文件路径。\n"
        "- 必须逐字包含以下测试命令：\n"
        "  ./.venv/bin/python -B -m unittest discover -s tests\n"
        "  ./.venv/bin/python -B examples/run_long_task_acceptance.py\n"
        "- 明确说明模型必须自动发现，不允许写死 NEWAPI_MODEL。\n"
        "- 如果提到 NATS 或 Cosmos，只能用于说明“不引入/不使用”。\n"
        "- 不要包含任何真实密钥。\n"
    )


def build_code_task_agent(config: NewApiChannelConfig, model: str, index: int, capability: str) -> Agent:
    model_config = NewApiChannelConfig(
        base_url=config.base_url,
        api_key=config.api_key,
        model=model,
        timeout_seconds=config.timeout_seconds,
    )
    agnet = NewApiAgnet(model_config, agent_id=f"real-code-agnet-{index}", capability=capability)

    def run(task: Task, shared_state: dict[str, str]) -> tuple[str, float]:
        content = agnet.chat(
            system_prompt=(
                "You are a senior coding agent participating in a swarm acceptance test. "
                "Analyze the provided real code context and propose a concrete patch plan. "
                "Do not include secrets."
            ),
            user_prompt=(
                f"Task kind: {task.kind}\n"
                f"Task input:\n{task.input}\n\n"
                f"Shared state keys: {', '.join(sorted(shared_state.keys()))}"
            ),
        )
        return content, score_code_task_output(content)

    return Agent(id=f"real-code-agnet-{index}", capability=capability, run=run)


def score_code_task_output(content: str) -> float:
    lowered = content.lower()
    checks = [
        "swarm_minimal/newapi_agnet.py" in content,
        "examples/run_long_task_acceptance.py" in content,
        "tests/test_newapi_agnet.py" in content or "tests/test_minimal_swarm.py" in content,
        "python" in lowered and "unittest" in lowered,
        "discover" in lowered or "模型发现" in content,
        "NEWAPI_MODEL" in content,
        "timeout" in lowered or "超时" in content,
        "retry" in lowered or "重试" in content or "降级" in content,
    ]
    return min(0.97, 0.49 + sum(1 for item in checks if item) * 0.06)


def collect_report(
    *,
    store: PostgresRedisBlobSwarmStore,
    run_id: str,
    result,
    selected_models: list[str],
    task_ids: list[str],
    stream_before: int,
) -> dict[str, object]:
    task_rows = fetch_task_rows(store, task_ids)
    pg_scores = fetch_pg_scores(store, task_ids)
    redis_scores = {task_id: store.redis.zscore("swarm:pheromones", task_id) for task_id in task_ids}
    stream_after = store.redis.xlen("swarm:events")
    shared_state = fetch_shared_state(store, f"run:{run_id}:status")
    convergence = fetch_convergence(store, run_id)
    artifact_path = convergence["artifact_path"] if convergence else ""
    blob_exists = bool(artifact_path and store.container.get_blob_client(artifact_path).exists())
    accepted = result.accepted_output

    checks = [
        {
            "name": "model_discovery_three_distinct_models",
            "passed": len(set(selected_models)) == 3,
            "evidence": selected_models,
        },
        {
            "name": "task_pool_pg_done",
            "passed": len(task_rows) == 3 and all(row["status"] == "done" for row in task_rows),
            "evidence": task_rows,
        },
        {
            "name": "pheromone_scores_pg_and_redis",
            "passed": all(pg_scores.get(task_id, 0) > 0 for task_id in task_ids)
            and all((redis_scores.get(task_id) or 0) > 0 for task_id in task_ids),
            "evidence": {"pg_scores": pg_scores, "redis_scores": redis_scores},
        },
        {
            "name": "shared_state_converged",
            "passed": shared_state == "converged",
            "evidence": shared_state,
        },
        {
            "name": "result_convergence_pg_and_blob",
            "passed": bool(convergence)
            and convergence["completed_tasks"] == 3
            and convergence["accepted_score"] >= 0.75
            and blob_exists,
            "evidence": {"convergence": convergence, "blob_exists": blob_exists},
        },
        {
            "name": "redis_stream_events",
            "passed": stream_after - stream_before >= 10,
            "evidence": {"before": stream_before, "after": stream_after, "delta": stream_after - stream_before},
        },
        {
            "name": "accepted_output_references_real_files",
            "passed": count_referenced_target_files(accepted) >= 3,
            "evidence": referenced_target_files(accepted),
        },
        {
            "name": "accepted_output_includes_test_commands",
            "passed": "./.venv/bin/python -B -m unittest discover -s tests" in accepted
            and "./.venv/bin/python -B examples/run_long_task_acceptance.py" in accepted,
            "evidence": "must include exact unittest and long-task commands",
        },
        {
            "name": "accepted_output_keeps_model_discovery",
            "passed": ("模型发现" in accepted or "discover" in accepted.lower())
            and "NEWAPI_MODEL" in accepted
            and ("不" in accepted or "not" in accepted.lower()),
            "evidence": "must explicitly reject fixed NEWAPI_MODEL for three-model flow",
        },
        {
            "name": "accepted_output_no_nats_or_cosmos_requirement",
            "passed": no_required_nats_or_cosmos(accepted),
            "evidence": "MVP must remain PostgreSQL/Redis/Blob based",
        },
    ]
    status = "PASS" if all(check["passed"] for check in checks) else "FAIL"
    return {
        "task": CODE_TASK,
        "acceptance_criteria": ACCEPTANCE_CRITERIA,
        "summary": {
            "status": status,
            "run_id": run_id,
            "accepted_score": result.accepted_score,
            "completed_tasks": result.completed_tasks,
            "accepted_task_id": result.accepted_task_id,
            "artifact_path": artifact_path,
        },
        "selected_models": selected_models,
        "checks": checks,
        "accepted_output_preview": accepted[:900],
    }


def referenced_target_files(text: str) -> list[str]:
    return [item for item in CODE_TASK["target_files"] if item in text]


def count_referenced_target_files(text: str) -> int:
    return len(referenced_target_files(text))


def no_required_nats_or_cosmos(text: str) -> bool:
    lowered = text.lower()
    if "nats" not in lowered and "cosmos" not in lowered:
        return True
    negative_markers = [
        "不引入",
        "不使用",
        "无需",
        "不要",
        "no ",
        "without",
        "not use",
    ]
    for name in ["nats", "cosmos"]:
        if name not in lowered:
            continue
        if not any(marker in lowered[max(0, lowered.find(name) - 24) : lowered.find(name) + 24] for marker in negative_markers):
            return False
    return True


def fetch_task_rows(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> list[dict[str, object]]:
    with store.pg.cursor() as cur:
        cur.execute(
            "select id, kind, status, claimed_by, score from swarm_tasks where id = any(%s) order by kind",
            (task_ids,),
        )
        return [
            {"id": row[0], "kind": row[1], "status": row[2], "claimed_by": row[3], "score": row[4]}
            for row in cur.fetchall()
        ]


def fetch_pg_scores(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> dict[str, float]:
    with store.pg.cursor() as cur:
        cur.execute("select task_id, score from swarm_pheromones where task_id = any(%s)", (task_ids,))
        return {row[0]: row[1] for row in cur.fetchall()}


def fetch_shared_state(store: PostgresRedisBlobSwarmStore, key: str) -> str | None:
    with store.pg.cursor() as cur:
        cur.execute("select value from swarm_shared_state where key = %s", (key,))
        row = cur.fetchone()
        return row[0] if row else None


def fetch_convergence(store: PostgresRedisBlobSwarmStore, run_id: str) -> dict[str, object] | None:
    with store.pg.cursor() as cur:
        cur.execute(
            """
            select run_id, completed_tasks, accepted_score, accepted_task_id, artifact_path
            from swarm_convergence
            where run_id = %s
            """,
            (run_id,),
        )
        row = cur.fetchone()
        if not row:
            return None
        return {
            "run_id": row[0],
            "completed_tasks": row[1],
            "accepted_score": row[2],
            "accepted_task_id": row[3],
            "artifact_path": row[4],
        }


if __name__ == "__main__":
    main()