Mutual_Fund/service/memory/long_term.py

"""MySQL + Milvus + Neo4j 客户长期记忆同步服务。"""

from __future__ import annotations

from datetime import datetime
from inspect import isawaitable
from typing import Any, Callable

from tool.llm import llm
from tool.confidence import BaseConfidenceCalcTool

from repositories.memory_unit import MemoryUnitRepo
from service.memory.milvus_memory import MilvusMemoryStore
from service.memory.neo4j_memory import Neo4jMemoryStore
from service.memory.schemas import MemoryUnitDTO


class LongTermMemoryService:
    """以 MySQL 为主体事实源，向 Milvus 和 Neo4j 同步镜像。"""

    def __init__(
        self,
        *,
        milvus_store: MilvusMemoryStore | None = None,
        neo4j_store: Neo4jMemoryStore | None = None,
        repository_factory=MemoryUnitRepo,
        embedder: Callable[[str], Any] | None = None,
    ):
        self.milvus_store = milvus_store or MilvusMemoryStore()
        self.neo4j_store = neo4j_store or Neo4jMemoryStore()
        self.repository_factory = repository_factory
        self.embedder = embedder or llm.embed_one
        self.confidence_tool = BaseConfidenceCalcTool()

    async def save(self, db, memory: MemoryUnitDTO) -> tuple[MemoryUnitDTO, list[str]]:
        """保存主体并尽力同步两个外部索引，返回记忆和 warnings。"""
        repo = self.repository_factory(db)
        existing = await repo.find_exact(
            memory.customer_id,
            memory.memory_type.value,
            memory.tag,
        )
        warnings: list[str] = []
        if existing is not None:
            existing = await repo.merge_evidence(
                existing,
                content=memory.content,
                source=memory.source.value,
                evidence_count=max(memory.evidence_count or 1, 1),
                evidence_ref=memory.evidence_ref,
            )
            entity = existing
        else:
            # final_score 是召回重排阶段的临时分数，不属于 MySQL 主体字段。
            values = memory.model_dump(
                mode="json",
                exclude={"id", "milvus_id", "graph_node_id", "final_score"},
            )
            now = datetime.now()
            values["last_verified_at"] = values.get("last_verified_at") or now
            values["update_time"] = values.get("update_time") or now
            values["memory_type"] = memory.memory_type.value
            values["source"] = memory.source.value
            confidence_result, confidence_warning = self._calculate_confidence(memory)
            warnings.extend(confidence_warning)
            values.update(confidence_result)
            entity = await repo.add_memory(values)

        if existing is not None:
            confidence_result, confidence_warning = self._calculate_confidence(entity)
            warnings.extend(confidence_warning)
            await repo.update_sync_status(entity.id, **confidence_result)
            for key, value in confidence_result.items():
                setattr(entity, key, value)

        try:
            vector = self.embedder(entity.content)
            if isawaitable(vector):
                vector = await vector
            milvus_id = await self.milvus_store.upsert(entity, vector)
            await repo.update_sync_status(
                entity.id, milvus_id=milvus_id, milvus_sync_status="success"
            )
            entity.milvus_id = milvus_id
            entity.milvus_sync_status = "success"
        except Exception as exc:
            warnings.append(f"milvus_sync_failed:{type(exc).__name__}")
            await repo.update_sync_status(
                entity.id,
                milvus_sync_status="failed",
                sync_retry_count=(entity.sync_retry_count or 0) + 1,
                last_sync_error=str(exc)[:500],
            )

        try:
            graph_id = await self.neo4j_store.upsert(entity)
            await repo.update_sync_status(
                entity.id, graph_node_id=graph_id, neo4j_sync_status="success"
            )
            entity.graph_node_id = graph_id
            entity.neo4j_sync_status = "success"
        except Exception as exc:
            warnings.append(f"neo4j_sync_failed:{type(exc).__name__}")
            await repo.update_sync_status(
                entity.id,
                neo4j_sync_status="failed",
                sync_retry_count=(entity.sync_retry_count or 0) + 1,
                last_sync_error=str(exc)[:500],
            )
        return self._to_dto(entity), warnings

    def _calculate_confidence(self, memory) -> tuple[dict[str, Any], list[str]]:
        """计算记忆置信度；异常时强制降级为候选记忆。"""
        source = memory.source.value if hasattr(memory.source, "value") else memory.source
        memory_type = (
            memory.memory_type.value
            if hasattr(memory.memory_type, "value")
            else memory.memory_type
        )
        create_time = getattr(memory, "create_time", None)
        age_days = max(0, (datetime.now() - create_time).days) if create_time else 0
        try:
            result = self.confidence_tool.evaluate(
                tag=memory.tag,
                source=source,
                evidence_count=memory.evidence_count or 0,
                age_days=age_days,
                memory_type=memory_type,
            )
            result.pop("age_days", None)
            result.pop("threshold", None)
            result["confidence_update_time"] = datetime.now()
            return result, []
        except Exception as exc:
            return {
                "status": "candidate",
                "confidence_reason": "置信度计算失败，降级保存为候选记忆",
                "confidence_version": BaseConfidenceCalcTool.VERSION,
                "confidence_update_time": datetime.now(),
            }, [f"confidence_calculation_failed:{type(exc).__name__}"]

    async def recall(
        self,
        db,
        customer_id: int,
        *,
        memory_type: str | None = None,
        tag: str | None = None,
        limit: int = 100,
    ) -> tuple[list[MemoryUnitDTO], list[str]]:
        """按客户、类型、标签和有效期召回主体记忆。"""
        entities = await self.repository_factory(db).list_for_customer(
            customer_id, memory_type=memory_type, tag=tag, limit=limit
        )
        return [self._to_dto(entity) for entity in entities], []

    async def refresh_confidence(
        self, db, customer_id: int, *, limit: int = 500
    ) -> int:
        """按当前时间刷新有效记忆的置信度并持久化结果。"""
        repo = self.repository_factory(db)
        entities = await repo.list_for_confidence_refresh(customer_id, limit=limit)
        refreshed = 0
        for entity in entities:
            confidence_result, _ = self._calculate_confidence(entity)
            await repo.update_confidence(entity.id, **confidence_result)
            refreshed += 1
        return refreshed

    async def retry_pending(self, db, *, limit: int = 100) -> dict[str, int]:
        """重试 MySQL 中缺少外部索引或同步失败的记忆。"""
        entities = await self.repository_factory(db).list_pending_sync(limit)
        success = 0
        failed = 0
        for entity in entities:
            dto = self._to_dto(entity)
            _, warnings = await self.save(db, dto)
            if warnings:
                failed += 1
            else:
                success += 1
        return {"success": success, "failed": failed}

    @staticmethod
    def _to_dto(entity) -> MemoryUnitDTO:
        """将 ORM 实体转换为跨层 DTO。"""
        data = {
            "id": entity.id,
            "customer_id": entity.customer_id,
            "session_id": entity.session_id,
            "agent_run_id": entity.agent_run_id,
            "memory_type": entity.memory_type,
            "tag": entity.tag,
            "content": entity.content,
            "info_type": entity.info_type,
            "source": entity.source,
            "evidence_ref": entity.evidence_ref or [],
            "source_confidence": float(entity.source_confidence or 0),
            "confidence": float(entity.confidence or 0),
            "historical_accuracy": float(entity.historical_accuracy or 0),
            "confidence_version": getattr(entity, "confidence_version", None),
            "confidence_reason": getattr(entity, "confidence_reason", None),
            "confidence_update_time": getattr(entity, "confidence_update_time", None),
            "evidence_count": entity.evidence_count or 0,
            "recall_count": entity.recall_count or 0,
            "status": entity.status,
            "valid_from": entity.valid_from,
            "valid_until": entity.valid_until,
            "last_verified_at": entity.last_verified_at,
            "milvus_id": entity.milvus_id,
            "graph_node_id": entity.graph_node_id,
        }
        return MemoryUnitDTO.model_validate(data)


__all__ = ["LongTermMemoryService"]