Mutual_Fund/service/memory/short_term.py

"""Redis 短期会话记忆。"""

from __future__ import annotations

import asyncio
import json
import time
import uuid
from datetime import datetime, timezone
from inspect import isawaitable
from typing import Any, Callable

from config.database.redis import client as redis_client

from .schemas import ShortTermMessage


class ShortTermMemoryError(RuntimeError):
    """短期记忆操作失败。"""


class SessionExpiredError(ShortTermMemoryError):
    """会话已超过最长生命周期。"""


async def _config(config_getter, key: str, default):
    """读取项目配置，并将数据库配置值转换为默认值类型。"""
    value = config_getter(key, str(default))
    if isawaitable(value):
        value = await value
    return type(default)(value)


class ShortTermMemory:
    """管理当前客服会话的 Redis 消息、Token 预算和生命周期。"""

    MESSAGE_KEY = "session:{session_id}:messages"
    META_KEY = "session:{session_id}:meta"
    ALLOWED_ROLES = frozenset({"user", "assistant", "system"})

    def __init__(
        self,
        redis=None,
        *,
        config_getter=None,
        token_counter: Callable[[str], int] | None = None,
        clock=time.time,
        fail_soft: bool = True,
    ):
        """创建短期记忆服务，默认复用项目级 Redis 客户端。"""
        self.redis = redis or redis_client()
        self.config_getter = config_getter or (lambda _key, default: default)
        self.token_counter = token_counter or (lambda text: max(1, len(text) // 4))
        self.clock = clock
        self.fail_soft = fail_soft
        self.last_warnings: list[str] = []
        self._locks: dict[str, asyncio.Lock] = {}

    @classmethod
    def message_key(cls, session_id: str) -> str:
        """生成会话消息列表 Key。"""
        return cls.MESSAGE_KEY.format(session_id=session_id)

    @classmethod
    def meta_key(cls, session_id: str) -> str:
        """生成会话元数据 Hash Key。"""
        return cls.META_KEY.format(session_id=session_id)

    def _lock(self, session_id: str) -> asyncio.Lock:
        """获取进程内会话锁，避免并发截断互相覆盖。"""
        return self._locks.setdefault(session_id, asyncio.Lock())

    async def append_message(
        self,
        session_id: str,
        role: str,
        content: str,
        *,
        message_id: str | None = None,
        agent_run_id: str | None = None,
        tool_calls: list[dict[str, Any]] | None = None,
    ) -> ShortTermMessage | None:
        """追加消息、刷新 TTL，并按 Token 预算从旧到新截断。"""
        self.last_warnings = []
        if role not in self.ALLOWED_ROLES:
            raise ValueError(f"非法消息角色: {role}")
        if not isinstance(content, str) or not content.strip():
            raise ValueError("消息内容不能为空")

        message = ShortTermMessage(
            message_id=message_id or uuid.uuid4().hex,
            session_id=session_id,
            role=role,
            content=content,
            token_count=self.token_counter(content),
            agent_run_id=agent_run_id,
            tool_calls=tool_calls or [],
            create_time=datetime.fromtimestamp(self.clock(), tz=timezone.utc),
        )
        try:
            async with self._lock(session_id):
                await self._ensure_meta(session_id)
                await self.redis.rpush(
                    self.message_key(session_id), self._serialize(message)
                )
                await self._touch(session_id)
                await self._trim(session_id)
            return message
        except SessionExpiredError:
            raise
        except Exception as exc:
            return self._degrade("append_message", exc)

    async def load_messages(self, session_id: str) -> list[ShortTermMessage]:
        """按写入顺序读取当前会话消息，并刷新空闲 TTL。"""
        self.last_warnings = []
        try:
            if not await self._session_is_active(session_id):
                return []
            raw_messages = await self.redis.lrange(self.message_key(session_id), 0, -1)
            await self._touch(session_id)
            return [self._deserialize(raw) for raw in raw_messages if raw]
        except Exception as exc:
            return self._degrade("load_messages", exc) or []

    async def get_message_count(self, session_id: str) -> int:
        """返回当前会话消息数量。"""
        return len(await self.load_messages(session_id))

    async def get_token_count(self, session_id: str) -> int:
        """返回当前会话消息的 Token 估算总数。"""
        return sum(message.token_count for message in await self.load_messages(session_id))

    async def clear_session(self, session_id: str) -> None:
        """清理会话消息和元数据。"""
        self.last_warnings = []
        try:
            await self.redis.delete(self.message_key(session_id), self.meta_key(session_id))
        except Exception as exc:
            self._degrade("clear_session", exc)

    async def _ensure_meta(self, session_id: str) -> None:
        """初始化会话元数据，并固定 24 小时绝对过期时间。"""
        meta_key = self.meta_key(session_id)
        meta = await self.redis.hgetall(meta_key)
        now = self.clock()
        if meta:
            absolute_expire_at = float(meta.get("absolute_expire_at", now))
            if absolute_expire_at <= now:
                await self.clear_session(session_id)
                raise SessionExpiredError("会话已超过最长生命周期")
            return
        max_lifetime = await _config(
            self.config_getter, "agent.customer.session.max_lifetime", 86400
        )
        await self.redis.hset(
            meta_key,
            mapping={
                "created_at": str(now),
                "absolute_expire_at": str(now + max_lifetime),
            },
        )
        await self.redis.expire(meta_key, max_lifetime)

    async def _session_is_active(self, session_id: str) -> bool:
        """检查会话元数据是否存在且未达到绝对过期时间。"""
        meta = await self.redis.hgetall(self.meta_key(session_id))
        if not meta:
            return False
        if float(meta.get("absolute_expire_at", 0)) <= self.clock():
            await self.clear_session(session_id)
            return False
        return True

    async def _touch(self, session_id: str) -> None:
        """刷新空闲 TTL，但不超过绝对过期时间。"""
        meta = await self.redis.hgetall(self.meta_key(session_id))
        if not meta:
            return
        remaining = int(float(meta["absolute_expire_at"]) - self.clock())
        if remaining <= 0:
            await self.clear_session(session_id)
            raise SessionExpiredError("会话已超过最长生命周期")
        idle_ttl = await _config(
            self.config_getter, "agent.customer.session.ttl", 1800
        )
        archive_grace = await _config(
            self.config_getter, "agent.customer.session.archive_grace", 60
        )
        ttl = min(idle_ttl + archive_grace, remaining)
        await self.redis.expire(self.message_key(session_id), ttl)
        await self.redis.expire(self.meta_key(session_id), remaining)

    async def _trim(self, session_id: str) -> None:
        """保留最新消息，确保最新一条消息不会因超预算被删除。"""
        limit = await _config(
            self.config_getter, "agent.customer.session_max_token", 4096
        )
        key = self.message_key(session_id)
        raw_messages = [raw for raw in await self.redis.lrange(key, 0, -1) if raw]
        total = 0
        keep_from = len(raw_messages)
        for index in range(len(raw_messages) - 1, -1, -1):
            message = self._deserialize(raw_messages[index])
            total += message.token_count
            keep_from = index
            if total > limit:
                keep_from = index + 1
                break
        if raw_messages and keep_from == len(raw_messages):
            keep_from = len(raw_messages) - 1
        await self.redis.ltrim(key, keep_from, -1)

    @staticmethod
    def _serialize(message: ShortTermMessage) -> str:
        """将消息转换为 Redis List 中的 JSON 字符串。"""
        return json.dumps(message.model_dump(mode="json"), ensure_ascii=False)

    @staticmethod
    def _deserialize(raw: str | bytes) -> ShortTermMessage:
        """将 Redis JSON 字符串恢复为消息 DTO。"""
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")
        return ShortTermMessage.model_validate(json.loads(raw))

    def _degrade(self, operation: str, exc: Exception):
        """记录降级原因；fail_soft 模式下不让 Redis 故障击穿客服请求。"""
        warning = f"short_term_{operation}_degraded:{type(exc).__name__}"
        self.last_warnings = [warning]
        if not self.fail_soft:
            raise ShortTermMemoryError(warning) from exc
        return None


__all__ = ["SessionExpiredError", "ShortTermMemory", "ShortTermMemoryError"]