ailine/app/memory/mem0_client.py

"""
Mem0 记忆层客户端封装模块
负责 Mem0 的初始化、检索和存储
"""

import asyncio
from typing import Optional, List, Dict, Any
from mem0 import AsyncMemory

from app.config import QDRANT_URL, QDRANT_COLLECTION_NAME, LLAMACPP_EMBEDDING_URL, LLAMACPP_API_KEY
from app.logger import info, warning, error


class Mem0Client:
    """Mem0 异步客户端封装类"""

    def __init__(self, llm_instance):
        """
        初始化 Mem0 客户端

        Args:
            llm_instance: LangChain LLM 实例（用于事实提取）
        """
        self.llm = llm_instance
        self.mem0: Optional[AsyncMemory] = None
        self._initialized = False

    async def initialize(self):
        """异步初始化 Mem0 客户端，并进行实际连接测试"""
        if self._initialized:
            return

        try:
            # Mem0 配置
            config = {
                "vector_store": {
                    "provider": "qdrant",
                    "config": {
                        "url": QDRANT_URL,                     # 直接使用完整 URL
                        "collection_name": QDRANT_COLLECTION_NAME,
                        "embedding_model_dims": 768,
                    }
                },
                "llm": {
                    "provider": "langchain",
                    "config": {
                        "model": self.llm
                    }
                },
                "embedder": {
                    "provider": "openai",
                    "config": {
                        "model": "embeddinggemma-300M-Q8_0",
                        "api_key": LLAMACPP_API_KEY,
                        "openai_base_url": LLAMACPP_EMBEDDING_URL,
                    },
                },
                "version": "v1.1"
            }

            self.mem0 = AsyncMemory.from_config(config)
            info("✅ Mem0 配置加载成功，开始连接测试...")

            # 实际连接测试：调用一次 search 确保 Qdrant 和 Embedding 都可达
            await asyncio.wait_for(
                self.mem0.search("ping", user_id="test", limit=1),
                timeout=60.0
            )
            info("✅ Mem0 实际连接测试成功，初始化完成")
            self._initialized = True

        except asyncio.TimeoutError:
            error("❌ Mem0 连接测试超时 (10s)，请检查 Qdrant 或 Embedding 服务响应")
            self.mem0 = None
            self._initialized = False
        except Exception as e:
            error(f"❌ Mem0 初始化或连接测试失败: {e}")
            import traceback
            error(f"详细错误信息:\n{traceback.format_exc()}")
            self.mem0 = None
            self._initialized = False

    async def search_memories(self, query: str, user_id: str, limit: int = 5) -> List[str]:
        """
        检索相关记忆

        Args:
            query: 查询文本
            user_id: 用户 ID
            limit: 返回结果数量限制

        Returns:
            List[str]: 记忆事实列表
        """
        if not self.mem0:
            warning("⚠️ Mem0 未初始化，跳过记忆检索")
            return []

        try:
            memories = await asyncio.wait_for(
                self.mem0.search(query, user_id=user_id, limit=limit),
                timeout=30.0
            )

            if memories and "results" in memories:
                facts = [m["memory"] for m in memories["results"] if m.get("memory")]
                if facts:
                    info(f"🔍 [记忆检索] Mem0 返回 {len(facts)} 条记忆")
                    return facts

            info("🔍 [记忆检索] 未找到相关记忆")
            return []

        except asyncio.TimeoutError:
            warning("⚠️ Mem0 检索超时 (30s)，跳过本次记忆检索")
            return []
        except Exception as e:
            warning(f"⚠️ Mem0 检索失败: {e}")
            return []

    async def add_memories(self, messages: List[Dict[str, str]], user_id: str) -> bool:
        """
        添加记忆（自动提取事实并存储）

        Args:
            messages: 消息列表，格式为 [{"role": "user/assistant/system", "content": "..."}]
            user_id: 用户 ID

        Returns:
            bool: 是否成功
        """
        if not self.mem0:
            warning("⚠️ Mem0 未初始化，跳过记忆添加")
            return False

        try:
            await asyncio.wait_for(
                self.mem0.add(
                    messages,
                    user_id=user_id,
                    metadata={"type": "conversation"}
                ),
                timeout=60.0
            )
            info("📝 [记忆添加] 已提交给 Mem0 进行事实提取")
            return True

        except asyncio.TimeoutError:
            error("❌ Mem0 记忆添加超时 (60s)")
            return False
        except Exception as e:
            error(f"❌ Mem0 记忆添加失败: {e}")
            return False