group_fqcd_jr/app/infrastructure/milvus_knowledge_adapter.py

from typing import Any

from app.core.errors import ForbiddenAgentError, RecoverableAgentError
from app.core.knowledge_contracts import ALLOWED_KNOWLEDGE_COLLECTIONS


class MilvusKnowledgeClient:
    def __init__(self, uri: str, token: str | None = None) -> None:
        self._uri = uri
        self._token = token
        self._client: Any | None = None

    async def _ensure_client(self) -> Any:
        if self._client is None:
            from pymilvus import AsyncMilvusClient  # type: ignore[import-untyped]

            self._client = AsyncMilvusClient(uri=self._uri, token=self._token)
        return self._client

    async def search(
        self, collection: str, vector: list[float], top_k: int
    ) -> list[dict[str, Any]]:
        if collection not in ALLOWED_KNOWLEDGE_COLLECTIONS:
            raise ForbiddenAgentError("未授权的知识集合")
        if len(vector) != 1024 or not 1 <= top_k <= 20:
            raise RecoverableAgentError("知识检索参数无效")
        try:
            client = await self._ensure_client()
            # Lite 重启后集合默认未加载；远程 Milvus 对重复加载保持幂等。
            load_collection = getattr(client, "load_collection", None)
            if load_collection is not None:
                await load_collection(collection_name=collection)
            batches = await client.search(
                collection_name=collection,
                data=[vector],
                limit=top_k,
                output_fields=["knowledge_id", "title", "snippet", "tags", "version"],
                search_params={"metric_type": "COSINE"},
            )
        except Exception as exc:
            raise RecoverableAgentError("知识检索不可用") from exc
        return [
            normalized
            for batch in batches
            for hit in batch
            if (normalized := self._normalize_hit(hit)) is not None
        ]

    @staticmethod
    def _normalize_hit(hit: Any) -> dict[str, Any] | None:
        """统一 Milvus SDK 的平铺与 entity 包装命中格式。"""
        raw = dict(hit)
        entity = raw.get("entity")
        fields = entity if isinstance(entity, dict) else raw
        knowledge_id = fields.get("knowledge_id")
        snippet = fields.get("snippet")
        score = raw.get("score", raw.get("distance", fields.get("score")))
        if (
            not isinstance(knowledge_id, str)
            or not isinstance(snippet, str)
            or not isinstance(score, (int, float))
            or isinstance(score, bool)
        ):
            return None
        normalized: dict[str, Any] = {
            "knowledge_id": knowledge_id,
            "snippet": snippet,
            "score": float(score),
        }
        for field in ("title", "tags", "version"):
            value = fields.get(field)
            if value is not None:
                normalized[field] = value
        return normalized