zhifu-platform/scripts/evaluate_agent_runs.py

import argparse
import asyncio
import hashlib
import json
from collections import Counter, defaultdict
from pathlib import Path
from time import time
from uuid import UUID

from sqlalchemy import select

from zhifu.app import create_app
from zhifu.evaluation import AgentManifest, Dataset, Judgment, digest, source_digest, write_report
from zhifu.modules.agent_runtime.budget_callback import CallAllowance
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt
from zhifu.modules.agent_runtime.worker import execute_one
from zhifu.modules.integrations.models import ExternalOperation


class PersistentAllowance(CallAllowance):
    def __init__(self, manifest, report, output, state):
        self.manifest, self.report, self.output, self.state = manifest, report, output, state
        super().__init__(manifest.max_calls - report["reserved_calls"])

    async def take(self):
        # 在模型调用前保存预留，未知结果及进程中断均消耗额度，重启不会增加限额。
        maximum = min(self.manifest.max_calls,
                      self.manifest.max_total_output_tokens // self.manifest.max_output_tokens)
        reserved = await self.state.eval(
            "local calls = tonumber(redis.call('HGET', KEYS[1], 'calls')); "
            "local deadline = tonumber(redis.call('HGET', KEYS[1], 'deadline')); "
            "if not calls or not deadline then return -3 end; "
            "if tonumber(redis.call('TIME')[1]) >= deadline then return -2 end; "
            "if calls >= tonumber(ARGV[1]) then return -1 end; "
            "return redis.call('HINCRBY', KEYS[1], 'calls', 1)",
            1, "evaluation-budget:" + self.report["manifest_sha256"], maximum,
        )
        if reserved < 0:
            raise ValueError("评测调用或时间限额已用完，或持久额度记录缺失")
        self.started += 1
        self.report["reserved_calls"] = reserved
        await asyncio.to_thread(write_report, self.output, self.report)


async def initialize_report(args, dataset, manifest, cases, state):
    manifest_hash = digest(manifest.model_dump(mode="json"))
    budget_key = "evaluation-budget:" + manifest_hash
    output_identity = str(args.output.resolve())
    # 同一清单绑定唯一输出位置，Redis 额度不设置过期时间，修改路径不能获得新额度。
    binding = await state.hget(budget_key, "output")
    if binding is not None and binding.decode() != output_identity:
        raise ValueError("同一评测清单已经绑定其他输出文件")
    if args.output.exists():
        report = json.loads(await asyncio.to_thread(args.output.read_text, encoding="utf-8"))
        if report["manifest_sha256"] != manifest_hash:
            raise ValueError("恢复执行时清单必须保持一致")
        if binding is None:
            raise ValueError("恢复报告缺少 Redis 持久额度记录")
    else:
        if binding is not None:
            raise ValueError("已登记评测的报告文件缺失，需要从备份恢复")
        report = {"manifest_sha256": manifest_hash, "dataset_version": dataset.version,
                  "generation_mode": "bounded" if args.execute else "read_only",
                  "reserved_calls": 0, "deadline": (
                      time() + manifest.total_timeout_seconds if args.execute else 0),
                  "cases": {key: {**case, "status": "NOT_RUN"} for key, case in cases.items()}}
        await state.hset(budget_key, mapping={"output": output_identity, "calls": 0,
                                            "deadline": report["deadline"],
                                            "mode": report["generation_mode"]})
    mode = (await state.hget(budget_key, "mode")).decode()
    if args.execute and mode == "read_only":
        mode = "bounded"
        await state.hset(budget_key, mapping={"mode": mode,
            "deadline": time() + manifest.total_timeout_seconds})
    report["generation_mode"] = mode
    report["reserved_calls"] = int(await state.hget(budget_key, "calls"))
    report["deadline"] = float(await state.hget(budget_key, "deadline"))
    return report


def load_inputs(dataset_path, manifest_path):
    content = dataset_path.read_bytes()
    dataset = Dataset.model_validate_json(content)
    manifest = AgentManifest.model_validate_json(manifest_path.read_bytes())
    if hashlib.sha256(content).hexdigest() != manifest.dataset_sha256:
        raise ValueError("评测资料摘要不一致")
    if source_digest() != manifest.source_sha256:
        raise ValueError("评测固定代码版本摘要不一致")
    cases = {case["id"]: case for case in dataset.cases()}
    for binding in manifest.cases:
        if binding.case_id not in cases or cases[binding.case_id]["split"] != manifest.split:
            raise ValueError("运行清单引用错误集合或未知评测编号")
        binding.review.verify()
        binding.environment_evidence.verify()
    return dataset, manifest, cases


def category_report(rows):
    groups = defaultdict(Counter)
    for row in rows.values():
        groups[f"{row['split']}/{row['category']}"][row["status"]] += 1
    return {key: dict(value) for key, value in groups.items()}


async def observe(sessions, binding, manifest, case):
    async with sessions() as session:
        run = await session.get(AgentRun, binding.run_id.bytes)
        if run is None or run.tenant_id != binding.tenant_id.bytes:
            raise ValueError("评测运行不存在或租户不一致")
        if run.graph_version != manifest.graph_version or run.model_name != manifest.model_name:
            raise ValueError("运行图版本或模型版本不一致")
        if not run.release_snapshot or digest(run.release_snapshot) != manifest.release_sha256:
            raise ValueError("运行发布参数摘要不一致")
        if run.release_snapshot.get("max_output_tokens") != manifest.max_output_tokens:
            raise ValueError("发布参数的单次输出上限不一致")
        prompts = [message["data"]["content"] for message in run.messages
                   if message["type"] == "human"]
        if not prompts or prompts[0] != case["prompt"]:
            raise ValueError("运行的初始问题与评测资料不一致")
        attempts = (await session.scalars(select(ModelAttempt).where(
            ModelAttempt.run_id == run.id
        ).order_by(ModelAttempt.created_at, ModelAttempt.id))).all()
        observation = {"run_id": str(binding.run_id), "run_status": run.status,
                       "error_code": run.error_code, "messages": run.messages,
                       "waiting": run.waiting,
                       "attempts": [{"id": str(UUID(bytes=item.id)), "status": item.status,
                                     "usage": item.usage} for item in attempts]}
        operations = (await session.scalars(select(ExternalOperation).where(
            ExternalOperation.tenant_id == run.tenant_id,
            ExternalOperation.ticket_id == run.business_ticket_id,
        ).order_by(ExternalOperation.id))).all() if run.business_ticket_id else []
        observation["external_operations"] = [{"id": str(UUID(bytes=item.id)),
            "status": item.status, "proposal_version": item.proposal_version}
            for item in operations]
        return observation


async def evaluate(args, dataset, manifest, cases):
    manifest_hash = digest(manifest.model_dump(mode="json"))
    output = args.output
    judgments = {}
    if args.judgments:
        for value in json.loads(await asyncio.to_thread(args.judgments.read_bytes)):
            judgment = Judgment.model_validate(value)
            await asyncio.to_thread(judgment.verify)
            if judgment.case_id in judgments:
                raise ValueError("结果审核编号重复")
            judgments[judgment.case_id] = judgment
        if not judgments.keys() <= {binding.case_id for binding in manifest.cases}:
            raise ValueError("结果审核包含清单之外的评测")
    app = create_app()
    async with app.router.lifespan_context(app):
        lock = app.state.state.lock(
            "evaluation:" + manifest_hash, timeout=manifest.total_timeout_seconds + 30,
            blocking_timeout=0, thread_local=False
        )
        if not await lock.acquire():
            raise RuntimeError("同一评测清单已有执行进程")
        try:
            report = await initialize_report(args, dataset, manifest, cases, app.state.state)
            if args.execute:
                for role in ("agent", "memory"):
                    async for _ in app.state.state.scan_iter(match=f"monitor:worker:{role}:*"):
                        raise RuntimeError("受限评测要求停止常驻生成工作进程")
            allowance = PersistentAllowance(manifest, report, output, app.state.state)
            await asyncio.to_thread(write_report, output, report)
            for binding in manifest.cases:
                row = report["cases"][binding.case_id]
                observation = await observe(app.state.sessions, binding, manifest, row)
                # 已有终态和业务等待只读取证据；重新执行只允许当前排队状态。
                if args.execute and observation["run_status"] == "QUEUED":
                    remaining = report["deadline"] - time()
                    if remaining <= 0:
                        raise TimeoutError("评测总时间已达到上限")
                    try:
                        async with asyncio.timeout(remaining):
                            await execute_one(
                                app.state.sessions, app.state.state,
                                AgentSettings(max_output_tokens=manifest.max_output_tokens),
                                app.state.knowledge.search, allowance=allowance,
                                run_ids=[binding.run_id.bytes],
                            )
                    finally:
                        observation = await observe(app.state.sessions, binding, manifest, row)
                        row.update(observation=observation, status=(
                            "REVIEW_REQUIRED" if observation["attempts"] else "NOT_RUN"))
                        row["observation_sha256"] = digest(observation)
                        await asyncio.to_thread(write_report, output, report)
                row.update(observation=observation)
                row["observation_sha256"] = digest(observation)
                row["status"] = "REVIEW_REQUIRED" if observation["attempts"] else "NOT_RUN"
                if binding.case_id in judgments:
                    judgment = judgments[binding.case_id]
                    if (judgment.observation_sha256 != row["observation_sha256"]
                            or not observation["attempts"]
                            or observation["run_status"] in {"QUEUED", "RUNNING"}):
                        raise ValueError("结果审核必须对应已执行并暂停或结束的当前运行证据")
                    row["status"] = judgment.verdict
                    row["judgment"] = judgment.model_dump(mode="json")
                elif row.get("judgment"):
                    previous = Judgment.model_validate(row["judgment"])
                    await asyncio.to_thread(previous.verify)
                    if previous.observation_sha256 == row["observation_sha256"]:
                        row["status"] = previous.verdict
                    else:
                        row.pop("judgment")
                report["categories"] = category_report(report["cases"])
                report["recorded_attempts"] = sum(len(item.get("observation", {}).get(
                    "attempts", [])) for item in report["cases"].values())
                attempts = [attempt for item in report["cases"].values()
                            for attempt in item.get("observation", {}).get("attempts", [])]
                report["usage"] = {key: sum((item["usage"] or {}).get(key, 0)
                                            for item in attempts)
                                   for key in ("input_tokens", "output_tokens", "total_tokens")}
                report["attempts_without_usage"] = sum(item["usage"] is None for item in attempts)
                await asyncio.to_thread(write_report, output, report)
        finally:
            await lock.release()
    print(json.dumps({"output": str(output), "reserved_calls": report["reserved_calls"],
                      "categories": report["categories"]}, ensure_ascii=False))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="读取或受限执行已审核的 Agent 评测清单")
    parser.add_argument("--dataset", type=Path, default=Path("evaluations/agent-scenarios.json"))
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--execute", action="store_true", help="显式允许模型请求")
    parser.add_argument("--judgments", type=Path, help="人工结果审核 JSON 列表")
    parser.add_argument("--validate-only", action="store_true")
    args = parser.parse_args()
    dataset, manifest, cases = load_inputs(args.dataset, args.manifest)
    if args.validate_only:
        print(json.dumps({"valid": True, "cases": len(manifest.cases), "generation_attempts": 0}))
    else:
        asyncio.run(evaluate(args, dataset, manifest, cases))