Mutual_Fund/repositories/memory_unit.py

"""memory_unit 长期记忆仓储。"""

from __future__ import annotations

from datetime import datetime

from sqlalchemy import select

from model.memory_unit import MemoryUnit
from repositories.base import BaseRepository


ACTIVE_STATUSES = ("candidate", "confirmed")
INACTIVE_STATUSES = ("expired", "rejected", "archived")


class MemoryUnitRepo(BaseRepository):
    """负责 MySQL 主体记忆的隔离、去重和同步状态管理。"""

    model = MemoryUnit

    async def add_memory(self, values: dict) -> MemoryUnit:
        """先写入 MySQL 主体并返回数据库 ID。"""
        obj = MemoryUnit(**values)
        self.db.add(obj)
        await self.db.commit()
        await self.db.refresh(obj)
        return obj

    async def find_exact(
        self, customer_id: int, memory_type: str, tag: str
    ) -> MemoryUnit | None:
        """按客户、类型和标签查找当前有效记忆。"""
        return await self.db.scalar(
            select(MemoryUnit).where(
                MemoryUnit.customer_id == customer_id,
                MemoryUnit.memory_type == memory_type,
                MemoryUnit.tag == tag,
                MemoryUnit.status.not_in(INACTIVE_STATUSES),
            )
        )

    async def list_for_customer(
        self,
        customer_id: int,
        *,
        memory_type: str | None = None,
        tag: str | None = None,
        now: datetime | None = None,
        limit: int = 100,
    ) -> list[MemoryUnit]:
        """查询默认可召回的客户记忆，过滤拒绝、归档和过期记录。"""
        now = now or datetime.now()
        conditions = [
            MemoryUnit.customer_id == customer_id,
            MemoryUnit.status.in_(ACTIVE_STATUSES),
            (MemoryUnit.valid_until.is_(None) | (MemoryUnit.valid_until > now)),
        ]
        if memory_type:
            conditions.append(MemoryUnit.memory_type == memory_type)
        if tag:
            conditions.append(MemoryUnit.tag == tag)
        statement = (
            select(MemoryUnit)
            .where(*conditions)
            .order_by(MemoryUnit.update_time.desc(), MemoryUnit.id.desc())
            .limit(limit)
        )
        return list((await self.db.scalars(statement)).all())

    async def list_for_customer_by_ids(
        self,
        customer_id: int,
        ids: list[int],
        *,
        memory_type: str | None = None,
        tag: str | None = None,
        now: datetime | None = None,
    ) -> list[MemoryUnit]:
        """按 ID 集合取回本人有效记忆，用于向量召回命中后的主体回表。"""
        if not ids:
            return []
        now = now or datetime.now()
        conditions = [
            MemoryUnit.customer_id == customer_id,
            MemoryUnit.id.in_(ids),
            MemoryUnit.status.in_(ACTIVE_STATUSES),
            (MemoryUnit.valid_until.is_(None) | (MemoryUnit.valid_until > now)),
        ]
        if memory_type:
            conditions.append(MemoryUnit.memory_type == memory_type)
        if tag:
            conditions.append(MemoryUnit.tag == tag)
        statement = (
            select(MemoryUnit)
            .where(*conditions)
            .order_by(MemoryUnit.update_time.desc(), MemoryUnit.id.desc())
        )
        return list((await self.db.scalars(statement)).all())

    async def merge_evidence(
        self,
        memory: MemoryUnit,
        *,
        content: str | None = None,
        source: str | None = None,
        evidence_count: int = 1,
        evidence_ref: list[dict] | None = None,
    ) -> MemoryUnit:
        """更新最新内容并合并证据，不改变客户隔离范围。"""
        if content:
            memory.content = content
        if source:
            memory.source = source
        memory.evidence_count = (memory.evidence_count or 0) + evidence_count
        if evidence_ref:
            memory.evidence_ref = [*(memory.evidence_ref or []), *evidence_ref]
        memory.last_verified_at = datetime.now()
        memory.update_time = datetime.now()
        await self.db.commit()
        await self.db.refresh(memory)
        return memory

    async def list_for_confidence_refresh(
        self, customer_id: int, limit: int = 500
    ) -> list[MemoryUnit]:
        """返回需要重新计算时间衰减置信度的有效记忆。"""
        statement = (
            select(MemoryUnit)
            .where(
                MemoryUnit.customer_id == customer_id,
                MemoryUnit.status.in_(ACTIVE_STATUSES),
            )
            .order_by(MemoryUnit.id)
            .limit(limit)
        )
        return list((await self.db.scalars(statement)).all())

    async def update_confidence(self, memory_id: int, **values) -> None:
        """只更新置信度字段，不改变记忆内容和证据。"""
        memory = await self.db.get(MemoryUnit, memory_id)
        if memory is None:
            return
        for key, value in values.items():
            setattr(memory, key, value)
        await self.db.commit()

    async def update_sync_status(self, memory_id: int, **values) -> None:
        """更新向量/图谱索引 ID 和同步状态。"""
        memory = await self.db.get(MemoryUnit, memory_id)
        if memory is None:
            return
        for key, value in values.items():
            setattr(memory, key, value)
        memory.update_time = datetime.now()
        await self.db.commit()

    async def list_pending_sync(self, limit: int = 100) -> list[MemoryUnit]:
        """返回存在失败或空索引的记忆，供重试入口使用。"""
        statement = (
            select(MemoryUnit)
            .where(
                (MemoryUnit.milvus_id.is_(None) | (MemoryUnit.milvus_sync_status == "failed"))
                | (MemoryUnit.graph_node_id.is_(None) | (MemoryUnit.neo4j_sync_status == "failed"))
            )
            .order_by(MemoryUnit.id)
            .limit(limit)
        )
        return list((await self.db.scalars(statement)).all())


__all__ = ["MemoryUnitRepo"]