MemoryAgents/backend/app/adapters/mem0.py

from __future__ import annotations

from typing import Any

from ..config import settings
from ..llm import LLMUnavailable
from ..schemas import ChatResponse, SystemDescriptor
from .base import AdapterUnavailable, MemoryAgent


class Mem0Agent(MemoryAgent):
    id = "mem0"

    def __init__(self) -> None:
        super().__init__()
        self._memory: Any | None = None
        self._import_error: str | None = None
        try:
            from mem0 import Memory  # type: ignore

            self._memory_class = Memory
        except Exception as exc:
            self._memory_class = None
            self._import_error = str(exc)
        embedding_ready = bool(
            settings.embedding_api_key
            and settings.embedding_base_url
            and settings.embedding_model
        )
        ready = bool(self._memory_class and settings.mem0_enabled and embedding_ready)
        self.descriptor = SystemDescriptor(
            id="mem0",
            name="Mem0",
            paradigm="自动提取与冲突消解",
            description="将对话提炼为长期记忆，并按逻辑用户标识执行检索和更新。",
            available=ready,
            mode="real-sdk" if ready else "unavailable",
            status="ready" if ready else "not-configured",
            package="mem0ai",
            setup_hint=None if ready else (
                "当前仅保留 Mem0 适配骨架，请保持 MEM0_ENABLED=false。"
                "完成公共记忆投影、删除、重置、冲突和隔离闭环后再启用。"
            ),
        )

    async def _client(self) -> Any:
        if not self._memory_class or not settings.mem0_enabled:
            raise AdapterUnavailable(self.descriptor.setup_hint or "Mem0 不可用")
        if self._memory is None:
            if not settings.embedding_api_key or not settings.embedding_model:
                raise AdapterUnavailable(self.descriptor.setup_hint or "Mem0 缺少 Embedding 配置")
            llm_provider = "deepseek" if "deepseek" in settings.llm_base_url.lower() else "openai"
            llm_config = {"model": settings.llm_model, "api_key": settings.llm_api_key}
            llm_config["deepseek_base_url" if llm_provider == "deepseek" else "openai_base_url"] = settings.llm_base_url
            self._memory = self._memory_class.from_config({
                "llm": {"provider": llm_provider, "config": llm_config},
                "embedder": {
                    "provider": "openai",
                    "config": {
                        "model": settings.embedding_model,
                        "api_key": settings.embedding_api_key,
                        "openai_base_url": settings.embedding_base_url,
                        "embedding_dims": settings.embedding_dimensions,
                    },
                },
                "vector_store": {
                    "provider": "pgvector",
                    "config": {
                        "connection_string": settings.database_url,
                        "collection_name": "mem0_memory",
                        "embedding_model_dims": settings.embedding_dimensions,
                    },
                },
            })
        return self._memory

    async def chat(self, message: str) -> ChatResponse:
        memory = await self._client()
        add_result = memory.add(
            [{"role": "user", "content": message}],
            user_id=settings.workspace_id,
        )
        if hasattr(add_result, "__await__"):
            add_result = await add_result
        search_result = memory.search(message, user_id=settings.workspace_id, limit=5)
        if hasattr(search_result, "__await__"):
            search_result = await search_result
        context = search_result.get("results", search_result if isinstance(search_result, list) else [])
        context_text = "\n".join(f"- {item.get('memory', item)}" for item in context)
        try:
            answer = await self.llm.chat(
                "你是使用 Mem0 记忆的 AI 编程助手。只把 Mem0 返回的内容当作记忆证据。",
                f"Mem0 召回：\n{context_text or '（暂无）'}\n\n用户：{message}",
            )
            mode = "real-sdk+llm"
        except LLMUnavailable:
            answer = f"Mem0 已完成写入和检索。\n\n召回内容：\n{context_text or '（暂无）'}"
            mode = "real-sdk"
        await self._audit("MEM0/add+search", details={"raw_add": str(add_result)[:1000]})
        return ChatResponse(
            system="mem0", answer=answer, mode=mode,
            memory_context=await self.memories(), memory_events=[{"event": "add", "result": add_result}],
            audit_events=await self.audit(),
        )