memory/mem0_client.py

"""
mem0 client wrapper.

Responsibility boundary (do not blur this):
    - mem0       -> user preferences + long-term conversational memory
                    (e.g. "prefers beginner explanations", "likes step-by-step examples")
    - ChromaDB   -> DSA knowledge embeddings (rag/retriever.py) — NEVER touched here
    - SQLite     -> raw message/session transcripts (history/) — NEVER touched here

mem0 runs in local/open-source mode: a local Chroma vector store (separate
collection from the knowledge base) + a local sentence-transformers
embedder (same model as the KB for consistency, configurable independently
if needed).

Fact extraction (infer=True) would use an LLM to summarize what to
remember. To avoid depending on the (expensive, 4-bit-quantized) local
Mistral model just to log a preference, this wrapper uses infer=False by
default — messages are stored as-is and semantic search over them still
works for retrieval. Set infer=True if you want mem0 to run LLM-based
fact extraction using the same provider configured in llm/generate.py.
"""

import os

import config
from logs.logger import get_logger

logger = get_logger(__name__)

_MEM0_COLLECTION_NAME = "user_preferences_memory"


def _build_config_dict() -> dict:
    os.makedirs(config.MEM0_LOCAL_STORAGE_DIR, exist_ok=True)
    return {
        "vector_store": {
            "provider": "chroma",
            "config": {
                "collection_name": _MEM0_COLLECTION_NAME,
                "path": os.path.join(config.MEM0_LOCAL_STORAGE_DIR, "chroma_db"),
            },
        },
        "embedder": {
            "provider": "huggingface",
            "config": {
                "model": f"sentence-transformers/{config.EMBEDDING_MODEL_NAME}",
            },
        },
        "history_db_path": os.path.join(config.MEM0_LOCAL_STORAGE_DIR, "mem0_history.db"),
    }


class Mem0Client:
    """
    Thin wrapper around mem0.Memory scoped to a single responsibility:
    user preferences and long-term conversational memory.
    """

    def __init__(self, infer: bool = False):
        self.infer = infer
        self._memory = None
        self._enabled = config.USE_MEM0

        if not self._enabled:
            logger.info("USE_MEM0 is False — Mem0Client will be a no-op.")

    def _get_memory(self):
        if self._memory is None:
            from mem0 import Memory

            logger.info("Initializing mem0 (local chroma + huggingface embedder)...")
            self._memory = Memory.from_config(_build_config_dict())
        return self._memory

    def add_interaction(self, user_id: str, user_message: str, assistant_message: str = None) -> None:
        """
        Store a turn of conversation for long-term memory purposes (NOT
        transcript storage — that's history/service.py + SQLite).

        Only call this for things worth remembering long-term (preferences,
        recurring topics of interest, stated skill level) — not every raw
        message. The router/chat layer decides when this is worth calling;
        this wrapper doesn't filter content itself.
        """
        if not self._enabled:
            return

        messages = [{"role": "user", "content": user_message}]
        if assistant_message:
            messages.append({"role": "assistant", "content": assistant_message})

        try:
            self._get_memory().add(messages, user_id=user_id, infer=self.infer)
        except Exception:
            logger.exception("mem0 add_interaction failed for user_id=%s", user_id)

    def get_relevant_context(self, user_id: str, query: str, limit: int = 5) -> list:
        """
        Returns a list of plain-text memory strings relevant to `query`,
        ready to hand to llm/prompts.py's mem0_context parameter.

        Returns [] if mem0 is disabled, the user has no memories yet, or a
        lookup error occurs — callers should treat that as "no memory
        context available", not as an error condition.
        """
        if not self._enabled:
            return []

        try:
            results = self._get_memory().search(query, user_id=user_id, limit=limit)
        except Exception:
            logger.exception("mem0 get_relevant_context failed for user_id=%s", user_id)
            return []

        # mem0's search() returns {"results": [{"memory": "...", "score": ...}, ...]}
        # in v2, or a plain list in some versions — handle both.
        items = results.get("results", results) if isinstance(results, dict) else results
        return [item.get("memory", "") for item in items if item.get("memory")]

    def delete_all_for_user(self, user_id: str) -> None:
        """Useful for account deletion / a 'forget me' feature."""
        if not self._enabled:
            return
        try:
            self._get_memory().delete_all(user_id=user_id)
        except Exception:
            logger.exception("mem0 delete_all_for_user failed for user_id=%s", user_id)


_CLIENT_SINGLETON = None


def get_mem0_client() -> Mem0Client:
    global _CLIENT_SINGLETON
    if _CLIENT_SINGLETON is None:
        _CLIENT_SINGLETON = Mem0Client()
    return _CLIENT_SINGLETON