MoviePilot/app/agent/memory/__init__.py

"""
对话记忆管理器
"""

import asyncio
from datetime import datetime
from typing import Dict, List, Optional

from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict
from pydantic import BaseModel, ConfigDict, Field, field_serializer

from app.agent.history.message import legacy_identity
from app.agent.tools.result import messages_for_persistence
from app.application.messaging.chat import (
    AgentChatPersistenceService,
    AgentChatRecord,
    AgentChatService,
    get_configured_agent_chat_persistence,
    get_configured_agent_chat_service,
)
from app.runtime.log import logger
from app.runtime.settings import get_runtime_setting


class ConversationMemory(BaseModel):
    """
    对话记忆模型

    消息字段依赖 langchain_core，因此模型归属智能体记忆模块；普通 schema 导入不加载 LangChain。
    """

    session_id: str = Field(description="会话ID")
    user_id: Optional[str] = Field(default=None, description="用户ID")
    owner_user_id: Optional[str] = Field(default=None, description="持久化会话原始归属，缓存复用时再次校验")
    owner_username: Optional[str] = Field(default=None, description="无 ID 旧记录的用户名归属")
    messages: List[BaseMessage] = Field(default_factory=list, description="消息列表")
    updated_at: datetime = Field(default_factory=datetime.now, description="更新时间")

    model_config = ConfigDict()

    @field_serializer("updated_at", when_used="json")
    def serialize_datetime(self, value: datetime) -> str:
        """
        把记忆更新时间序列化为标准日期字符串。
        """
        return value.isoformat()


class MemoryManager:
    """
    对话记忆管理器
    """

    def __init__(
        self,
        chat: AgentChatService | None = None,
        persistence: AgentChatPersistenceService | None = None,
    ) -> None:
        """
        创建记忆缓存，并保存组合根显式注入的会话能力。
        """
        self._chat = chat
        self._persistence = persistence
        # 内存中的会话记忆缓存
        self.memory_cache: Dict[str, ConversationMemory] = {}
        # 内存缓存清理任务
        self.cleanup_task: Optional[asyncio.Task[None]] = None

    def configure(
        self,
        chat: AgentChatService,
        persistence: AgentChatPersistenceService,
    ) -> None:
        """
        在 Agent 启动前绑定唯一会话查询与写入服务。
        """
        self._chat = chat
        self._persistence = persistence

    def _chat_service(self) -> AgentChatService:
        """
        返回显式注入服务，兼容测试未装配时使用既有应用服务。
        """
        return self._chat or get_configured_agent_chat_service()

    def _chat_persistence(self) -> AgentChatPersistenceService:
        """
        返回显式注入写服务，兼容测试未装配时使用既有应用服务。
        """
        return self._persistence or get_configured_agent_chat_persistence()

    def initialize(self) -> None:
        """
        初始化记忆管理器
        """
        try:
            if self.cleanup_task and not self.cleanup_task.done():
                return
            # 启动内存缓存清理任务（Redis通过TTL自动过期）
            self.cleanup_task = asyncio.create_task(
                self._cleanup_expired_memories()
            )
            logger.info("对话记忆管理器初始化完成")

        except Exception as e:
            logger.warning(f"Redis连接失败，将使用内存存储: {e}")

    async def close(self) -> None:
        """
        关闭记忆管理器
        """
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
            self.cleanup_task = None

        logger.info("对话记忆管理器已关闭")

    @staticmethod
    def _get_memory_key(session_id: str, user_id: Optional[str]) -> str:
        """
        计算内存Key
        """
        return f"{user_id}:{session_id}" if user_id else session_id

    def get_memory(self, session_id: str, user_id: str) -> Optional[ConversationMemory]:
        """
        获取内存中的记忆
        """
        cache_key = self._get_memory_key(session_id, user_id)
        return self.memory_cache.get(cache_key)

    @staticmethod
    def _chat_lookup_params(
        session_id: str,
        user_id: str,
    ) -> tuple[dict[str, str], ...]:
        """
        返回用户专属会话优先、旧无用户会话兜底的查询顺序。
        """
        return (
            {"session_id": session_id, "user_id": user_id},
            {"session_id": session_id},
        )

    def _restore_agent_messages(
        self,
        *,
        session_id: str,
        user_id: str,
        chat: Optional[AgentChatRecord],
        allow_cross_owner: bool = False,
        legacy_username: Optional[str] = None,
    ) -> List[BaseMessage]:
        """
        仅恢复当前用户或无归属旧记录，再清理工具图像并回填缓存。
        """
        if not chat or not chat.agent_messages:
            return []
        # 无用户条件的兼容查询只用于无归属旧记录，不能扩大到其他明确归属账号。
        owner_id = getattr(chat, "user_id", None)
        owner_username = getattr(chat, "username", None)
        foreign_owner = (
            str(owner_id) != str(user_id) if owner_id
            else bool(owner_username and owner_username != legacy_username)
        )
        if foreign_owner and not allow_cross_owner:
            return []
        try:
            messages = messages_for_persistence(messages_from_dict(chat.agent_messages))
            for position, message in enumerate(messages):
                if not message.id:
                    message.id = legacy_identity(session_id, position, message)
        except Exception as e:
            logger.debug(f"恢复持久化Agent消息失败: {e}")
            return []
        memory = ConversationMemory(
            session_id=session_id,
            user_id=user_id,
            owner_user_id=owner_id,
            owner_username=owner_username,
            messages=messages,
        )
        self.save_memory(memory)
        return memory.messages

    def get_agent_messages(
            self, session_id: str, user_id: str, *, allow_cross_owner: bool = False,
            legacy_username: Optional[str] = None,
    ) -> List[BaseMessage]:
        """
        为Agent获取最近的消息。

        优先使用内存缓存，缓存不存在时从数据库恢复上一轮持久化的原始 messages。
        """
        memory = self.get_memory(session_id, user_id)
        if memory:
            return self._authorized_cached_messages(memory, user_id, allow_cross_owner, legacy_username)

        try:
            service = self._chat_service()
            chat = None
            for lookup_params in self._chat_lookup_params(session_id, user_id):
                chat = service.get_sync(**lookup_params)
                if chat:
                    break
        except Exception as e:
            logger.debug(f"读取持久化Agent会话失败: {e}")
            return []
        return self._restore_agent_messages(
            session_id=session_id,
            user_id=user_id,
            chat=chat,
            allow_cross_owner=allow_cross_owner,
            legacy_username=legacy_username,
        )

    @staticmethod
    def _authorized_cached_messages(
        memory: ConversationMemory, user_id: str, allow_cross_owner: bool,
        legacy_username: Optional[str],
    ) -> List[BaseMessage]:
        """
        恢复授权不会永久附着于缓存，每次跨归属读取都需要显式许可。
        """
        foreign_owner = (
            memory.owner_user_id != str(user_id) if memory.owner_user_id
            else bool(memory.owner_username and memory.owner_username != legacy_username)
        )
        if foreign_owner and not allow_cross_owner:
            return []
        return memory.messages

    async def async_get_agent_messages(
        self, session_id: str, user_id: str, *, allow_cross_owner: bool = False,
        legacy_username: Optional[str] = None,
    ) -> List[BaseMessage]:
        """
        异步恢复 Agent 消息，查询与会话应用服务保持同一异步端口。
        """
        memory = self.get_memory(session_id, user_id)
        if memory:
            return self._authorized_cached_messages(memory, user_id, allow_cross_owner, legacy_username)

        try:
            service = self._chat_service()
            chat = None
            for lookup_params in self._chat_lookup_params(session_id, user_id):
                chat = await service.get(**lookup_params)
                if chat:
                    break
        except Exception as e:
            logger.debug(f"读取持久化Agent会话失败: {e}")
            return []
        return self._restore_agent_messages(
            session_id=session_id,
            user_id=user_id,
            chat=chat,
            allow_cross_owner=allow_cross_owner,
            legacy_username=legacy_username,
        )

    def _update_agent_messages(
        self,
        *,
        session_id: str,
        user_id: str,
        messages: List[BaseMessage],
    ) -> None:
        """
        统一更新同步和异步写路径共享的内存消息状态。
        """
        memory = self.get_memory(session_id, user_id)
        if not memory:
            memory = ConversationMemory(session_id=session_id, user_id=user_id)
        memory.messages = messages
        memory.updated_at = datetime.now()
        self.save_memory(memory)

    def save_agent_messages(
            self, session_id: str, user_id: str, messages: List[BaseMessage]
    ) -> None:
        """
        当前会话在缓存保留原图，持久化副本只保存工具图片的失效说明。
        """
        self._update_agent_messages(
            session_id=session_id,
            user_id=user_id,
            messages=messages,
        )
        try:
            self._chat_service().save_agent_messages(
                session_id=session_id,
                user_id=user_id,
                messages=messages_to_dict(messages_for_persistence(messages)),
            )
        except Exception as e:
            logger.debug(f"持久化Agent消息失败: {e}")

    async def async_save_agent_messages(
        self, session_id: str, user_id: str, messages: List[BaseMessage]
    ) -> None:
        """
        缓存保留原图，去除工具图片的副本经有界数据库 worker 持久化。
        """
        self._update_agent_messages(
            session_id=session_id,
            user_id=user_id,
            messages=messages,
        )
        try:
            persistence = self._chat_persistence()
            await persistence.async_save_agent_messages(
                session_id=session_id,
                user_id=user_id,
                messages=messages_to_dict(messages_for_persistence(messages)),
            )
        except Exception as e:
            logger.debug(f"持久化Agent消息失败: {e}")

    def save_memory(self, memory: ConversationMemory) -> None:
        """
        保存记忆到内存缓存

        注意：Redis中的记忆通过TTL机制自动过期，这里只更新内存缓存，Redis会在下次访问时自动过期
        """
        cache_key = self._get_memory_key(memory.session_id, memory.user_id)
        self.memory_cache[cache_key] = memory

    def clear_memory(self, session_id: str, user_id: str) -> None:
        """
        清空会话记忆
        """
        cache_key = self._get_memory_key(session_id, user_id)
        if cache_key in self.memory_cache:
            del self.memory_cache[cache_key]

        logger.info(f"会话记忆已清空: session_id={session_id}, user_id={user_id}")

    async def _cleanup_expired_memories(self) -> None:
        """
        清理内存中过期记忆的后台任务

        注意：Redis中的记忆通过TTL机制自动过期，这里只清理内存缓存
        """
        while True:
            try:
                # 每小时清理一次
                await asyncio.sleep(3600)

                current_time = datetime.now()
                expired_sessions = []

                # 只检查内存缓存中的过期记忆
                # Redis中的记忆会通过TTL自动过期，无需手动处理
                for cache_key, memory in self.memory_cache.items():
                    if (
                            current_time - memory.updated_at
                    ).days > get_runtime_setting('LLM_MEMORY_RETENTION_DAYS'):
                        expired_sessions.append(cache_key)

                # 只清理内存缓存，不删除Redis中的键（Redis会自动过期）
                for cache_key in expired_sessions:
                    if cache_key in self.memory_cache:
                        del self.memory_cache[cache_key]

                if expired_sessions:
                    logger.info(f"清理了{len(expired_sessions)}个过期内存会话记忆")

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"清理记忆时发生错误: {e}")


memory_manager = MemoryManager()