File size: 5,153 Bytes
8a2dcce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """
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
|