zhifu-platform/scripts/evaluate_retrieval.py

import argparse
import asyncio
import json
import os
from collections import defaultdict
from pathlib import Path
from time import perf_counter
from uuid import UUID

import httpx

from zhifu.evaluation import RetrievalManifest, digest, write_report


async def evaluate(args, manifest):
    report = {"manifest_sha256": digest(manifest.model_dump(mode="json")),
              "dataset_version": manifest.dataset_version, "origin": manifest.origin,
              "generation_attempts": 0, "request_limit": len(manifest.queries),
              "total_timeout_seconds": manifest.total_timeout_seconds,
              "embedding_model": manifest.embedding_model,
              "embedding_revision": manifest.embedding_revision, "results": [],
              "status": "RUNNING"}
    started = perf_counter()
    # 仅检索现有发布资料。写入和发布由正常业务接口完成，评测令牌仅需读取权限。
    try:
        async with asyncio.timeout(manifest.total_timeout_seconds), httpx.AsyncClient(
            base_url=args.base_url, timeout=30, trust_env=False,
            headers={"Authorization": "Bearer " + os.environ[args.token_env]},
        ) as client:
            for item in manifest.queries:
                began = perf_counter()
                response = await client.post(
                    f"/openapi/v1/stores/{manifest.store_id}/knowledge/search",
                    json={"query": item.query, "limit": 10},
                )
                response.raise_for_status()
                found = []
                for document in response.json()["items"]:
                    identifier = UUID(document["document_id"])
                    if manifest.documents.get(identifier) != document["version"]:
                        raise ValueError("检索结果出现未登记资料或版本变化")
                    if identifier not in found:
                        found.append(identifier)
                recall = len(set(found) & set(item.expected)) / len(set(item.expected))
                report["results"].append({"query": item.query, "category": item.category,
                    "expected": [str(key) for key in item.expected],
                    "retrieved": [str(key) for key in found], "recall_at_10": recall,
                    "seconds": perf_counter() - began})
                await asyncio.to_thread(write_report, args.output, report)
        categories = defaultdict(list)
        for result in report["results"]:
            categories[result["category"]].append(result["recall_at_10"])
        report["categories"] = {key: {"count": len(values),
            "recall_at_10": sum(values) / len(values)} for key, values in categories.items()}
        report["recall_at_10"] = sum(row["recall_at_10"] for row in report["results"]) / len(
            report["results"])
        report["status"] = "PASSED" if report["recall_at_10"] >= 0.9 else "FAILED"
    finally:
        if report["status"] == "RUNNING":
            report["status"] = "INCOMPLETE"
        report["seconds"] = perf_counter() - started
        await asyncio.to_thread(write_report, args.output, report)
    print(json.dumps({key: value for key, value in report.items() if key != "results"}))
    if report["status"] != "PASSED":
        raise ValueError("Recall@10 未达到 90%")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="真实知识检索 Recall@10 验收")
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--base-url", default="http://127.0.0.1:18000")
    parser.add_argument("--token-env", default="EVALUATION_ACCESS_TOKEN")
    args = parser.parse_args()
    manifest = RetrievalManifest.model_validate_json(args.manifest.read_bytes())
    if manifest.review:
        manifest.review.verify()
    if manifest.deployment_review:
        manifest.deployment_review.verify()
    asyncio.run(evaluate(args, manifest))