ailine/backend/app/memory/mem0_client.py

from ..config import LLM_API_KEY
from ..config import VLLM_BASE_URL
import time
"""
Mem0 记忆层客户端封装模块
负责 Mem0 的初始化、检索和存储
"""

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

from ..config import (
    QDRANT_URL,QDRANT_COLLECTION_NAME,QDRANT_API_KEY,
    VLLM_BASE_URL, LLM_API_KEY,
    LLAMACPP_EMBEDDING_URL, LLAMACPP_API_KEY
)
from ..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
                        "api_key": QDRANT_API_KEY,
                        "collection_name": QDRANT_COLLECTION_NAME,
                        "embedding_model_dims": 1024,
                    }
                },
                "llm": {
                    "provider": "openai",
                    "config": {
                        "model": "LLM_MODEL",
                        "api_key": LLM_API_KEY,
                        "openai_base_url": VLLM_BASE_URL,  
                        "temperature": 0.1,
                        "max_tokens": 2000,
                    }
                },
                "embedder": {
                    "provider": "openai",
                    "config": {
                        "model": "Qwen3-Embedding-0.6B-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, user_id):
      if not self.mem0:
        return False
      try:
        start = time.time()
        info(f"📝 开始 Mem0 add，消息数: {len(messages)}")
        await asyncio.wait_for(
            self.mem0.add(messages, user_id=user_id, metadata={"type": "conversation"}),
            timeout=60.0
        )
        info(f"✅ Mem0 add 完成，耗时: {time.time() - start:.2f}s")
        return True
      except asyncio.TimeoutError:
        error(f"❌ Mem0 记忆添加超时 (60s)，已等待 {time.time() - start:.2f}s")
        return False