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


TARGET_FILES = [
    "swarm_minimal/core.py",
    "swarm_minimal/newapi_agnet.py",
    "swarm_minimal/azure_store.py",
    "examples/run_full_live_test.py",
    "examples/run_long_task_acceptance.py",
    "tests/test_newapi_agnet.py",
    "tests/test_minimal_swarm.py",
    "README.md",
]


MULTITASK_SCENARIO = {
    "title": "真实全面场景：把 swarm-minimal 演进成可处理大规模代码任务的生产级多 Agnet 调度原型",
    "description": (
        "当前代码已有 PostgreSQL/Redis/Blob/NewAPI 三模型最小闭环。"
        "现在要评估它能否支撑复杂真实场景：一个大规模代码库改造任务被拆成多个子任务，"
        "多个 Agnet 长时间思考后分别给出架构、算法、数据一致性、失败恢复、补丁计划和验收方案。"
        "输出必须面向真实代码改造，而不是泛泛测试。"
    ),
}


SUBTASKS = [
    {
        "capability": "complex_architecture",
        "title": "架构拆分",
        "ask": "设计模块边界和执行流，说明现有哪些文件要改，如何支持多任务拆分和多模型 Agnet。",
    },
    {
        "capability": "complex_algorithm",
        "title": "复杂调度算法",
        "ask": "设计支持上千代码任务的调度/抢占/信息素评分/收敛算法，给出关键数据结构和复杂度。",
    },
    {
        "capability": "data_consistency",
        "title": "数据一致性",
        "ask": "设计 PostgreSQL、Redis、Blob 之间的一致性、幂等、outbox、artifact 写入和恢复策略。",
    },
    {
        "capability": "failure_recovery",
        "title": "失败恢复",
        "ask": "设计模型超时、部分失败、重试、降级、死信、租约过期和重复执行的处理方式。",
    },
    {
        "capability": "code_patch_plan",
        "title": "代码补丁计划",
        "ask": "给出具体文件级补丁计划，必须引用目标文件，包含测试文件如何补。",
    },
    {
        "capability": "acceptance_plan",
        "title": "验收方案",
        "ask": "给出真实验收标准、命令、指标、失败判定和规模化压测方式。",
    },
]


ACCEPTANCE_CRITERIA = [
    "必须自动发现至少 3 个模型，并使用 3 个互不相同的模型分担子任务。",
    "必须拆出 6 个真实工程子任务，并全部完成。",
    "所有子任务必须写入 PostgreSQL task pool，状态为 done。",
    "所有子任务必须在 PostgreSQL 和 Redis pheromone score 中有正分。",
    "共享状态必须收敛为 converged。",
    "最终收敛结果必须写入 PostgreSQL，并存在 Blob artifact。",
    "Redis Stream 必须新增至少 3*N+1 条事件。",
    "所有模型输出合并后必须引用至少 5 个真实目标文件。",
    "输出必须包含复杂算法/数据结构/复杂度说明。",
    "输出必须包含大规模代码场景要素，例如上千任务、并发、租约、重试或队列。",
    "输出必须包含可执行测试或验收命令。",
    "输出必须明确三模型流程不能写死 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 < 360:
        newapi_config = NewApiChannelConfig(
            base_url=newapi_config.base_url,
            api_key=newapi_config.api_key,
            model=newapi_config.model,
            timeout_seconds=360,
        )

    store = PostgresRedisBlobSwarmStore(azure_config)
    try:
        store.ensure_schema()
        discovered_models = discover_newapi_models(newapi_config)
        models = select_models_for_complex_acceptance(newapi_config, discovered_models, count=3)
        run_id = uuid4().hex
        stream_before = store._run_redis(lambda redis: redis.xlen("swarm:events"))
        code_context = build_code_context()

        store.shared_state[f"run:{run_id}:goal"] = MULTITASK_SCENARIO["title"]
        store.shared_state[f"run:{run_id}:status"] = "running"
        store.shared_state[f"run:{run_id}:subtask_count"] = str(len(SUBTASKS))

        agents: list[Agent] = []
        task_ids: list[str] = []
        task_models: dict[str, str] = {}
        for index, subtask in enumerate(SUBTASKS):
            model = models[index % len(models)]
            task = Task(
                kind=subtask["capability"],
                input=build_subtask_prompt(subtask, code_context, model),
            )
            store.add_task(task)
            task_ids.append(task.id)
            task_models[task.id] = model
            fallback_models = [candidate for candidate in models if candidate != model]
            agents.append(build_subtask_agent(newapi_config, model, fallback_models, index + 1, subtask["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,
            discovered_models=discovered_models,
            task_ids=task_ids,
            task_models=task_models,
            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 select_models_for_complex_acceptance(
    config: NewApiChannelConfig,
    discovered_models: list[str],
    *,
    count: int,
) -> list[str]:
    """Choose responsive models from discovery without relying on NEWAPI_MODEL."""

    candidates = select_distinct_models(sorted(discovered_models, key=model_priority), count=len(set(discovered_models)))
    responsive: list[str] = []
    for model in candidates:
        probe_config = NewApiChannelConfig(
            base_url=config.base_url,
            api_key=config.api_key,
            model=model,
            timeout_seconds=min(config.timeout_seconds, 30),
        )
        try:
            NewApiAgnet(probe_config, agent_id="probe", capability="probe").chat(
                system_prompt="Return only OK.",
                user_prompt="Health probe for model selection. Return OK.",
            )
        except Exception:
            continue
        responsive.append(model)
        if len(responsive) == count:
            return responsive
    if len(responsive) >= count:
        return responsive[:count]
    return select_distinct_models(candidates, count=count)


def model_priority(model: str) -> tuple[int, str]:
    lowered = model.lower()
    if "flash" in lowered:
        return (0, model)
    if "haiku" in lowered:
        return (1, model)
    if "sonnet" in lowered:
        return (2, model)
    if "mini" in lowered:
        return (3, model)
    if "pro" in lowered:
        return (8, model)
    if "opus" in lowered:
        return (9, model)
    return (5, model)


def build_code_context() -> str:
    return """
Current code map:
- swarm_minimal/core.py: Task, Agent, InMemorySwarmStore, SwarmCoordinator, sequential task claiming, score-based convergence.
- swarm_minimal/newapi_agnet.py: NewApiChannelConfig, model discovery through /v1/models /models /model, model selection, chat calls.
- swarm_minimal/azure_store.py: PostgreSQL task pool, pheromone table, shared state, observations, convergence, outbox; Redis Stream and sorted-set; Blob artifact.
- examples/run_full_live_test.py: Azure-backed three-model live test.
- examples/run_long_task_acceptance.py: Long task acceptance, 120s timeout, checks model discovery/task pool/scores/shared state/convergence/Redis stream.
- tests/test_newapi_agnet.py: Mock HTTP tests for model discovery and three model Agnets.
- tests/test_minimal_swarm.py: Minimal swarm, Azure resource, env parsing and redaction tests.
- README.md: User-facing runbook and Azure resource mapping.
""".strip()


def build_subtask_prompt(subtask: dict[str, str], code_context: str, model: str) -> str:
    return (
        f"{MULTITASK_SCENARIO['title']}\n\n"
        f"总任务：{MULTITASK_SCENARIO['description']}\n\n"
        f"当前子任务：{subtask['title']}\n"
        f"子任务要求：{subtask['ask']}\n"
        f"当前模型：{model}\n\n"
        "目标文件：\n"
        + "\n".join(f"- {item}" for item in TARGET_FILES)
        + "\n\n总体验收标准：\n"
        + "\n".join(f"- {item}" for item in ACCEPTANCE_CRITERIA)
        + "\n\n代码上下文：\n"
        + code_context
        + "\n\n输出要求：\n"
        "- 中文，控制在 900 字以内。\n"
        "- 必须引用相关真实文件路径。\n"
        "- 必须给出工程化细节，不要泛泛而谈。\n"
        "- 如果是算法子任务，必须写复杂度或数据结构。\n"
        "- 必须包含至少一个可执行命令，例如 ./.venv/bin/python -B -m unittest discover -s tests。\n"
        "- 必须说明三模型流程依赖模型发现，不允许写死 NEWAPI_MODEL。\n"
        "- 不要把 NATS 或 Cosmos 作为 MVP 必需依赖。\n"
        "- 不要包含任何真实密钥。\n"
    )


def build_subtask_agent(
    config: NewApiChannelConfig,
    model: str,
    fallback_models: list[str],
    index: int,
    capability: str,
) -> Agent:
    def run(task: Task, shared_state: dict[str, str]) -> tuple[str, float]:
        content = chat_with_model_fallback(config, model, fallback_models, index, capability, task, shared_state)
        return content, score_output(content, capability)

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


def chat_with_model_fallback(
    config: NewApiChannelConfig,
    primary_model: str,
    fallback_models: list[str],
    index: int,
    capability: str,
    task: Task,
    shared_state: dict[str, str],
) -> str:
    errors: list[str] = []
    for candidate in [primary_model, *fallback_models]:
        model_config = NewApiChannelConfig(
            base_url=config.base_url,
            api_key=config.api_key,
            model=candidate,
            timeout_seconds=config.timeout_seconds,
        )
        agnet = NewApiAgnet(model_config, agent_id=f"complex-agnet-{index}", capability=capability)
        try:
            content = agnet.chat(
                system_prompt=(
                    "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."
                ),
                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()))}"
                ),
            )
        except Exception as exc:
            errors.append(f"{candidate}:{exc.__class__.__name__}")
            continue
        prefix = f"primary_model={primary_model}; used_model={candidate}; model_selection=discovered_models_not_NEWAPI_MODEL"
        if errors:
            prefix += "; fallback_after=" + ",".join(errors)
        return prefix + "\n" + content
    raise RuntimeError("all model attempts failed: " + ",".join(errors))


def score_output(content: str, capability: str) -> float:
    lowered = content.lower()
    checks = [
        count_referenced_files(content) >= 2,
        "NEWAPI_MODEL" in content,
        "discover" in lowered or "模型发现" in content,
        "python" in lowered and "unittest" in lowered,
        "postgres" in lowered or "postgresql" in lowered,
        "redis" in lowered,
        "blob" in lowered,
        "retry" in lowered or "重试" in content or "降级" in content,
    ]
    if capability == "complex_algorithm":
        checks.extend(["o(" in lowered or "复杂度" in content, "queue" in lowered or "heap" in lowered or "队列" in content])
    if capability == "acceptance_plan":
        checks.extend(["pass" in lowered or "验收" in content, "压测" in content or "load" in lowered])
    return min(0.98, 0.44 + sum(1 for item in checks if item) * 0.06)


def collect_report(
    *,
    store: PostgresRedisBlobSwarmStore,
    run_id: str,
    result,
    selected_models: list[str],
    discovered_models: list[str],
    task_ids: list[str],
    task_models: dict[str, str],
    stream_before: int,
) -> dict[str, object]:
    task_rows = fetch_task_rows(store, task_ids)
    outputs = {row["id"]: row["output"] or "" for row in task_rows}
    merged_output = "\n\n".join(outputs.values())
    pg_scores = fetch_pg_scores(store, task_ids)
    redis_scores = {
        task_id: store._run_redis(lambda redis, current_task_id=task_id: redis.zscore("swarm:pheromones", current_task_id))
        for task_id in task_ids
    }
    stream_after = store._run_redis(lambda redis: 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())
    expected_event_delta = 3 * len(task_ids) + 1

    checks = [
        {
            "name": "three_distinct_models_from_discovery",
            "passed": len(discovered_models) >= 3
            and len(set(selected_models)) == 3
            and set(selected_models).issubset(set(discovered_models)),
            "evidence": {"discovered": discovered_models, "selected": selected_models},
        },
        {
            "name": "six_subtasks_all_done_in_pg",
            "passed": len(task_rows) == 6 and all(row["status"] == "done" for row in task_rows),
            "evidence": compact_task_rows(task_rows, task_models),
        },
        {
            "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": "convergence_pg_and_blob",
            "passed": bool(convergence)
            and convergence["completed_tasks"] == 6
            and convergence["accepted_score"] >= 0.75
            and blob_exists,
            "evidence": {"convergence": convergence, "blob_exists": blob_exists},
        },
        {
            "name": "redis_stream_event_volume",
            "passed": stream_after - stream_before >= expected_event_delta,
            "evidence": {"before": stream_before, "after": stream_after, "delta": stream_after - stream_before, "expected_min": expected_event_delta},
        },
        {
            "name": "references_real_code_files",
            "passed": count_referenced_files(merged_output) >= 5,
            "evidence": referenced_files(merged_output),
        },
        {
            "name": "contains_complex_algorithm_discussion",
            "passed": ("o(" in merged_output.lower() or "复杂度" in merged_output)
            and any(term in merged_output.lower() for term in ["queue", "heap", "priority", "队列", "优先级"]),
            "evidence": "requires complexity and data structure discussion",
        },
        {
            "name": "contains_large_scale_operational_scenario",
            "passed": any(term in merged_output for term in ["上千", "1000", "千级", "大规模"])
            and any(term in merged_output for term in ["并发", "租约", "重试", "队列"]),
            "evidence": "requires large-scale and operational terms",
        },
        {
            "name": "contains_test_or_acceptance_commands",
            "passed": "./.venv/bin/python -B -m unittest discover -s tests" in merged_output
            or "python -B -m unittest discover -s tests" in merged_output,
            "evidence": "requires executable unittest command",
        },
        {
            "name": "keeps_model_discovery_not_fixed_model",
            "passed": ("NEWAPI_MODEL" in merged_output)
            and ("模型发现" in merged_output or "discover" in merged_output.lower())
            and ("不" in merged_output or "not" in merged_output.lower()),
            "evidence": "must reject fixed NEWAPI_MODEL",
        },
        {
            "name": "no_required_nats_or_cosmos",
            "passed": no_required_nats_or_cosmos(merged_output),
            "evidence": "NATS/Cosmos may only appear as rejected dependencies",
        },
    ]
    status = "PASS" if all(check["passed"] for check in checks) else "FAIL"
    return {
        "task": MULTITASK_SCENARIO,
        "subtasks": SUBTASKS,
        "acceptance_criteria": ACCEPTANCE_CRITERIA,
        "summary": {
            "status": status,
            "run_id": run_id,
            "completed_tasks": result.completed_tasks,
            "accepted_score": result.accepted_score,
            "accepted_task_id": result.accepted_task_id,
            "artifact_path": artifact_path,
        },
        "discovered_models": discovered_models,
        "selected_models": selected_models,
        "checks": checks,
        "accepted_output_preview": result.accepted_output[:1200],
    }


def compact_task_rows(task_rows: list[dict[str, object]], task_models: dict[str, str]) -> list[dict[str, object]]:
    return [
        {
            "id": row["id"],
            "kind": row["kind"],
            "status": row["status"],
            "claimed_by": row["claimed_by"],
            "score": row["score"],
            "model": task_models.get(row["id"], "<unknown>"),
        }
        for row in task_rows
    ]


def referenced_files(text: str) -> list[str]:
    return [item for item in TARGET_FILES if item in text]


def count_referenced_files(text: str) -> int:
    return len(referenced_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",
        "reject",
    ]
    for name in ["nats", "cosmos"]:
        position = lowered.find(name)
        if position == -1:
            continue
        nearby = lowered[max(0, position - 64) : position + 64]
        if not any(marker in nearby for marker in negative_markers):
            return False
    return True


def fetch_task_rows(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> list[dict[str, object]]:
    def operation(cur) -> list[dict[str, object]]:
        cur.execute(
            """
            select id, kind, status, claimed_by, score, output
            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],
                "output": row[5],
            }
            for row in cur.fetchall()
        ]

    return store._run_pg(operation)


def fetch_pg_scores(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> dict[str, float]:
    def operation(cur) -> dict[str, float]:
        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()}

    return store._run_pg(operation)


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

    return store._run_pg(operation)


def fetch_convergence(store: PostgresRedisBlobSwarmStore, run_id: str) -> dict[str, object] | None:
    def operation(cur) -> dict[str, object] | None:
        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],
        }

    return store._run_pg(operation)


if __name__ == "__main__":
    main()