Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6d64bbc4cf |
+92
-39
@@ -9,7 +9,7 @@ import uuid
|
|||||||
import warnings
|
import warnings
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, Optional
|
||||||
|
|
||||||
from pydantic import ValidationError
|
from pydantic import ValidationError
|
||||||
|
|
||||||
@@ -479,8 +479,7 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
msg_content = message_dict["content"]
|
msg_content = message_dict["content"]
|
||||||
msg_embeddings = self.embedding_model.embed(msg_content, "add")
|
msg_embeddings = self.embedding_model.embed(msg_content, "add")
|
||||||
# Pass embeddings as a dict so _create_memory can reuse the cached embedding
|
mem_id = self._create_memory(msg_content, msg_embeddings, per_msg_meta)
|
||||||
mem_id = self._create_memory(msg_content, {msg_content: msg_embeddings}, per_msg_meta)
|
|
||||||
|
|
||||||
returned_memories.append(
|
returned_memories.append(
|
||||||
{
|
{
|
||||||
@@ -609,9 +608,6 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
event_type = resp.get("event")
|
event_type = resp.get("event")
|
||||||
if event_type == "ADD":
|
if event_type == "ADD":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = self.embedding_model.embed(action_text, "add")
|
|
||||||
memory_id = self._create_memory(
|
memory_id = self._create_memory(
|
||||||
data=action_text,
|
data=action_text,
|
||||||
existing_embeddings=new_message_embeddings,
|
existing_embeddings=new_message_embeddings,
|
||||||
@@ -619,9 +615,6 @@ class Memory(MemoryBase):
|
|||||||
)
|
)
|
||||||
returned_memories.append({"id": memory_id, "memory": action_text, "event": event_type})
|
returned_memories.append({"id": memory_id, "memory": action_text, "event": event_type})
|
||||||
elif event_type == "UPDATE":
|
elif event_type == "UPDATE":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = self.embedding_model.embed(action_text, "update")
|
|
||||||
self._update_memory(
|
self._update_memory(
|
||||||
memory_id=temp_uuid_mapping[resp.get("id")],
|
memory_id=temp_uuid_mapping[resp.get("id")],
|
||||||
data=action_text,
|
data=action_text,
|
||||||
@@ -1160,13 +1153,10 @@ class Memory(MemoryBase):
|
|||||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "sync"})
|
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "sync"})
|
||||||
return self.db.get_history(memory_id)
|
return self.db.get_history(memory_id)
|
||||||
|
|
||||||
def _create_memory(self, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||||
logger.debug(f"Creating memory with {data=}")
|
logger.debug(f"Creating memory with {data=}")
|
||||||
# existing_embeddings may be a dict (preferred) or a precomputed vector
|
if data in existing_embeddings:
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = self.embedding_model.embed(data, memory_action="add")
|
embeddings = self.embedding_model.embed(data, memory_action="add")
|
||||||
memory_id = str(uuid.uuid4())
|
memory_id = str(uuid.uuid4())
|
||||||
@@ -1230,7 +1220,7 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
def _update_memory(self, memory_id, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
|
||||||
logger.info(f"Updating memory with {data=}")
|
logger.info(f"Updating memory with {data=}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -1263,10 +1253,8 @@ class Memory(MemoryBase):
|
|||||||
if "role" not in new_metadata and "role" in existing_memory.payload:
|
if "role" not in new_metadata and "role" in existing_memory.payload:
|
||||||
new_metadata["role"] = existing_memory.payload["role"]
|
new_metadata["role"] = existing_memory.payload["role"]
|
||||||
|
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
if data in existing_embeddings:
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = self.embedding_model.embed(data, "update")
|
embeddings = self.embedding_model.embed(data, "update")
|
||||||
|
|
||||||
@@ -1332,6 +1320,47 @@ class Memory(MemoryBase):
|
|||||||
)
|
)
|
||||||
capture_event("mem0.reset", self, {"sync_type": "sync"})
|
capture_event("mem0.reset", self, {"sync_type": "sync"})
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
"""Release all resources held by this Memory instance.
|
||||||
|
|
||||||
|
Shuts down telemetry background threads, closes the SQLite
|
||||||
|
history-store connection, and closes the vector store client
|
||||||
|
if it exposes a close() method.
|
||||||
|
|
||||||
|
Safe to call multiple times. Use as a context manager for
|
||||||
|
automatic cleanup::
|
||||||
|
|
||||||
|
with Memory.from_config(config) as m:
|
||||||
|
m.add(...)
|
||||||
|
"""
|
||||||
|
from mem0.memory.telemetry import shutdown_telemetry
|
||||||
|
|
||||||
|
try:
|
||||||
|
shutdown_telemetry()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error shutting down telemetry", exc_info=True)
|
||||||
|
|
||||||
|
if hasattr(self, "db") and self.db is not None:
|
||||||
|
try:
|
||||||
|
self.db.close()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error closing database", exc_info=True)
|
||||||
|
|
||||||
|
if hasattr(self, "vector_store") and self.vector_store is not None:
|
||||||
|
_close = getattr(self.vector_store, "close", None)
|
||||||
|
if callable(_close):
|
||||||
|
try:
|
||||||
|
_close()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error closing vector store", exc_info=True)
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||||
|
self.close()
|
||||||
|
return False
|
||||||
|
|
||||||
def chat(self, query):
|
def chat(self, query):
|
||||||
raise NotImplementedError("Chat function not implemented yet.")
|
raise NotImplementedError("Chat function not implemented yet.")
|
||||||
|
|
||||||
@@ -1535,8 +1564,7 @@ class AsyncMemory(MemoryBase):
|
|||||||
|
|
||||||
msg_content = message_dict["content"]
|
msg_content = message_dict["content"]
|
||||||
msg_embeddings = await asyncio.to_thread(self.embedding_model.embed, msg_content, "add")
|
msg_embeddings = await asyncio.to_thread(self.embedding_model.embed, msg_content, "add")
|
||||||
# Pass embeddings as a dict so _create_memory can reuse the cached embedding
|
mem_id = await self._create_memory(msg_content, msg_embeddings, per_msg_meta)
|
||||||
mem_id = await self._create_memory(msg_content, {msg_content: msg_embeddings}, per_msg_meta)
|
|
||||||
|
|
||||||
returned_memories.append(
|
returned_memories.append(
|
||||||
{
|
{
|
||||||
@@ -1664,11 +1692,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
event_type = resp.get("event")
|
event_type = resp.get("event")
|
||||||
|
|
||||||
if event_type == "ADD":
|
if event_type == "ADD":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = await asyncio.to_thread(
|
|
||||||
self.embedding_model.embed, action_text, "add"
|
|
||||||
)
|
|
||||||
task = asyncio.create_task(
|
task = asyncio.create_task(
|
||||||
self._create_memory(
|
self._create_memory(
|
||||||
data=action_text,
|
data=action_text,
|
||||||
@@ -1678,11 +1701,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
)
|
)
|
||||||
memory_tasks.append((task, resp, "ADD", None))
|
memory_tasks.append((task, resp, "ADD", None))
|
||||||
elif event_type == "UPDATE":
|
elif event_type == "UPDATE":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = await asyncio.to_thread(
|
|
||||||
self.embedding_model.embed, action_text, "update"
|
|
||||||
)
|
|
||||||
task = asyncio.create_task(
|
task = asyncio.create_task(
|
||||||
self._update_memory(
|
self._update_memory(
|
||||||
memory_id=temp_uuid_mapping[resp["id"]],
|
memory_id=temp_uuid_mapping[resp["id"]],
|
||||||
@@ -2252,13 +2270,10 @@ class AsyncMemory(MemoryBase):
|
|||||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "async"})
|
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "async"})
|
||||||
return await asyncio.to_thread(self.db.get_history, memory_id)
|
return await asyncio.to_thread(self.db.get_history, memory_id)
|
||||||
|
|
||||||
async def _create_memory(self, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
async def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||||
logger.debug(f"Creating memory with {data=}")
|
logger.debug(f"Creating memory with {data=}")
|
||||||
# existing_embeddings may be a dict (preferred) or a precomputed vector
|
if data in existing_embeddings:
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, memory_action="add")
|
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, memory_action="add")
|
||||||
|
|
||||||
@@ -2341,7 +2356,7 @@ class AsyncMemory(MemoryBase):
|
|||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
async def _update_memory(self, memory_id, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
async def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
|
||||||
logger.info(f"Updating memory with {data=}")
|
logger.info(f"Updating memory with {data=}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -2375,10 +2390,8 @@ class AsyncMemory(MemoryBase):
|
|||||||
if "role" not in new_metadata and "role" in existing_memory.payload:
|
if "role" not in new_metadata and "role" in existing_memory.payload:
|
||||||
new_metadata["role"] = existing_memory.payload["role"]
|
new_metadata["role"] = existing_memory.payload["role"]
|
||||||
|
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
if data in existing_embeddings:
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, "update")
|
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, "update")
|
||||||
|
|
||||||
@@ -2450,5 +2463,45 @@ class AsyncMemory(MemoryBase):
|
|||||||
)
|
)
|
||||||
capture_event("mem0.reset", self, {"sync_type": "async"})
|
capture_event("mem0.reset", self, {"sync_type": "async"})
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
"""Release all resources held by this AsyncMemory instance.
|
||||||
|
|
||||||
|
Shuts down telemetry background threads, closes the SQLite
|
||||||
|
history-store connection, and closes the vector store client
|
||||||
|
if it exposes a close() method.
|
||||||
|
|
||||||
|
Safe to call multiple times. Use as an async context manager::
|
||||||
|
|
||||||
|
async with await AsyncMemory.from_config(config) as m:
|
||||||
|
await m.add(...)
|
||||||
|
"""
|
||||||
|
from mem0.memory.telemetry import shutdown_telemetry
|
||||||
|
|
||||||
|
try:
|
||||||
|
shutdown_telemetry()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error shutting down telemetry", exc_info=True)
|
||||||
|
|
||||||
|
if hasattr(self, "db") and self.db is not None:
|
||||||
|
try:
|
||||||
|
self.db.close()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error closing database", exc_info=True)
|
||||||
|
|
||||||
|
if hasattr(self, "vector_store") and self.vector_store is not None:
|
||||||
|
_close = getattr(self.vector_store, "close", None)
|
||||||
|
if callable(_close):
|
||||||
|
try:
|
||||||
|
_close()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error closing vector store", exc_info=True)
|
||||||
|
|
||||||
|
async def __aenter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||||
|
self.close()
|
||||||
|
return False
|
||||||
|
|
||||||
async def chat(self, query):
|
async def chat(self, query):
|
||||||
raise NotImplementedError("Chat function not implemented yet.")
|
raise NotImplementedError("Chat function not implemented yet.")
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
|
import atexit
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import platform
|
import platform
|
||||||
import sys
|
import sys
|
||||||
|
import threading
|
||||||
|
|
||||||
from posthog import Posthog
|
from posthog import Posthog
|
||||||
|
|
||||||
@@ -57,18 +59,62 @@ class AnonymousTelemetry:
|
|||||||
self.posthog.shutdown()
|
self.posthog.shutdown()
|
||||||
|
|
||||||
|
|
||||||
client_telemetry = AnonymousTelemetry()
|
# Thread-safe lazy singletons
|
||||||
|
_telemetry_lock = threading.Lock()
|
||||||
|
_oss_telemetry = None
|
||||||
|
_client_telemetry = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_oss_telemetry(vector_store=None):
|
||||||
|
"""Return the singleton OSS telemetry instance, creating it on first call.
|
||||||
|
|
||||||
|
The vector_store parameter is only used on the first call for
|
||||||
|
get_or_create_user_id(). Subsequent calls reuse the existing instance
|
||||||
|
since the user ID is stable once created.
|
||||||
|
"""
|
||||||
|
global _oss_telemetry
|
||||||
|
if _oss_telemetry is None:
|
||||||
|
with _telemetry_lock:
|
||||||
|
if _oss_telemetry is None:
|
||||||
|
_oss_telemetry = AnonymousTelemetry(vector_store=vector_store)
|
||||||
|
return _oss_telemetry
|
||||||
|
|
||||||
|
|
||||||
|
def _get_client_telemetry():
|
||||||
|
"""Return the singleton client telemetry instance, creating it on first call."""
|
||||||
|
global _client_telemetry
|
||||||
|
if _client_telemetry is None:
|
||||||
|
with _telemetry_lock:
|
||||||
|
if _client_telemetry is None:
|
||||||
|
_client_telemetry = AnonymousTelemetry()
|
||||||
|
return _client_telemetry
|
||||||
|
|
||||||
|
|
||||||
|
def shutdown_telemetry():
|
||||||
|
"""Shut down all telemetry singletons. Safe to call multiple times.
|
||||||
|
|
||||||
|
Singletons are lazily re-created on next use, so this is safe even
|
||||||
|
when multiple Memory instances exist.
|
||||||
|
"""
|
||||||
|
global _oss_telemetry, _client_telemetry
|
||||||
|
with _telemetry_lock:
|
||||||
|
if _oss_telemetry is not None:
|
||||||
|
_oss_telemetry.close()
|
||||||
|
_oss_telemetry = None
|
||||||
|
if _client_telemetry is not None:
|
||||||
|
_client_telemetry.close()
|
||||||
|
_client_telemetry = None
|
||||||
|
|
||||||
|
|
||||||
|
atexit.register(shutdown_telemetry)
|
||||||
|
|
||||||
|
|
||||||
def capture_event(event_name, memory_instance, additional_data=None):
|
def capture_event(event_name, memory_instance, additional_data=None):
|
||||||
if not MEM0_TELEMETRY:
|
if not MEM0_TELEMETRY:
|
||||||
return
|
return
|
||||||
|
|
||||||
oss_telemetry = AnonymousTelemetry(
|
vector_store = getattr(memory_instance, "_telemetry_vector_store", None)
|
||||||
vector_store=memory_instance._telemetry_vector_store
|
oss_telemetry = _get_oss_telemetry(vector_store=vector_store)
|
||||||
if hasattr(memory_instance, "_telemetry_vector_store")
|
|
||||||
else None,
|
|
||||||
)
|
|
||||||
|
|
||||||
event_data = {
|
event_data = {
|
||||||
"collection": memory_instance.collection_name,
|
"collection": memory_instance.collection_name,
|
||||||
@@ -92,6 +138,8 @@ def capture_client_event(event_name, instance, additional_data=None):
|
|||||||
if not MEM0_TELEMETRY:
|
if not MEM0_TELEMETRY:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
client_telemetry = _get_client_telemetry()
|
||||||
|
|
||||||
event_data = {
|
event_data = {
|
||||||
"function": f"{instance.__class__.__module__}.{instance.__class__.__name__}",
|
"function": f"{instance.__class__.__module__}.{instance.__class__.__name__}",
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ class OutputData(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class MongoDB(VectorStoreBase):
|
class MongoDB(VectorStoreBase):
|
||||||
VECTOR_TYPE = "vector"
|
VECTOR_TYPE = "knnVector"
|
||||||
SIMILARITY_METRIC = "cosine"
|
SIMILARITY_METRIC = "cosine"
|
||||||
|
|
||||||
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
|
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
|
||||||
@@ -69,16 +69,17 @@ class MongoDB(VectorStoreBase):
|
|||||||
else:
|
else:
|
||||||
search_index_model = SearchIndexModel(
|
search_index_model = SearchIndexModel(
|
||||||
name=self.index_name,
|
name=self.index_name,
|
||||||
type="vectorSearch",
|
|
||||||
definition={
|
definition={
|
||||||
"fields": [
|
"mappings": {
|
||||||
{
|
"dynamic": False,
|
||||||
"type": self.VECTOR_TYPE,
|
"fields": {
|
||||||
"path": "embedding",
|
"embedding": {
|
||||||
"numDimensions": self.embedding_model_dims,
|
"type": self.VECTOR_TYPE,
|
||||||
"similarity": self.SIMILARITY_METRIC,
|
"dimensions": self.embedding_model_dims,
|
||||||
}
|
"similarity": self.SIMILARITY_METRIC,
|
||||||
]
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
collection.create_search_index(search_index_model)
|
collection.create_search_index(search_index_model)
|
||||||
@@ -140,7 +141,7 @@ class MongoDB(VectorStoreBase):
|
|||||||
"$vectorSearch": {
|
"$vectorSearch": {
|
||||||
"index": self.index_name,
|
"index": self.index_name,
|
||||||
"limit": limit,
|
"limit": limit,
|
||||||
"numCandidates": min(limit * 20, 10000),
|
"numCandidates": limit,
|
||||||
"queryVector": vectors,
|
"queryVector": vectors,
|
||||||
"path": "embedding",
|
"path": "embedding",
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -438,137 +438,3 @@ async def test_async_update_nonexistent_memory_raises_error(mock_sqlite, mock_ll
|
|||||||
await memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
|
await memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
|
||||||
|
|
||||||
mock_vector_store.update.assert_not_called()
|
mock_vector_store.update.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
|
||||||
@patch('mem0.memory.storage.SQLiteManager')
|
|
||||||
def test_add_infer_false_embeds_once(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723: adding with infer=False should not trigger duplicate embedding calls.
|
|
||||||
|
|
||||||
Root cause: _create_memory expected a dict for existing_embeddings but received a raw list[float],
|
|
||||||
causing the cache check `data in existing_embeddings` to always fail and trigger a redundant embed.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = []
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.get.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
mock_llm_factory.return_value = MagicMock()
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("foo", user_id="test_user", infer=False)
|
|
||||||
|
|
||||||
assert embedder.embed.call_count == 1
|
|
||||||
mock_vector_store.insert.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
|
||||||
@patch('mem0.memory.storage.SQLiteManager')
|
|
||||||
def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723 (infer=True path): when the LLM rewrites a fact during the
|
|
||||||
ADD action, the embedding should be computed once and cached, not computed again inside _create_memory.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = []
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.get.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
# LLM extracts fact "User likes Python", then ADD action rewrites to "The user enjoys Python"
|
|
||||||
mock_llm = MagicMock()
|
|
||||||
mock_llm.generate_response.side_effect = [
|
|
||||||
json.dumps({"facts": ["User likes Python"]}),
|
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user enjoys Python", "event": "ADD", "old_memory": None}]}),
|
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("I like Python", user_id="test_user", infer=True)
|
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
|
||||||
# 1. For the extracted fact "User likes Python" (search)
|
|
||||||
# 2. For the rewritten text "The user enjoys Python" (pre-cached before _create_memory)
|
|
||||||
# It should NOT be called a 3rd time inside _create_memory
|
|
||||||
assert embedder.embed.call_count == 2
|
|
||||||
mock_vector_store.insert.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
|
||||||
@patch('mem0.memory.storage.SQLiteManager')
|
|
||||||
def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723 (infer=True UPDATE path): when the LLM rewrites a fact during
|
|
||||||
an UPDATE action, the embedding should be computed once and cached, not computed again inside _update_memory.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
# Existing memory that will be matched for update
|
|
||||||
existing_memory = MockVectorMemory(
|
|
||||||
memory_id="existing-mem-id",
|
|
||||||
payload={
|
|
||||||
"data": "User likes Python",
|
|
||||||
"hash": "abc123",
|
|
||||||
"created_at": "2025-01-01T00:00:00+00:00",
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = [existing_memory]
|
|
||||||
mock_vector_store.get.return_value = existing_memory
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.update.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
# LLM extracts fact "User loves Python now", then UPDATE action rewrites to "The user loves Python"
|
|
||||||
mock_llm = MagicMock()
|
|
||||||
mock_llm.generate_response.side_effect = [
|
|
||||||
json.dumps({"facts": ["User loves Python now"]}),
|
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user loves Python", "event": "UPDATE", "old_memory": "User likes Python"}]}),
|
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("I love Python now", user_id="test_user", infer=True)
|
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
|
||||||
# 1. For the extracted fact "User loves Python now" (search)
|
|
||||||
# 2. For the rewritten text "The user loves Python" (pre-cached before _update_memory)
|
|
||||||
# It should NOT be called a 3rd time inside _update_memory
|
|
||||||
assert embedder.embed.call_count == 2
|
|
||||||
mock_vector_store.update.assert_called_once()
|
|
||||||
|
|||||||
@@ -0,0 +1,164 @@
|
|||||||
|
"""Tests for Memory resource cleanup: close() and context manager."""
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def memory_instance():
|
||||||
|
"""Create a Memory instance with all dependencies mocked."""
|
||||||
|
with (
|
||||||
|
patch("mem0.memory.main.EmbedderFactory") as mock_embedder,
|
||||||
|
patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
|
||||||
|
patch("mem0.memory.main.LlmFactory") as mock_llm,
|
||||||
|
patch("mem0.memory.telemetry.capture_event"),
|
||||||
|
patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
|
||||||
|
):
|
||||||
|
mock_embedder.create.return_value = MagicMock()
|
||||||
|
mock_embedder.create.return_value.config.embedding_dims = 1536
|
||||||
|
mock_vector_store.create.return_value = MagicMock()
|
||||||
|
mock_llm.create.return_value = MagicMock()
|
||||||
|
mock_graph_store.create.return_value = MagicMock()
|
||||||
|
|
||||||
|
from mem0.memory.main import Memory
|
||||||
|
m = Memory()
|
||||||
|
yield m
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemoryClose:
|
||||||
|
"""Verify Memory.close() releases resources."""
|
||||||
|
|
||||||
|
def test_close_shuts_down_telemetry(self, memory_instance):
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry") as mock_shutdown:
|
||||||
|
memory_instance.close()
|
||||||
|
mock_shutdown.assert_called_once()
|
||||||
|
|
||||||
|
def test_close_closes_db(self, memory_instance):
|
||||||
|
memory_instance.db = MagicMock()
|
||||||
|
memory_instance.close()
|
||||||
|
memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
def test_close_closes_vector_store_if_available(self, memory_instance):
|
||||||
|
mock_vs = MagicMock(spec=["close"])
|
||||||
|
memory_instance.vector_store = mock_vs
|
||||||
|
memory_instance.close()
|
||||||
|
mock_vs.close.assert_called_once()
|
||||||
|
|
||||||
|
def test_close_skips_vector_store_without_close(self, memory_instance):
|
||||||
|
mock_vs = MagicMock(spec=[])
|
||||||
|
memory_instance.vector_store = mock_vs
|
||||||
|
memory_instance.close() # should not raise
|
||||||
|
|
||||||
|
def test_close_is_idempotent(self, memory_instance):
|
||||||
|
memory_instance.close()
|
||||||
|
memory_instance.close() # should not raise
|
||||||
|
|
||||||
|
def test_close_continues_on_db_error(self, memory_instance):
|
||||||
|
memory_instance.db = MagicMock()
|
||||||
|
memory_instance.db.close.side_effect = Exception("db error")
|
||||||
|
memory_instance.close() # should not raise
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemoryContextManager:
|
||||||
|
"""Verify Memory works as a context manager."""
|
||||||
|
|
||||||
|
def test_enter_returns_self(self, memory_instance):
|
||||||
|
assert memory_instance.__enter__() is memory_instance
|
||||||
|
|
||||||
|
def test_exit_calls_close(self, memory_instance):
|
||||||
|
memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
memory_instance.__exit__(None, None, None)
|
||||||
|
memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
def test_does_not_suppress_exceptions(self, memory_instance):
|
||||||
|
assert memory_instance.__exit__(ValueError, ValueError("x"), None) is False
|
||||||
|
|
||||||
|
def test_with_statement(self, memory_instance):
|
||||||
|
memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
with memory_instance:
|
||||||
|
pass
|
||||||
|
memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
def test_with_statement_on_exception(self, memory_instance):
|
||||||
|
memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
with memory_instance:
|
||||||
|
raise RuntimeError("boom")
|
||||||
|
memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def async_memory_instance():
|
||||||
|
"""Create an AsyncMemory instance with all dependencies mocked."""
|
||||||
|
with (
|
||||||
|
patch("mem0.memory.main.EmbedderFactory") as mock_embedder,
|
||||||
|
patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
|
||||||
|
patch("mem0.memory.main.LlmFactory") as mock_llm,
|
||||||
|
patch("mem0.memory.telemetry.capture_event"),
|
||||||
|
patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
|
||||||
|
):
|
||||||
|
mock_embedder.create.return_value = MagicMock()
|
||||||
|
mock_embedder.create.return_value.config.embedding_dims = 1536
|
||||||
|
mock_vector_store.create.return_value = MagicMock()
|
||||||
|
mock_llm.create.return_value = MagicMock()
|
||||||
|
mock_graph_store.create.return_value = MagicMock()
|
||||||
|
|
||||||
|
from mem0.memory.main import AsyncMemory
|
||||||
|
m = AsyncMemory()
|
||||||
|
yield m
|
||||||
|
|
||||||
|
|
||||||
|
class TestAsyncMemoryClose:
|
||||||
|
"""Verify AsyncMemory.close() releases resources."""
|
||||||
|
|
||||||
|
def test_close_shuts_down_telemetry(self, async_memory_instance):
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry") as mock_shutdown:
|
||||||
|
async_memory_instance.close()
|
||||||
|
mock_shutdown.assert_called_once()
|
||||||
|
|
||||||
|
def test_close_closes_db(self, async_memory_instance):
|
||||||
|
async_memory_instance.db = MagicMock()
|
||||||
|
async_memory_instance.close()
|
||||||
|
async_memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
def test_close_is_idempotent(self, async_memory_instance):
|
||||||
|
async_memory_instance.close()
|
||||||
|
async_memory_instance.close() # should not raise
|
||||||
|
|
||||||
|
|
||||||
|
class TestAsyncMemoryContextManager:
|
||||||
|
"""Verify AsyncMemory works as an async context manager."""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_aenter_returns_self(self, async_memory_instance):
|
||||||
|
result = await async_memory_instance.__aenter__()
|
||||||
|
assert result is async_memory_instance
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_aexit_calls_close(self, async_memory_instance):
|
||||||
|
async_memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
await async_memory_instance.__aexit__(None, None, None)
|
||||||
|
async_memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_with_statement(self, async_memory_instance):
|
||||||
|
async_memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
async with async_memory_instance:
|
||||||
|
pass
|
||||||
|
async_memory_instance.db.close.assert_called_once()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_with_statement_on_exception(self, async_memory_instance):
|
||||||
|
async_memory_instance.db = MagicMock()
|
||||||
|
with patch("mem0.memory.telemetry.shutdown_telemetry"):
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
async with async_memory_instance:
|
||||||
|
raise RuntimeError("boom")
|
||||||
|
async_memory_instance.db.close.assert_called_once()
|
||||||
+81
-17
@@ -30,9 +30,10 @@ class TestTelemetryDisabled:
|
|||||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", False):
|
with patch.object(telemetry_module, "MEM0_TELEMETRY", False):
|
||||||
mock_instance = MagicMock()
|
mock_instance = MagicMock()
|
||||||
mock_client_telemetry = MagicMock()
|
mock_client_telemetry = MagicMock()
|
||||||
with patch.object(telemetry_module, "client_telemetry", mock_client_telemetry):
|
telemetry_module._client_telemetry = mock_client_telemetry
|
||||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||||
mock_client_telemetry.capture_event.assert_not_called()
|
mock_client_telemetry.capture_event.assert_not_called()
|
||||||
|
telemetry_module._client_telemetry = None
|
||||||
|
|
||||||
def test_instance_capture_event_noop_when_posthog_is_none(self):
|
def test_instance_capture_event_noop_when_posthog_is_none(self):
|
||||||
"""AnonymousTelemetry.capture_event() should be a no-op when posthog is None."""
|
"""AnonymousTelemetry.capture_event() should be a no-op when posthog is None."""
|
||||||
@@ -69,26 +70,89 @@ class TestTelemetryEnabled:
|
|||||||
assert at.user_id == "test-user"
|
assert at.user_id == "test-user"
|
||||||
|
|
||||||
def test_capture_event_sends_when_enabled(self):
|
def test_capture_event_sends_when_enabled(self):
|
||||||
"""capture_event() should create AnonymousTelemetry and call capture when enabled."""
|
"""capture_event() should use singleton AnonymousTelemetry and call capture when enabled."""
|
||||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||||
with patch("mem0.memory.telemetry.AnonymousTelemetry") as mock_cls:
|
mock_at = MagicMock()
|
||||||
mock_at = MagicMock()
|
telemetry_module._oss_telemetry = mock_at
|
||||||
mock_cls.return_value = mock_at
|
mock_memory = MagicMock()
|
||||||
mock_memory = MagicMock()
|
mock_memory.config.graph_store.config = None
|
||||||
mock_memory.config.graph_store.config = None
|
mock_memory.api_version = "v1"
|
||||||
mock_memory.api_version = "v1"
|
telemetry_module.capture_event("test.event", mock_memory)
|
||||||
telemetry_module.capture_event("test.event", mock_memory)
|
mock_at.capture_event.assert_called_once()
|
||||||
mock_at.capture_event.assert_called_once()
|
telemetry_module._oss_telemetry = None
|
||||||
|
|
||||||
def test_capture_client_event_sends_when_enabled(self):
|
def test_capture_client_event_sends_when_enabled(self):
|
||||||
"""capture_client_event() should call client_telemetry.capture_event when enabled."""
|
"""capture_client_event() should call client_telemetry.capture_event when enabled."""
|
||||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||||
mock_client_telemetry = MagicMock()
|
mock_client_telemetry = MagicMock()
|
||||||
with patch.object(telemetry_module, "client_telemetry", mock_client_telemetry):
|
telemetry_module._client_telemetry = mock_client_telemetry
|
||||||
mock_instance = MagicMock()
|
mock_instance = MagicMock()
|
||||||
mock_instance.user_email = "test@example.com"
|
mock_instance.user_email = "test@example.com"
|
||||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||||
mock_client_telemetry.capture_event.assert_called_once()
|
mock_client_telemetry.capture_event.assert_called_once()
|
||||||
|
telemetry_module._client_telemetry = None
|
||||||
|
|
||||||
|
def test_capture_event_reuses_singleton(self):
|
||||||
|
"""capture_event() must reuse a single AnonymousTelemetry, not create one per call."""
|
||||||
|
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||||
|
with patch("mem0.memory.telemetry.AnonymousTelemetry") as mock_cls:
|
||||||
|
mock_at = MagicMock()
|
||||||
|
mock_cls.return_value = mock_at
|
||||||
|
telemetry_module._oss_telemetry = None
|
||||||
|
mock_memory = MagicMock()
|
||||||
|
mock_memory.config.graph_store.config = None
|
||||||
|
mock_memory.api_version = "v1"
|
||||||
|
|
||||||
|
telemetry_module.capture_event("event1", mock_memory)
|
||||||
|
telemetry_module.capture_event("event2", mock_memory)
|
||||||
|
|
||||||
|
assert mock_cls.call_count == 1
|
||||||
|
assert mock_at.capture_event.call_count == 2
|
||||||
|
|
||||||
|
telemetry_module._oss_telemetry = None
|
||||||
|
|
||||||
|
|
||||||
|
class TestTelemetryShutdown:
|
||||||
|
"""Verify shutdown_telemetry() cleans up all singleton instances."""
|
||||||
|
|
||||||
|
def test_shutdown_calls_close_on_singletons(self):
|
||||||
|
mock_oss = MagicMock()
|
||||||
|
mock_client = MagicMock()
|
||||||
|
telemetry_module._oss_telemetry = mock_oss
|
||||||
|
telemetry_module._client_telemetry = mock_client
|
||||||
|
|
||||||
|
telemetry_module.shutdown_telemetry()
|
||||||
|
|
||||||
|
mock_oss.close.assert_called_once()
|
||||||
|
mock_client.close.assert_called_once()
|
||||||
|
assert telemetry_module._oss_telemetry is None
|
||||||
|
assert telemetry_module._client_telemetry is None
|
||||||
|
|
||||||
|
def test_shutdown_is_safe_when_no_singletons(self):
|
||||||
|
telemetry_module._oss_telemetry = None
|
||||||
|
telemetry_module._client_telemetry = None
|
||||||
|
telemetry_module.shutdown_telemetry() # should not raise
|
||||||
|
|
||||||
|
|
||||||
|
class TestNoThreadLeak:
|
||||||
|
"""Verify capture_event does not leak threads."""
|
||||||
|
|
||||||
|
def test_multiple_capture_events_single_posthog_instance(self):
|
||||||
|
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||||
|
with patch("mem0.memory.telemetry.Posthog") as mock_posthog:
|
||||||
|
with patch("mem0.memory.telemetry.get_or_create_user_id", return_value="test"):
|
||||||
|
telemetry_module._oss_telemetry = None
|
||||||
|
|
||||||
|
mock_memory = MagicMock()
|
||||||
|
mock_memory.config.graph_store.config = None
|
||||||
|
mock_memory.api_version = "v1"
|
||||||
|
|
||||||
|
for _ in range(10):
|
||||||
|
telemetry_module.capture_event("test.event", mock_memory)
|
||||||
|
|
||||||
|
assert mock_posthog.call_count == 1
|
||||||
|
|
||||||
|
telemetry_module._oss_telemetry = None
|
||||||
|
|
||||||
|
|
||||||
class TestTelemetryEnvVar:
|
class TestTelemetryEnvVar:
|
||||||
|
|||||||
@@ -48,16 +48,17 @@ def test_initalize_create_col(mongo_vector_fixture):
|
|||||||
search_index_model = args[0].document
|
search_index_model = args[0].document
|
||||||
assert search_index_model == {
|
assert search_index_model == {
|
||||||
"name": "test_collection_vector_index",
|
"name": "test_collection_vector_index",
|
||||||
"type": "vectorSearch",
|
|
||||||
"definition": {
|
"definition": {
|
||||||
"fields": [
|
"mappings": {
|
||||||
{
|
"dynamic": False,
|
||||||
"type": "vector",
|
"fields": {
|
||||||
"path": "embedding",
|
"embedding": {
|
||||||
"numDimensions": 1536,
|
"type": "knnVector",
|
||||||
"similarity": "cosine",
|
"dimensions": 1536,
|
||||||
}
|
"similarity": "cosine",
|
||||||
]
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
assert mongo_vector.collection == mock_collection
|
assert mongo_vector.collection == mock_collection
|
||||||
@@ -94,7 +95,7 @@ def test_search(mongo_vector_fixture):
|
|||||||
"$vectorSearch": {
|
"$vectorSearch": {
|
||||||
"index": "test_collection_vector_index",
|
"index": "test_collection_vector_index",
|
||||||
"limit": 2,
|
"limit": 2,
|
||||||
"numCandidates": 40,
|
"numCandidates": 2,
|
||||||
"queryVector": query_vector,
|
"queryVector": query_vector,
|
||||||
"path": "embedding",
|
"path": "embedding",
|
||||||
},
|
},
|
||||||
|
|||||||
Reference in New Issue
Block a user