75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
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
|