66 lines
2.8 KiB
Python
66 lines
2.8 KiB
Python
|
"""长期记忆画像向量的 Milvus 客户端。
|
||
|
|
|
||
|
|
与 `MilvusKnowledgeWriter` 分开：那个适配器的 `upsert` 签名绑定知识集合的
|
||
|
|
`knowledge_id` 主键与字段表，而画像投影的主键是 `memory_uuid`、且需要先 `query`
|
||
|
|
按版本判重。两者共用一套连接口径（惰性连接 + 失败一律 `RecoverableAgentError`），
|
||
|
|
但不硬把两种 schema 塞进同一个类。
|
||
|
|
|
||
|
|
与召回侧的客户端也分开：写路径不与检索进程共用连接（读写物理隔离，向量库故障
|
||
|
|
不能从写路径传染到问答主链路），与 `get_milvus_knowledge_writer` 的取向一致。
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from app.core.errors import RecoverableAgentError
|
||
|
|
|
||
|
|
|
||
|
|
class MilvusProfileVectorClient:
|
||
|
|
"""满足 `MilvusProfileProjection` 所需的 `query` / `upsert` / `load_collection`。"""
|
||
|
|
|
||
|
|
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 load_collection(self, *, collection_name: str) -> None:
|
||
|
|
"""把集合载入内存后再查/写。
|
||
|
|
|
||
|
|
集合不存在时抛 `RecoverableAgentError`：由 outbox 退避重试并最终判死信，
|
||
|
|
而不是静默跳过——"集合没建"是装配问题，必须可见。
|
||
|
|
"""
|
||
|
|
client = await self._ensure()
|
||
|
|
try:
|
||
|
|
await client.load_collection(collection_name=collection_name)
|
||
|
|
except Exception as exc:
|
||
|
|
raise RecoverableAgentError("画像向量集合不可用") from exc
|
||
|
|
|
||
|
|
async def query(self, **kwargs: Any) -> list[dict[str, Any]]:
|
||
|
|
"""按 filter 查询；返回空列表表示无匹配（不是错误）。"""
|
||
|
|
client = await self._ensure()
|
||
|
|
try:
|
||
|
|
raw = await client.query(**kwargs)
|
||
|
|
except Exception as exc:
|
||
|
|
raise RecoverableAgentError("画像向量查询失败") from exc
|
||
|
|
return list(raw or [])
|
||
|
|
|
||
|
|
async def upsert(self, **kwargs: Any) -> Any:
|
||
|
|
client = await self._ensure()
|
||
|
|
try:
|
||
|
|
return await client.upsert(**kwargs)
|
||
|
|
except RecoverableAgentError:
|
||
|
|
raise
|
||
|
|
except Exception as exc:
|
||
|
|
raise RecoverableAgentError("画像向量写入失败") from exc
|