hg-agents/backend/modules/memory/manager.py

import json
import asyncio
import logging
from datetime import datetime
from redis.asyncio import Redis
from config import settings

logger = logging.getLogger(__name__)

_memory_manager: "MemoryManager | None" = None


def get_memory_manager() -> "MemoryManager":
    global _memory_manager
    if _memory_manager is None:
        raise RuntimeError("MemoryManager 未初始化，请先调用 init_memory_manager()")
    return _memory_manager


async def init_memory_manager():
    global _memory_manager
    redis = Redis.from_url(settings.REDIS_URL, decode_responses=True)
    await redis.ping()
    _memory_manager = MemoryManager(redis)


class MemoryManager:
    KEY_PREFIX = "mem"
    DEFAULT_TTL = 604800
    SESSION_INDEX_TTL = 2592000
    MAX_HISTORY = 40

    def __init__(self, redis: Redis):
        self.redis = redis

    async def inject_memory(
        self,
        user_id: str,
        flow_id: str,
        session_id: str,
        context: dict,
    ):
        messages = await self._get_recent_messages(user_id, flow_id, session_id)
        summary = await self._get_summary(user_id, flow_id, session_id)

        context["_memory_context"] = {
            "recent_messages": list(reversed(messages)),
            "summary": summary,
            "session_id": session_id,
        }

    async def record_exchange(
        self,
        user_id: str,
        flow_id: str,
        session_id: str,
        user_msg: str,
        assistant_msg: str,
        flow_name: str = "",
    ):
        key = self._msg_key(user_id, flow_id, session_id)
        ts = datetime.utcnow().isoformat()

        try:
            async with self.redis.pipeline() as pipe:
                pipe.lpush(key,
                    json.dumps({"role": "assistant", "content": assistant_msg, "ts": ts}, ensure_ascii=False),
                    json.dumps({"role": "user", "content": user_msg, "ts": ts}, ensure_ascii=False),
                )
                pipe.ltrim(key, 0, self.MAX_HISTORY - 1)
                pipe.expire(key, self.DEFAULT_TTL)

                pipe.hset(self._meta_key(user_id, flow_id, session_id), mapping={
                    "flow_name": flow_name,
                    "last_active_at": ts,
                })
                pipe.expire(self._meta_key(user_id, flow_id, session_id), self.DEFAULT_TTL)

                pipe.sadd(f"{self.KEY_PREFIX}:{user_id}:sessions", session_id)
                pipe.expire(f"{self.KEY_PREFIX}:{user_id}:sessions", self.SESSION_INDEX_TTL)

                await pipe.execute()
        except Exception as e:
            logger.warning(f"记录记忆失败: {e}")

        asyncio.create_task(self._maybe_summarize(user_id, flow_id, session_id))

    async def get_conversation_history(
        self, user_id: str, flow_id: str, session_id: str, limit: int = 20
    ) -> list[dict]:
        messages = await self._get_recent_messages(user_id, flow_id, session_id, limit)
        return list(reversed(messages))

    async def delete_session(self, user_id: str, session_id: str):
        try:
            patterns = await self.redis.keys(f"{self.KEY_PREFIX}:{user_id}:*:{session_id}:*")
            async with self.redis.pipeline() as pipe:
                if patterns:
                    pipe.delete(*patterns)
                pipe.srem(f"{self.KEY_PREFIX}:{user_id}:sessions", session_id)
                await pipe.execute()
        except Exception as e:
            logger.warning(f"清除记忆失败: {e}")

    async def list_user_sessions(self, user_id: str) -> list[dict]:
        try:
            session_ids = await self.redis.smembers(f"{self.KEY_PREFIX}:{user_id}:sessions")
        except Exception:
            return []

        sessions = []
        for sid in session_ids:
            try:
                keys = await self.redis.keys(f"{self.KEY_PREFIX}:{user_id}:*:{sid}:meta")
                for k in keys:
                    meta = await self.redis.hgetall(k)
                    parts = k.split(":")
                    flow_id = parts[2] if len(parts) > 2 else ""
                    sessions.append({
                        "session_id": sid,
                        "flow_id": flow_id,
                        "flow_name": meta.get("flow_name", ""),
                        "last_active_at": meta.get("last_active_at", ""),
                    })
            except Exception:
                continue

        return sorted(sessions, key=lambda s: s.get("last_active_at", ""), reverse=True)

    async def _get_recent_messages(
        self, user_id: str, flow_id: str, session_id: str, limit: int = None
    ) -> list[dict]:
        limit = limit or self.MAX_HISTORY
        try:
            key = self._msg_key(user_id, flow_id, session_id)
            raw = await self.redis.lrange(key, 0, limit - 1)
            result = []
            for m in raw:
                try:
                    result.append(json.loads(m))
                except json.JSONDecodeError:
                    continue
            return result
        except Exception:
            return []

    async def _get_summary(self, user_id: str, flow_id: str, session_id: str) -> str:
        try:
            key = f"{self.KEY_PREFIX}:{user_id}:{flow_id}:{session_id}:summary"
            val = await self.redis.get(key)
            return val or ""
        except Exception:
            return ""

    async def _maybe_summarize(self, user_id: str, flow_id: str, session_id: str):
        try:
            key = self._msg_key(user_id, flow_id, session_id)
            count = await self.redis.llen(key)
            if count < 30:
                return

            summary_key = f"{self.KEY_PREFIX}:{user_id}:{flow_id}:{session_id}:summary"
            existing = await self.redis.get(summary_key)
            if existing:
                return

            recent = await self._get_recent_messages(user_id, flow_id, session_id, 20)
            dialogue = "\n".join(
                f"{m['role']}: {m['content'][:500]}" for m in reversed(recent[:10])
            )

            import httpx
            api_base = settings.LLM_API_BASE.rstrip("/")
            async with httpx.AsyncClient(timeout=30) as client:
                resp = await client.post(
                    f"{api_base}/chat/completions",
                    json={
                        "model": settings.LLM_MODEL,
                        "messages": [{
                            "role": "user",
                            "content": f"请用一段话简要总结以下对话的关键内容。保留人名、任务、决策、时间等关键信息。\n\n{dialogue}"
                        }],
                        "max_tokens": 200,
                    },
                    headers={"Authorization": f"Bearer {settings.LLM_API_KEY}"},
                )
                data = resp.json()
                summary = data.get("choices", [{}])[0].get("message", {}).get("content", "")

                if summary:
                    await self.redis.setex(summary_key, 2592000, summary)
        except Exception:
            pass

    @staticmethod
    def _msg_key(user_id: str, flow_id: str, session_id: str) -> str:
        return f"mem:{user_id}:{flow_id}:{session_id}:messages"

    @staticmethod
    def _meta_key(user_id: str, flow_id: str, session_id: str) -> str:
        return f"mem:{user_id}:{flow_id}:{session_id}:meta"