zhifu-platform/tests/retrieval_evaluation.py

import asyncio
import hashlib
import json
from pathlib import Path
from time import perf_counter
from uuid import UUID, uuid4

from test_after_sales import scenario  # noqa: F401

from zhifu.modules.identity.access_models import StoreGrant
from zhifu.modules.knowledge.config import KnowledgeSettings
from zhifu.modules.knowledge.vector import connection, search


async def publish_dataset(context):
    dataset = json.loads(
        await asyncio.to_thread(Path("evaluations/retrieval.json").read_text, encoding="utf-8")
    )
    async with context["app"].state.sessions.begin() as session:
        session.add(
            StoreGrant(
                tenant_id=context["tenant"].bytes,
                principal_id=context["accounts"]["agent"].bytes,
                store_id=context["store"].bytes,
                permission="knowledge:write",
            )
        )
    identifiers = {}
    for document in dataset["documents"]:
        identifier = str(uuid4())
        identifiers[identifier] = document["key"]
        path = f"/admin/v1/stores/{context['store']}/knowledge/{identifier}"
        response = await context["call"](
            "PUT",
            path,
            "agent",
            {
                "expected_version": 0,
                "title": document["title"],
                "text": document["text"],
            },
        )
        assert response.status_code == 200, response.text
        response = await context["call"](
            "POST", path + "/publish", "agent", {"expected_version": 1}
        )
        assert response.status_code == 200, response.text
    return dataset, identifiers


async def test_retrieval_quality(scenario):  # noqa: F811
    try:
        dataset, identifiers = await publish_dataset(scenario)
        settings = KnowledgeSettings()
        results = []
        async with asyncio.timeout(180):
            await evaluate_queries(scenario, dataset, identifiers, results)
        report = {
            "dataset": dataset["version"],
            "dataset_sha256": hashlib.sha256(await asyncio.to_thread(
                Path("evaluations/retrieval.json").read_bytes)).hexdigest(),
            "origin": "local_regression",
            "production_acceptance": "NOT_RUN",
            "embedding_model": settings.embedding_model,
            "embedding_revision": settings.embedding_revision,
            "generation_attempts": 0,
            "request_limit": 51,
            "query_timeout_seconds": 180,
            "documents": len(identifiers),
            "queries": len(results),
            "recall": {
                str(k): sum(row["rank"] is not None and row["rank"] <= k for row in results)
                / len(results)
                for k in (1, 3, 5, 10)
            },
            "mrr_at_10": sum(1 / row["rank"] if row["rank"] else 0 for row in results)
            / len(results),
            "results": results,
        }
        await asyncio.to_thread(
            Path("docs/verification/retrieval-at10-results.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        assert report["recall"]["10"] >= 0.9, report
        assert report["mrr_at_10"] >= 0.8, report
    finally:
        with connection() as (client, collection):
            client.delete(collection, filter=f'tenant == "{scenario["tenant"].hex}"')


async def evaluate_queries(scenario, dataset, identifiers, results):  # noqa: F811
    path = f"/openapi/v1/stores/{scenario['store']}/knowledge/search"
    for limit in (0, 11):
        response = await scenario["call"](
            "POST", path, payload={"query": "换货", "limit": limit})
        assert response.status_code == 422
    default = await scenario["call"]("POST", path, payload={"query": "换货"})
    assert default.status_code == 200
    assert len(default.json()["items"]) == 5
    for question in dataset["queries"]:
        started = perf_counter()
        response = await scenario["call"](
            "POST", path, payload={"query": question["query"], "limit": 10})
        elapsed = perf_counter() - started
        assert response.status_code == 200, response.text
        keys = list(dict.fromkeys(
            identifiers[item["document_id"]] for item in response.json()["items"]))
        rank = keys.index(question["expected"]) + 1 if question["expected"] in keys else None
        reference = await search(
            scenario["tenant"].bytes, scenario["store"].bytes, question["query"])
        reference_keys = list(dict.fromkeys(
            identifiers[str(UUID(hex=item["entity"]["document"]))] for item in reference[:10]))
        assert keys == reference_keys, {
            "query": question["query"], "pooled": keys, "reference": reference_keys}
        results.append({**question, "retrieved": keys, "rank": rank, "seconds": elapsed,
                        "reference_retrieved": reference_keys})