Mutual_Fund/service/memory/milvus_memory.py

"""客户长期记忆的 Milvus 向量镜像。"""

from __future__ import annotations

from typing import Any

from pymilvus import AsyncMilvusClient, DataType

from config.database.milvus import client as configured_client
from rag.embedding import EMBEDDING_DIMENSION


CUSTOMER_MEMORY_COLLECTION = "customer_memory"


def build_memory_schema():
    """构造客户记忆向量集合结构，维度来自 LLM_EMBED_DIMENSIONS。"""
    schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
    schema.add_field("memory_id", DataType.VARCHAR, is_primary=True, max_length=128)
    schema.add_field("customer_id", DataType.INT64)
    schema.add_field("memory_type", DataType.VARCHAR, max_length=32)
    schema.add_field("tag", DataType.VARCHAR, max_length=64)
    schema.add_field("content", DataType.VARCHAR, max_length=2048)
    schema.add_field("status", DataType.VARCHAR, max_length=16)
    schema.add_field("vector", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIMENSION)
    return schema


def build_memory_index_params():
    """构造客户记忆向量索引。"""
    params = AsyncMilvusClient.prepare_index_params()
    params.add_index(
        field_name="vector",
        index_type="HNSW",
        metric_type="COSINE",
        params={"M": 16, "efConstruction": 200},
    )
    return params


class MilvusMemoryStore:
    """封装客户记忆向量写入、查询和删除。"""

    def __init__(self, client=None, *, collection_name=CUSTOMER_MEMORY_COLLECTION):
        self.client = client or configured_client()
        self.collection_name = collection_name

    async def ensure_collection(self) -> None:
        """创建集合或校验已有集合的向量维度。"""
        if not await self.client.has_collection(self.collection_name):
            await self.client.create_collection(
                collection_name=self.collection_name,
                schema=build_memory_schema(),
                index_params=build_memory_index_params(),
            )
            return
        desc = await self.client.describe_collection(self.collection_name)
        for field in desc.get("fields", []):
            if field.get("name") == "vector":
                dim = field.get("params", {}).get("dim")
                if dim is not None and int(dim) != EMBEDDING_DIMENSION:
                    raise RuntimeError(
                        f"Milvus collection {self.collection_name!r} vector dim={dim}, "
                        f"expected {EMBEDDING_DIMENSION}"
                    )

    async def upsert(self, memory, vector: list[float]) -> str:
        """写入一条客户记忆向量并返回 Milvus 主键。"""
        await self.ensure_collection()
        memory_id = str(memory.id)
        await self.client.delete(
            collection_name=self.collection_name,
            filter=f'memory_id == "{memory_id}"',
        )
        await self.client.insert(
            collection_name=self.collection_name,
            data=[
                {
                    "memory_id": memory_id,
                    "customer_id": memory.customer_id,
                    "memory_type": memory.memory_type,
                    "tag": memory.tag,
                    "content": memory.content,
                    "status": memory.status,
                    "vector": vector,
                }
            ],
        )
        return memory_id

    async def search(self, vector: list[float], customer_id: int, *, limit: int = 10) -> list[dict]:
        """按客户 ID 过滤向量查询，返回归一化的命中列表。"""
        await self.ensure_collection()
        raw = await self.client.search(
            collection_name=self.collection_name,
            data=[vector],
            limit=limit,
            filter=f"customer_id == {int(customer_id)}",
            output_fields=["memory_id", "customer_id", "memory_type", "tag", "content", "status"],
        )
        return self._normalize_hits(raw)

    @staticmethod
    def _normalize_hits(raw: Any) -> list[dict]:
        """将 pymilvus 返回结构收敛为 memory_id + distance 的扁平列表。"""
        hits: list[dict] = []
        for batch in raw or []:
            for hit in batch or []:
                if not isinstance(hit, dict):
                    continue
                entity = hit.get("entity") or {}
                memory_id = entity.get("memory_id") or hit.get("id")
                if memory_id is None:
                    continue
                try:
                    distance = float(hit.get("distance", 0.0))
                except (TypeError, ValueError):
                    distance = 0.0
                hits.append(
                    {
                        "memory_id": str(memory_id),
                        "distance": max(0.0, min(1.0, distance)),
                        "tag": entity.get("tag"),
                        "content": entity.get("content"),
                        "status": entity.get("status"),
                    }
                )
        return hits

    async def delete(self, memory_id: int | str) -> None:
        """删除一条客户记忆向量。"""
        await self.client.delete(
            collection_name=self.collection_name,
            filter=f'memory_id == "{memory_id}"',
        )


__all__ = [
    "CUSTOMER_MEMORY_COLLECTION",
    "MilvusMemoryStore",
    "build_memory_index_params",
    "build_memory_schema",
]