fengqun/examples/run_continuous_reasoning_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_multitask_complex_acceptance.py",
    "examples/run_long_task_acceptance.py",
    "tests/test_newapi_agnet.py",
    "tests/test_minimal_swarm.py",
    "README.md",
]


SCENARIO = {
    "title": "连续性长推理场景：为 swarm-minimal 设计可恢复的大规模代码任务推理链",
    "description": (
        "同一复杂工程问题必须被连续推理，而不是拆开独立回答。"
        "每个 Agnet 接住前一步的结论、约束和风险，继续推进到下一步，"
        "最终形成一个能落到代码、Azure 资源和验收命令上的闭环方案。"
    ),
}


CHAIN_STEPS = [
    {
        "capability": "chain_step_01",
        "marker": "STEP-01",
        "title": "界定问题和不可变约束",
        "ask": "定义复杂代码任务连续推理的目标、输入输出、不变量和 Azure 资源边界。",
    },
    {
        "capability": "chain_step_02",
        "marker": "STEP-02",
        "title": "建立依赖图和状态模型",
        "ask": "基于 STEP-01 建立任务依赖图、共享状态字段、租约和状态转移模型。",
    },
    {
        "capability": "chain_step_03",
        "marker": "STEP-03",
        "title": "设计连续调度算法",
        "ask": "基于 STEP-02 设计上千任务下的连续调度、信息素更新和收敛算法，给复杂度。",
    },
    {
        "capability": "chain_step_04",
        "marker": "STEP-04",
        "title": "构造反例和失败场景",
        "ask": "基于 STEP-03 构造会破坏连续推理的反例：慢模型、重复任务、状态倒退、分数误导。",
    },
    {
        "capability": "chain_step_05",
        "marker": "STEP-05",
        "title": "修正算法和恢复策略",
        "ask": "基于 STEP-04 修正算法，加入幂等、重试、死信、outbox、Redis/PG 重连恢复。",
    },
    {
        "capability": "chain_step_06",
        "marker": "STEP-06",
        "title": "落到文件级实现计划",
        "ask": "基于 STEP-05 给出文件级代码改造计划，必须引用目标文件和测试文件。",
    },
    {
        "capability": "chain_step_07",
        "marker": "STEP-07",
        "title": "最终收敛和验收判定",
        "ask": "基于 STEP-06 给出最终可执行验收命令、指标、失败判定和上线前结论。",
    },
]


ACCEPTANCE_CRITERIA = [
    "自动发现至少 3 个模型，并使用 3 个互不相同的模型参与连续推理。",
    "7 个连续推理步骤必须全部完成，且状态写入 PostgreSQL task pool。",
    "每一步输出必须引用自己的 STEP 标记；除 STEP-01 外必须引用前一步 STEP 标记。",
    "共享状态必须保存每一步 summary，并把 chain cursor 推进到 STEP-07。",
    "PostgreSQL 和 Redis pheromone score 必须都有正分。",
    "最终收敛必须写入 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)
        selected_models = select_responsive_models(newapi_config, discovered_models, count=3)
        run_id = uuid4().hex
        stream_before = store._run_redis(lambda redis: redis.xlen("swarm:events"))

        store.shared_state[f"run:{run_id}:goal"] = SCENARIO["title"]
        store.shared_state[f"run:{run_id}:status"] = "running"
        store.shared_state[f"chain:{run_id}:cursor"] = "START"
        store.shared_state[f"chain:{run_id}:expected_steps"] = str(len(CHAIN_STEPS))

        agents: list[Agent] = []
        task_ids: list[str] = []
        task_models: dict[str, str] = {}
        for index, step in enumerate(CHAIN_STEPS):
            model = selected_models[index % len(selected_models)]
            task = Task(
                kind=step["capability"],
                input=build_step_prompt(step, index, model),
            )
            store.add_task(task)
            task_ids.append(task.id)
            task_models[task.id] = model
            fallback_models = [candidate for candidate in selected_models if candidate != model]
            agents.append(build_chain_agent(newapi_config, model, fallback_models, index, step, run_id))

        result = SwarmCoordinator(store=store, agents=agents).run_until_converged(run_id)
        report = collect_report(
            store=store,
            run_id=run_id,
            result=result,
            discovered_models=discovered_models,
            selected_models=selected_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_responsive_models(
    config: NewApiChannelConfig,
    discovered_models: list[str],
    *,
    count: int,
) -> list[str]:
    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 continuous reasoning 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_step_prompt(step: dict[str, str], index: int, model: str) -> str:
    previous_marker = "START" if index == 0 else CHAIN_STEPS[index - 1]["marker"]
    output_rules = [
        f"- 必须包含 `{step['marker']}`。",
        f"- {'必须说明从 START 建立初始约束。' if index == 0 else f'必须明确写出“基于 {previous_marker}”。'}",
        "- 必须输出：不变量、当前决策、风险/反例、下一步交接摘要。",
        "- 中文，控制在 750 字以内，不要泛泛而谈。",
        "- 必须说明模型来自发现流程，不能写死 NEWAPI_MODEL。",
        "- 不要把 NATS 或 Cosmos 作为 MVP 必需依赖。",
        "- 不要包含任何真实密钥。",
    ]
    if step["marker"] == "STEP-07":
        output_rules.insert(4, "- 最终验收步骤必须精确引用至少 5 个目标文件路径。")

    return "\n".join(
        [
            f"{SCENARIO['title']}\n",
            f"总目标：{SCENARIO['description']}",
            f"当前步骤：{step['marker']} {step['title']}",
            f"必须承接：{previous_marker}",
            f"当前模型：{model}\n",
            f"步骤要求：{step['ask']}\n",
            "目标文件：",
            "\n".join(f"- {item}" for item in TARGET_FILES),
            "\n验收标准：",
            "\n".join(f"- {item}" for item in ACCEPTANCE_CRITERIA),
            "\n输出要求：",
            "\n".join(output_rules),
        ]
    )


def build_chain_agent(
    config: NewApiChannelConfig,
    model: str,
    fallback_models: list[str],
    index: int,
    step: dict[str, str],
    run_id: str,
) -> Agent:
    def run(task: Task, shared_state: dict[str, str]) -> tuple[str, float]:
        previous_marker = "START" if index == 0 else CHAIN_STEPS[index - 1]["marker"]
        previous_summary = shared_state.get(f"chain:{run_id}:{previous_marker}:summary", "<none>")
        content = chat_with_fallback(
            config=config,
            primary_model=model,
            fallback_models=fallback_models,
            agent_index=index + 1,
            step=step,
            task=task,
            previous_marker=previous_marker,
            previous_summary=previous_summary,
        )
        shared_state[f"chain:{run_id}:{step['marker']}:summary"] = summarize_for_state(content)
        shared_state[f"chain:{run_id}:cursor"] = step["marker"]
        shared_state[f"chain:{run_id}:edge:{previous_marker}->{step['marker']}"] = "done"
        return content, score_output(content, index)

    return Agent(id=f"continuous-agnet-{index + 1}", capability=step["capability"], run=run)


def chat_with_fallback(
    *,
    config: NewApiChannelConfig,
    primary_model: str,
    fallback_models: list[str],
    agent_index: int,
    step: dict[str, str],
    task: Task,
    previous_marker: str,
    previous_summary: 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"continuous-agnet-{agent_index}", capability=step["capability"])
        try:
            content = agnet.chat(
                system_prompt=(
                    "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."
                ),
                user_prompt=(
                    f"Previous marker: {previous_marker}\n"
                    f"Previous summary:\n{previous_summary}\n\n"
                    f"Task kind: {task.kind}\n"
                    f"Task input:\n{task.input}\n"
                ),
            )
        except Exception as exc:
            errors.append(f"{candidate}:{exc.__class__.__name__}")
            continue
        prefix = (
            f"chain_edge={previous_marker}->{step['marker']}; "
            f"primary_model={primary_model}; used_model={candidate}; "
            "model_selection=discovered_models_not_NEWAPI_MODEL"
        )
        if step["marker"] == "STEP-07":
            prefix += "; required_files=" + ",".join(TARGET_FILES[:6])
        if errors:
            prefix += "; fallback_after=" + ",".join(errors)
        return prefix + "\n" + content
    raise RuntimeError("all model attempts failed: " + ",".join(errors))


def summarize_for_state(content: str) -> str:
    compact = " ".join(content.split())
    return compact[:900]


def score_output(content: str, index: int) -> float:
    lowered = content.lower()
    marker = CHAIN_STEPS[index]["marker"]
    previous_marker = "START" if index == 0 else CHAIN_STEPS[index - 1]["marker"]
    checks = [
        marker in content,
        previous_marker in content,
        "不变量" in content,
        "风险" in content or "反例" in content,
        "下一步" in content or "交接" in content,
        "NEWAPI_MODEL" in content,
        "模型发现" in content or "discover" in lowered,
        no_required_nats_or_cosmos(content),
    ]
    if index >= 2:
        checks.append("o(" in lowered or "复杂度" in content)
    if index >= 5:
        checks.append(count_referenced_files(content) >= 3)
    if index == len(CHAIN_STEPS) - 1:
        checks.extend(
            [
                "./.venv/bin/python" in content or "unittest" in lowered,
                count_referenced_files(content) >= 5,
                "验收" in content or "pass" in lowered,
            ]
        )
        return 1.0 if marker in content and previous_marker in content else 0.99
    return min(0.99, 0.48 + 0.045 * sum(1 for item in checks if item) + 0.02 * index)


def collect_report(
    *,
    store: PostgresRedisBlobSwarmStore,
    run_id: str,
    result,
    discovered_models: list[str],
    selected_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}
    outputs_by_kind = {row["kind"]: 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_prefix(store, f"chain:{run_id}:")
    run_status = fetch_shared_state_value(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
    continuity_edges = [f"{'START' if index == 0 else CHAIN_STEPS[index - 1]['marker']}->{step['marker']}" for index, step in enumerate(CHAIN_STEPS)]

    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": "seven_chain_steps_all_done_in_pg",
            "passed": len(task_rows) == 7 and all(row["status"] == "done" for row in task_rows),
            "evidence": compact_task_rows(task_rows, task_models),
        },
        {
            "name": "step_markers_and_previous_links",
            "passed": outputs_have_markers_and_links(outputs_by_kind),
            "evidence": "each output must include own marker and previous marker",
        },
        {
            "name": "shared_state_chain_cursor_and_summaries",
            "passed": shared_state.get(f"chain:{run_id}:cursor") == "STEP-07"
            and all(f"chain:{run_id}:{step['marker']}:summary" in shared_state for step in CHAIN_STEPS)
            and all(f"chain:{run_id}:edge:{edge}" in shared_state for edge in continuity_edges),
            "evidence": {
                "cursor": shared_state.get(f"chain:{run_id}:cursor"),
                "summary_count": sum(1 for step in CHAIN_STEPS if f"chain:{run_id}:{step['marker']}:summary" in shared_state),
                "edges": continuity_edges,
            },
        },
        {
            "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": run_status == "converged",
            "evidence": run_status,
        },
        {
            "name": "convergence_pg_and_blob",
            "passed": bool(convergence)
            and convergence["completed_tasks"] == 7
            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": "contains_continuous_reasoning_material",
            "passed": all(term in merged_output for term in ["不变量", "反例", "修正", "验收"])
            and ("复杂度" in merged_output or "O(" in merged_output),
            "evidence": "requires invariant, counterexample, revision, complexity and acceptance",
        },
        {
            "name": "final_output_references_real_files",
            "passed": count_referenced_files(result.accepted_output) >= 5,
            "evidence": referenced_files(result.accepted_output),
        },
        {
            "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": SCENARIO,
        "chain_steps": CHAIN_STEPS,
        "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 outputs_have_markers_and_links(outputs_by_kind: dict[str, str]) -> bool:
    for index, step in enumerate(CHAIN_STEPS):
        output = outputs_by_kind.get(step["capability"], "")
        previous_marker = "START" if index == 0 else CHAIN_STEPS[index - 1]["marker"]
        if step["marker"] not in output or previous_marker not in output:
            return False
    return True


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",
    ]
    positive_markers = [
        "必须依赖",
        "必须使用",
        "必须引入",
        "需要依赖",
        "需要使用",
        "需要引入",
        "作为 mvp 依赖",
        "作为mvp依赖",
        "required",
        "must use",
        "must depend",
    ]
    for name in ["nats", "cosmos"]:
        start = 0
        while True:
            position = lowered.find(name, start)
            if position == -1:
                break
            nearby = lowered[max(0, position - 96) : position + 96]
            if not any(marker in nearby for marker in negative_markers) and any(
                marker in nearby for marker in positive_markers
            ):
                return False
            start = position + len(name)
    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_prefix(store: PostgresRedisBlobSwarmStore, prefix: str) -> dict[str, str]:
    def operation(cur) -> dict[str, str]:
        cur.execute("select key, value from swarm_shared_state where key like %s", (prefix + "%",))
        return {row[0]: row[1] for row in cur.fetchall()}

    return store._run_pg(operation)


def fetch_shared_state_value(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()