128 lines
5.3 KiB
Python
128 lines
5.3 KiB
Python
"""知识检索**读路径**的 Milvus 适配器（Task 6）。
|
||
|
||
与写路径的物理隔离（`app/infrastructure/milvus_knowledge_writer.py`）：
|
||
|
||
- 本模块**只读**（`search` + `describe_collection`），绝不 import 写路径模块 ——
|
||
检索进程不持有写客户端，向量库故障不会从写侧传染到问答主链路，反之亦然；
|
||
- 检索侧只接受 `ALLOWED_COLLECTIONS` 内的集合名，校验由调用方
|
||
（`KnowledgeRetrievalService._assert_collections_allowed`）在**任何网络调用之前**完成，
|
||
本适配器只做"第二道闸"式的兜底拒绝。
|
||
|
||
连接是**惰性**的：`__init__` 不连 Milvus，首次检索才 `import pymilvus` 并建
|
||
`AsyncMilvusClient`。任何连接/查询失败统一转成 `RecoverableAgentError`，
|
||
由 `KnowledgeRetrievalService` 捕获后降级 MySQL LIKE —— 本层不吞异常、不返回空结果，
|
||
避免把"向量库挂了"伪装成"知识库没有内容"。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from app.core.errors import ForbiddenAgentError, RecoverableAgentError
|
||
from app.core.knowledge_contracts import ALLOWED_COLLECTIONS
|
||
|
||
#: 返回给上层的标量字段（不含 `embedding`：向量不回传，省带宽也避免误用）。
|
||
DEFAULT_OUTPUT_FIELDS: tuple[str, ...] = (
|
||
"knowledge_id",
|
||
"title",
|
||
"snippet",
|
||
"tags",
|
||
"version",
|
||
"intent",
|
||
)
|
||
|
||
#: 一次检索最多返回的候选数下限/上限，防调用方传入异常 top_k。
|
||
MIN_TOP_K = 1
|
||
MAX_TOP_K = 50
|
||
|
||
|
||
class MilvusKnowledgeClient:
|
||
"""知识向量读边界：`search` 失败一律 `RecoverableAgentError`，由上层降级。"""
|
||
|
||
def __init__(self, uri: str, token: str = "") -> None:
|
||
self._uri = uri
|
||
self._token = token
|
||
self._client: Any = None
|
||
|
||
async def _ensure(self) -> Any:
|
||
if self._client is None:
|
||
try:
|
||
from pymilvus import AsyncMilvusClient # type: ignore[import-untyped]
|
||
except ImportError as exc: # pragma: no cover - 依赖已声明，缺装是环境问题
|
||
raise RecoverableAgentError("pymilvus 未安装，无法检索知识向量") from exc
|
||
try:
|
||
self._client = AsyncMilvusClient(uri=self._uri, token=self._token or None)
|
||
except Exception as exc:
|
||
raise RecoverableAgentError("Milvus 读客户端初始化失败") from exc
|
||
return self._client
|
||
|
||
async def search(
|
||
self,
|
||
*,
|
||
collection: str,
|
||
vector: list[float],
|
||
top_k: int,
|
||
output_fields: tuple[str, ...] | list[str] = DEFAULT_OUTPUT_FIELDS,
|
||
) -> list[dict[str, Any]]:
|
||
"""向量检索，返回**扁平化**的命中行（每行含标量字段与 `score`）。
|
||
|
||
pymilvus 返回的是 `[[{id, distance, entity}]]`（每查询一组），这里折叠为单层列表；
|
||
距离字段名可能是 `distance` 或 `score`，两种都吸收（COSINE 越大越相似）。
|
||
"""
|
||
if collection not in ALLOWED_COLLECTIONS:
|
||
raise ForbiddenAgentError(f"知识集合不在白名单内：{collection}")
|
||
if not vector:
|
||
raise RecoverableAgentError("检索向量不能为空")
|
||
client = await self._ensure()
|
||
try:
|
||
raw = await client.search(
|
||
collection_name=collection,
|
||
data=[list(vector)],
|
||
limit=max(MIN_TOP_K, min(int(top_k), MAX_TOP_K)),
|
||
output_fields=list(output_fields),
|
||
search_params={"metric_type": "COSINE"},
|
||
)
|
||
except Exception as exc:
|
||
raise RecoverableAgentError(f"知识向量检索失败：{collection}") from exc
|
||
return self.parse_hits(raw)
|
||
|
||
@staticmethod
|
||
def parse_hits(raw: Any) -> list[dict[str, Any]]:
|
||
"""把 pymilvus 的嵌套命中结构折叠为 `list[dict]`（纯函数，不抛异常）。"""
|
||
rows: list[dict[str, Any]] = []
|
||
for group in raw if isinstance(raw, list | tuple) else [raw]:
|
||
items = group if isinstance(group, list | tuple) else [group]
|
||
for item in items:
|
||
row = MilvusKnowledgeClient._as_row(item)
|
||
if row is not None:
|
||
rows.append(row)
|
||
return rows
|
||
|
||
@staticmethod
|
||
def _as_row(item: Any) -> dict[str, Any] | None:
|
||
if isinstance(item, dict):
|
||
fields = dict(item)
|
||
elif hasattr(item, "entity") or hasattr(item, "id"):
|
||
fields = {}
|
||
for name in (*DEFAULT_OUTPUT_FIELDS, "id", "distance", "score", "entity"):
|
||
if hasattr(item, name):
|
||
fields[name] = getattr(item, name)
|
||
else:
|
||
return None
|
||
entity = fields.pop("entity", None)
|
||
if isinstance(entity, dict):
|
||
merged = dict(entity)
|
||
merged.update({key: value for key, value in fields.items() if value is not None})
|
||
fields = merged
|
||
knowledge_id = fields.get("knowledge_id") or fields.get("id")
|
||
if knowledge_id is None or not str(knowledge_id).strip():
|
||
# 没有主键的行无法回表校验，丢弃而不是猜造标识（避免假引用）。
|
||
return None
|
||
fields["knowledge_id"] = str(knowledge_id).strip()
|
||
return fields
|
||
|
||
async def close(self) -> None:
|
||
if self._client is not None:
|
||
client, self._client = self._client, None
|
||
await client.close()
|