zhifu-platform/tests/retrieval_stability.py

import asyncio
import json
from pathlib import Path
from uuid import UUID, uuid4

from pymilvus import AnnSearchRequest, RRFRanker
from retrieval_evaluation import publish_dataset
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.knowledge.vector import connection, stable_candidates


async def test_retrieval_stability(scenario):  # noqa: F811
    client = scenario["app"].state.knowledge
    try:
        dataset, identifiers = await publish_dataset(scenario)
        originals = await asyncio.to_thread(
            client.milvus.query,
            client.settings.collection,
            filter=f'tenant == "{scenario["tenant"].hex}"',
            output_fields=[
                "id",
                "tenant",
                "store",
                "document",
                "version",
                "position",
                "text",
                "dense",
            ],
        )
        queries = await client.embeddings([item["query"] for item in dataset["queries"]])
        generations = []
        for generation in range(4):
            store = uuid4().hex
            keys = {}
            rows = []
            for original in originals:
                identifier = uuid4().hex
                keys[identifier] = identifiers[str(UUID(hex=original["document"]))]
                rows.append(
                    {
                        **original,
                        "id": f"{scenario['tenant'].hex}:{identifier}:1:0",
                        "store": store,
                        "document": identifier,
                    }
                )
            await asyncio.to_thread(client.milvus.upsert, client.settings.collection, rows)
            expression = f'tenant == "{scenario["tenant"].hex}" and store == "{store}"'
            results = []
            for question, vector in zip(dataset["queries"], queries, strict=True):
                branches = []
                for field, metric, data in (
                    ("dense", "COSINE", vector),
                    ("sparse", "BM25", question["query"]),
                ):
                    hits = await asyncio.to_thread(
                        client.milvus.search,
                        client.settings.collection,
                        [data],
                        anns_field=field,
                        search_params={"metric_type": metric},
                        filter=expression,
                        limit=30,
                        output_fields=["document"],
                    )
                    branches.append(
                        [
                            {"key": keys[hit["entity"]["document"]], "score": hit["distance"]}
                            for hit in hits[0]
                        ]
                    )
                reqs = [
                    AnnSearchRequest(
                        [vector], "dense", {"metric_type": "COSINE"}, 30, expr=expression
                    ),
                    AnnSearchRequest(
                        [question["query"]], "sparse", {"metric_type": "BM25"}, 30, expr=expression
                    ),
                ]
                repeats, ordered = [], []
                for _ in range(3):
                    hits = await asyncio.to_thread(
                        client.milvus.hybrid_search,
                        client.settings.collection,
                        reqs,
                        RRFRanker(),
                        limit=30,
                        output_fields=["document", "position", "text"],
                    )
                    repeats.append(
                        [
                            {"key": keys[hit["entity"]["document"]], "score": hit["distance"]}
                            for hit in hits[0]
                        ]
                    )
                    stable = stable_candidates(hits[0])
                    ordered.append([keys[hit["entity"]["document"]] for hit in stable])
                    assert [hit["distance"] for hit in stable] == sorted(
                        [hit["distance"] for hit in hits[0]], reverse=True
                    )
                results.append(
                    {
                        **question,
                        "dense": branches[0],
                        "sparse": branches[1],
                        "repeats": repeats,
                        "stable": ordered,
                    }
                )
            generations.append({"generation": generation, "queries": results})
        raw_changes, stable_changes = [], []
        for index, question in enumerate(dataset["queries"]):
            raw = {run["queries"][index]["repeats"][0][0]["key"] for run in generations}
            stable = {run["queries"][index]["stable"][0][0] for run in generations}
            if len(raw) > 1:
                raw_changes.append(question["query"])
            if len(stable) > 1:
                stable_changes.append(question["query"])
        report = {
            "dataset": dataset["version"],
            "generations": generations,
            "raw_top1_changes": raw_changes,
            "stable_top1_changes": stable_changes,
        }
        await asyncio.to_thread(
            Path("docs/verification/retrieval-stability.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        assert not stable_changes, stable_changes
        assert all(
            query["stable"][0] == query["stable"][1] == query["stable"][2]
            for run in generations
            for query in run["queries"]
        )
    finally:
        with connection() as (milvus, collection):
            milvus.delete(collection, filter=f'tenant == "{scenario["tenant"].hex}"')