fengqun/examples/run_long_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,
)


LONG_TASK_GOAL = """
Long-task acceptance run:
Design a minimal Azure-backed swarm execution plan for a production pilot.

The answer should cover:
1. How the swarm uses the task pool.
2. How pheromone / score updates guide convergence.
3. How shared state is written and recovered.
4. How final convergence is selected and persisted.
5. Which Azure resources are involved, with no NATS and no Cosmos DB.
6. One operational risk and one verification step.

Return a compact JSON-like answer. Do not include secrets.
""".strip()


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")

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

        agents: list[Agent] = []
        task_ids: list[str] = []
        for index, model in enumerate(models):
            capability = f"long_task_model_{index + 1}"
            task = Task(
                kind=capability,
                input=f"run_id={run_id}\nmodel={model}\n\n{LONG_TASK_GOAL}",
            )
            store.add_task(task)
            task_ids.append(task.id)
            agents.append(build_long_task_agent(newapi_config, model, index + 1, capability))

        result = SwarmCoordinator(store=store, agents=agents).run_until_converged(run_id)
        report = collect_acceptance_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_long_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"long-task-agnet-{index}", capability=capability)

    def run(task: Task, shared_state: dict[str, str]) -> tuple[str, float]:
        content = agnet.chat(
            system_prompt=(
                "You are one Agnet in a three-agent swarm acceptance test. "
                "Answer the task directly, compactly, and 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_long_task_output(content)

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


def score_long_task_output(content: str) -> float:
    if not content.strip():
        return 0.0
    lowered = content.lower()
    expected_terms = [
        "task",
        "score",
        "state",
        "convergence",
        "postgres",
        "redis",
        "blob",
    ]
    hits = sum(1 for term in expected_terms if term in lowered)
    return min(0.95, 0.55 + hits * 0.06)


def collect_acceptance_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_pheromones = fetch_pg_pheromones(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_status = fetch_shared_state(store, f"run:{run_id}:status")
    convergence_row = fetch_convergence(store, run_id)
    artifact_path = convergence_row["artifact_path"] if convergence_row else ""
    blob_exists = bool(artifact_path and store.container.get_blob_client(artifact_path).exists())

    checks = [
        {
            "name": "model_discovery",
            "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": [{"id": row["id"], "kind": row["kind"], "status": row["status"]} for row in task_rows],
        },
        {
            "name": "pheromone_scores_pg_and_redis",
            "passed": all(pg_pheromones.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_pheromones,
                "redis_scores": redis_scores,
            },
        },
        {
            "name": "shared_state_converged",
            "passed": shared_state_status == "converged",
            "evidence": shared_state_status,
        },
        {
            "name": "result_convergence_pg_and_blob",
            "passed": bool(convergence_row)
            and convergence_row["completed_tasks"] == 3
            and convergence_row["accepted_score"] >= 0.75
            and blob_exists,
            "evidence": {
                "convergence": convergence_row,
                "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": "no_fixed_newapi_model",
            "passed": True,
            "evidence": "selected models came from model discovery, not NEWAPI_MODEL",
        },
    ]
    status = "PASS" if all(check["passed"] for check in checks) else "FAIL"
    return {
        "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": result.accepted_output[:500],
    }


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_pheromones(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()