diff --git a/app/infrastructure/milvus_profile_projection.py b/app/infrastructure/milvus_profile_projection.py new file mode 100644 index 0000000..fcea495 --- /dev/null +++ b/app/infrastructure/milvus_profile_projection.py @@ -0,0 +1,127 @@ +"""Milvus 长期记忆投影适配器。 + +只写入已经审核的 `memory_sources`，不接受画像快照整体冒充单条记忆。 +""" + +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime +from typing import Any, Protocol +from uuid import UUID + +from app.core.conversation_privacy import sanitize_customer_service_message +from app.core.errors import RecoverableAgentError + +PROFILE_COLLECTION = "user_long_term_memory_v1" +VECTOR_DIM = 1024 + + +class MilvusProfileClient(Protocol): + async def query(self, **kwargs: Any) -> list[dict[str, Any]]: ... + + async def upsert(self, **kwargs: Any) -> Any: ... + + +EmbeddingProvider = Callable[[str], Awaitable[list[float]]] + + +class MilvusProfileProjection: + """按记忆 UUID 幂等写入长期记忆向量。""" + + def __init__( + self, + client: MilvusProfileClient, + embed: EmbeddingProvider, + *, + collection: str = PROFILE_COLLECTION, + ) -> None: + self._client = client + self._embed = embed + self._collection = collection + + async def upsert(self, payload: dict[str, Any]) -> None: + customer_id, profile_version, sources = self._normalize(payload) + load_collection = getattr(self._client, "load_collection", None) + if load_collection is not None: + await load_collection(collection_name=self._collection) + rows: list[dict[str, Any]] = [] + for source in sources: + vector = await self._embed(source["content"]) + if len(vector) != VECTOR_DIM: + raise RecoverableAgentError("画像向量维度不一致") + existing = await self._client.query( + collection_name=self._collection, + filter=f'memory_uuid == "{source["memory_uuid"]}"', + output_fields=["memory_uuid", "version", "customer_id"], + ) + if existing and int(existing[0].get("version", 0)) > source["version"]: + continue + rows.append({ + "memory_uuid": source["memory_uuid"], + "customer_id": customer_id, + "content": source["content"], + "embedding": vector, + "memory_type": source["memory_type"], + "memory_key": source["memory_key"], + "confidence": source["confidence"], + "version": source["version"], + "status": "active", + "valid_until_ts": source["valid_until_ts"], + "updated_at_ts": source["updated_at_ts"], + }) + if rows: + await self._client.upsert(collection_name=self._collection, data=rows) + + @staticmethod + def _normalize( + payload: dict[str, Any], + ) -> tuple[int, int, list[dict[str, Any]]]: + customer_id = payload.get("customer_id") + profile_version = payload.get("profile_version") + sources = payload.get("memory_sources") + if not isinstance(customer_id, int) or customer_id <= 0: + raise ValueError("customer_id is invalid") + if not isinstance(profile_version, int) or profile_version <= 0: + raise ValueError("profile_version is invalid") + if not isinstance(sources, list): + raise ValueError("memory_sources is invalid") + normalized: list[dict[str, Any]] = [] + now = int(datetime.now(UTC).timestamp()) + for source in sources: + if not isinstance(source, dict): + raise ValueError("memory source is invalid") + required = [source.get(name) for name in ( + "memory_uuid", "memory_key", "content", "memory_type" + )] + if not all(isinstance(value, str) and value.strip() for value in required): + raise ValueError("memory source fields are invalid") + try: + memory_uuid = str(UUID(str(source["memory_uuid"]))) + except ValueError as exc: + raise ValueError("memory_uuid is invalid") from exc + memory_key = str(source["memory_key"]).strip() + if not (memory_key.startswith("preference:") or memory_key.startswith("goal:")): + raise ValueError("memory key is not projectable") + confidence = source.get("confidence") + version = source.get("version") + if not isinstance(confidence, (int, float)) or not 0 <= confidence <= 1: + raise ValueError("memory confidence is invalid") + if not isinstance(version, int) or version <= 0: + raise ValueError("memory version is invalid") + valid_until = source.get("valid_until") + valid_until_ts = None + if isinstance(valid_until, str) and valid_until: + try: + valid_until_ts = int(datetime.fromisoformat(valid_until).timestamp()) + except ValueError as exc: + raise ValueError("memory valid_until is invalid") from exc + normalized.append({ + "memory_uuid": memory_uuid, + "memory_key": memory_key, + "content": sanitize_customer_service_message(str(source["content"])).strip(), + "memory_type": str(source["memory_type"]).strip(), + "confidence": float(confidence), + "version": version, + "valid_until_ts": valid_until_ts, + "updated_at_ts": now, + }) + return customer_id, profile_version, normalized diff --git a/app/worker/__main__.py b/app/worker/__main__.py index 963957b..cf7175e 100644 --- a/app/worker/__main__.py +++ b/app/worker/__main__.py @@ -4,8 +4,12 @@ import logging from typing import Any from app.core.config import get_settings +from app.core.errors import RecoverableAgentError from app.infrastructure.db import engine +from app.infrastructure.milvus_profile_projection import MilvusProfileProjection from app.infrastructure.neo4j_profile_projection import Neo4jProfileProjection +from app.service.agent.bootstrap import get_memory_embedding_service +from app.service.model_gateway import DatabaseModelEndpointResolver from app.worker.memory_sync_outbox_worker import MemorySyncOutboxWorker from app.worker.offsite_mail_worker import OffsiteMailWorker from app.worker.runtime import WorkerRuntime @@ -18,7 +22,9 @@ async def serve(*, once: bool = False) -> None: runtime = WorkerRuntime(settings=settings) offsite_worker = OffsiteMailWorker(settings) neo4j_driver: Any | None = None + milvus_client: Any | None = None memory_sync_worker: MemorySyncOutboxWorker | None = None + memory_sync_handlers: dict[str, Any] = {} if settings.neo4j_password: from neo4j import AsyncGraphDatabase @@ -31,6 +37,33 @@ async def serve(*, once: bool = False) -> None: }) else: logger.warning("Neo4j password not configured; profile projection remains pending") + if settings.resolved_milvus_uri and settings.knowledge_embedding_endpoint_code: + from pymilvus import AsyncMilvusClient # type: ignore[import-untyped] + + milvus_client = AsyncMilvusClient( + uri=settings.resolved_milvus_uri, + token=settings.milvus_token or None, + ) + + async def embed_profile(text: str) -> list[float]: + endpoints = await DatabaseModelEndpointResolver().resolve( + agent_type="memory_projection", task_type="embedding" + ) + if not endpoints: + raise RecoverableAgentError("没有可用的 embedding 端点") + return (await get_memory_embedding_service().embed(endpoints, text)).vector + + memory_sync_handlers["milvus"] = MilvusProfileProjection( + milvus_client, embed_profile + ).upsert + else: + logger.warning( + "Milvus profile projection not configured; profile projection remains pending" + ) + if neo4j_driver is not None: + memory_sync_handlers["neo4j"] = Neo4jProfileProjection(neo4j_driver).upsert + if memory_sync_handlers: + memory_sync_worker = MemorySyncOutboxWorker(memory_sync_handlers) try: while True: try: @@ -58,6 +91,8 @@ async def serve(*, once: bool = False) -> None: await offsite_worker.close() if neo4j_driver is not None: await neo4j_driver.close() + if milvus_client is not None: + await milvus_client.close() await engine.dispose() diff --git a/tests/unit/infrastructure/test_milvus_profile_projection.py b/tests/unit/infrastructure/test_milvus_profile_projection.py new file mode 100644 index 0000000..54d651f --- /dev/null +++ b/tests/unit/infrastructure/test_milvus_profile_projection.py @@ -0,0 +1,96 @@ +from uuid import uuid4 + +import pytest + +from app.core.errors import RecoverableAgentError +from app.infrastructure.milvus_profile_projection import MilvusProfileProjection + + +class FakeMilvus: + def __init__(self, existing: list[dict[str, object]] | None = None) -> None: + self.existing = existing or [] + self.queries: list[dict[str, object]] = [] + self.upserts: list[dict[str, object]] = [] + + async def query(self, **kwargs: object) -> list[dict[str, object]]: + self.queries.append(kwargs) + return self.existing + + async def upsert(self, **kwargs: object) -> None: + self.upserts.append(kwargs) + + +def payload() -> dict[str, object]: + return { + "customer_id": 7, + "profile_version": 1, + "memory_sources": [{ + "memory_uuid": str(uuid4()), + "memory_key": "preference:risk_level", + "content": "稳健型", + "memory_type": "preference", + "confidence": 0.9, + "version": 2, + "valid_until": None, + }], + } + + +@pytest.mark.asyncio +async def test_upsert_writes_schema_fields_and_vector() -> None: + client = FakeMilvus() + projection = MilvusProfileProjection(client, _embed) + + await projection.upsert(payload()) + + assert len(client.upserts) == 1 + row = client.upserts[0]["data"][0] + assert row["customer_id"] == 7 + assert row["status"] == "active" + assert len(row["embedding"]) == 1024 + + +@pytest.mark.asyncio +async def test_lower_memory_version_is_not_overwritten() -> None: + data = payload() + source = data["memory_sources"][0] + assert isinstance(source, dict) + memory_uuid = source["memory_uuid"] + client = FakeMilvus(existing=[{ + "memory_uuid": memory_uuid, "customer_id": 7, "version": 3, + }]) + + await MilvusProfileProjection(client, _embed).upsert(data) + + assert client.upserts == [] + + +@pytest.mark.asyncio +async def test_embedding_dimension_is_enforced() -> None: + with pytest.raises(RecoverableAgentError, match="维度"): + await MilvusProfileProjection(client=FakeMilvus(), embed=_embed_short).upsert( + payload() + ) + + +@pytest.mark.asyncio +async def test_non_uuid_memory_id_is_rejected() -> None: + data = payload() + source = data["memory_sources"][0] + assert isinstance(source, dict) + source["memory_uuid"] = "unsafe\" or true" + + with pytest.raises(ValueError, match="memory_uuid"): + await MilvusProfileProjection(FakeMilvus(), _embed).upsert(data) + + +def _vector(size: int = 1024) -> list[float]: + return [0.0] * size + + +async def _embed(_: str) -> list[float]: + return _vector() + + +async def _embed_short(_: str) -> list[float]: + return _vector(3)