JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(oss): port v3 pipeline with hybrid search, entity extraction, and additive scoring (#4805)

Browse Source 
Co-authored-by: Soumil Rathi <soumilrathi@gmail.com>
Co-authored-by: Saket Aryan <saketaryan2002@gmail.com>
Co-authored-by: chaithanyak42 <chaithanya.kumar42a@gmail.com>
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
soumil-rathi
2026-04-14 18:00:58 +05:30
committed by GitHub
co-authored by Soumil Rathi Saket Aryan chaithanyak42 kartik-mem0
parent 57f944e18a
commit a488e19044
120 changed files with 10107 additions and 17135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/llms/test_langchain.py
+2
View File
@@ -2,6 +2,8 @@ from unittest.mock import Mock
import pytest
pytest.importorskip("langchain", reason="langchain not installed")
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.langchain import LangchainLLM
tests/llms/test_vllm.py
+26 -20
View File
@@ -144,6 +144,7 @@ def create_mocked_memory():
mock_embedder = MagicMock()
mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
mock_embedder.embed_batch.return_value = [[0.1, 0.2, 0.3]]
mock_embedder_factory.return_value = mock_embedder
mock_vector_store = MagicMock()
@@ -151,9 +152,12 @@ def create_mocked_memory():
mock_vector_store.add.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_sqlite.return_value = MagicMock()
mock_db = MagicMock()
mock_db.get_last_messages.return_value = []
mock_sqlite.return_value = mock_db
memory = Memory()
memory.custom_instructions = None
memory.api_version = "v1.0"
return memory, mock_llm, mock_vector_store
@@ -170,6 +174,7 @@ def create_mocked_async_memory():
mock_embedder = MagicMock()
mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
mock_embedder.embed_batch.return_value = [[0.1, 0.2, 0.3]]
mock_embedder_factory.return_value = mock_embedder
mock_vector_store = MagicMock()
@@ -177,9 +182,12 @@ def create_mocked_async_memory():
mock_vector_store.add.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_sqlite.return_value = MagicMock()
mock_db = MagicMock()
mock_db.get_last_messages.return_value = []
mock_sqlite.return_value = mock_db
memory = AsyncMemory()
memory.custom_instructions = None
memory.api_version = "v1.0"
return memory, mock_llm, mock_vector_store
@@ -187,22 +195,21 @@ def create_mocked_async_memory():
def test_thinking_tags_sync():
"""Test thinking tags handling in Memory._add_to_vector_store (sync)."""
memory, mock_llm, mock_vector_store = create_mocked_memory()
# Mock LLM responses for both phases
# v3 pipeline: single LLM call returning ADD-only memories
mock_llm.generate_response.side_effect = [
' <think>Sync fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
' <think>Sync memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
'<think>Sync extraction</think>\n{"memory": [{"text": "Loves sci-fi", "attributed_to": "user"}]}'
]
mock_vector_store.search.return_value = []
result = memory._add_to_vector_store(
messages=[{"role": "user", "content": "I love sci-fi movies"}],
metadata={},
filters={},
metadata={},
filters={},
infer=True
)
assert len(result) == 1
assert result[0]["memory"] == "Loves sci-fi"
assert result[0]["event"] == "ADD"
@@ -213,13 +220,12 @@ def test_thinking_tags_sync():
async def test_async_thinking_tags_async():
"""Test thinking tags handling in AsyncMemory._add_to_vector_store."""
memory, mock_llm, mock_vector_store = create_mocked_async_memory()
# Directly mock llm.generate_response instead of via asyncio.to_thread
# v3 pipeline: single LLM call returning ADD-only memories
mock_llm.generate_response.side_effect = [
' <think>Async fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
' <think>Async memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
'<think>Async extraction</think>\n{"memory": [{"text": "Loves sci-fi", "attributed_to": "user"}]}'
]
# Mock asyncio.to_thread to call the function directly (bypass threading)
async def mock_to_thread(func, *args, **kwargs):
if func == mock_llm.generate_response:
@@ -230,15 +236,15 @@ async def test_async_thinking_tags_async():
return []
else:
return func(*args, **kwargs)
with patch('mem0.memory.main.asyncio.to_thread', side_effect=mock_to_thread):
result = await memory._add_to_vector_store(
messages=[{"role": "user", "content": "I love sci-fi movies"}],
metadata={},
effective_filters={},
metadata={},
effective_filters={},
infer=True
)
assert len(result) == 1
assert result[0]["memory"] == "Loves sci-fi"
assert result[0]["event"] == "ADD"
tests/memory/test_apache_age_e2e.py
-1105
View File
File diff suppressed because it is too large Load Diff
tests/memory/test_apache_age_memory.py
-226
View File
@@ -1,226 +0,0 @@
from unittest.mock import MagicMock, Mock, patch
# age and rank_bm25 are optional deps — mock them so tests run without install
_age_mock = Mock()
patch.dict("sys.modules", {
"age": _age_mock,
"age.models": Mock(),
"rank_bm25": Mock(),
}).start()
from mem0.memory.apache_age_memory import MemoryGraph, _cosine_similarity # noqa: E402
def _make_instance():
with patch.object(MemoryGraph, "__init__", return_value=None):
instance = MemoryGraph.__new__(MemoryGraph)
instance.llm_provider = "openai"
instance.llm = MagicMock()
instance.embedding_model = MagicMock()
instance.config = MagicMock()
instance.config.graph_store.custom_prompt = None
instance.ag = MagicMock()
instance.graph_name = "test_graph"
instance.threshold = 0.7
return instance
class TestCosineSimilarity:
"""Tests for the _cosine_similarity helper."""
def test_identical_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [1, 0, 0]) - 1.0) < 1e-6
def test_orthogonal_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [0, 1, 0])) < 1e-6
def test_zero_vector(self):
assert _cosine_similarity([0, 0, 0], [1, 2, 3]) == 0.0
class TestRetrieveNodesFromData:
"""Tests for _retrieve_nodes_from_data in Apache AGE MemoryGraph."""
def test_normal_entities_extracted(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "Alice", "entity_type": "person"},
{"entity": "hiking", "entity_type": "activity"},
]}}]
}
result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
assert result == {"alice": "person", "hiking": "activity"}
def test_malformed_entity_missing_entity_type_is_skipped(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "matrix multiplication", "entity_type": "task"},
{"entity": "task"},
{"entity": "ReLU", "entity_type": "task"},
]}}]
}
result = instance._retrieve_nodes_from_data("some text", {"user_id": "u1"})
assert "matrix_multiplication" in result
assert "relu" in result
assert "task" not in result
def test_missing_entities_key_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
}
result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
assert result == {}
def test_none_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": None}
result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
assert result == {}
class TestEstablishNodesRelationsFromData:
"""Tests for _establish_nodes_relations_from_data in Apache AGE MemoryGraph."""
def test_none_response_does_not_crash(self):
instance = _make_instance()
instance.llm.generate_response.return_value = None
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_empty_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": []}
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_valid_entities_returned(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "add_entities", "arguments": {"entities": [
{"source": "alice", "relationship": "loves", "destination": "hiking"}
]}}]
}
result = instance._establish_nodes_relations_from_data(
"Alice loves hiking", {"user_id": "u1"}, {"alice": "person"}
)
assert len(result) == 1
assert result[0]["source"] == "alice"
class TestRemoveSpacesFromEntities:
"""Tests for _remove_spaces_from_entities."""
def test_spaces_and_case(self):
instance = _make_instance()
entities = [{"source": "Alice Smith", "relationship": "Works At", "destination": "Big Corp"}]
result = instance._remove_spaces_from_entities(entities)
assert result[0]["source"] == "alice_smith"
assert result[0]["relationship"] == "works_at"
assert result[0]["destination"] == "big_corp"
class TestFindSimilarNode:
"""Tests for _find_similar_node."""
def test_returns_none_when_no_nodes(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result is None
def test_returns_best_match_above_threshold(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1"},
{"name": "bob", "embedding": [0.0, 1.0], "user_id": "u1"},
])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result["name"] == "alice"
def test_filters_by_agent_id(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1", "agent_id": "a2"},
])
result = instance._find_similar_node(
[1.0, 0.0], {"user_id": "u1", "agent_id": "a1"}, threshold=0.9
)
assert result is None
class TestDeleteAll:
"""Tests for delete_all."""
def test_calls_exec_cypher_and_commits(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
instance.delete_all({"user_id": "u1"})
instance._exec_cypher.assert_called_once()
instance.ag.commit.assert_called_once()
class TestGetAll:
"""Tests for get_all."""
def test_returns_formatted_results(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "alice", "relationship": "KNOWS", "target": "bob"},
{"source": "alice", "relationship": "LIKES", "target": "hiking"},
])
results = instance.get_all({"user_id": "u1"}, top_k=10)
assert len(results) == 2
assert results[0]["source"] == "alice"
assert results[0]["relationship"] == "KNOWS"
assert results[0]["target"] == "bob"
def test_passes_limit_to_cypher(self):
"""Limit is enforced via LIMIT in the Cypher query, not Python slicing."""
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "n0", "relationship": "R", "target": "m0"},
])
instance.get_all({"user_id": "u1"}, top_k=3)
# Verify limit was passed as a parameter to the query
cypher_stmt = instance._exec_cypher.call_args[0][0]
assert "LIMIT %s" in cypher_stmt
params = instance._exec_cypher.call_args[1].get("params") or instance._exec_cypher.call_args[0][2]
assert 3 in params
class TestAdd:
"""Tests for the add orchestration method."""
def test_add_returns_added_and_deleted(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._establish_nodes_relations_from_data = MagicMock(return_value=[
{"source": "alice", "relationship": "knows", "destination": "bob"}
])
instance._search_graph_db = MagicMock(return_value=[])
instance._get_delete_entities_from_search_output = MagicMock(return_value=[])
instance._delete_entities = MagicMock(return_value=[])
instance._add_entities = MagicMock(return_value=["added"])
result = instance.add("Alice knows Bob", {"user_id": "u1"})
assert "deleted_entities" in result
assert "added_entities" in result
assert result["added_entities"] == ["added"]
class TestSearch:
"""Tests for the search method."""
def test_returns_empty_when_no_search_output(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._search_graph_db = MagicMock(return_value=[])
result = instance.search("Who is Alice?", {"user_id": "u1"})
assert result == []
tests/memory/test_graph_memory_soft_delete.py
-315
View File
@@ -1,315 +0,0 @@
"""Tests for graph memory soft-delete behavior.
Verifies that _delete_entities marks relationships as invalid (soft-delete)
rather than permanently removing them, and that search/retrieval queries
exclude soft-deleted relationships by default.
See: https://github.com/mem0ai/mem0/issues/4187
"""
from unittest.mock import Mock, patch
# Mock optional deps at module level so the import works across all Python
# versions without triggering transitive C-extension reloads (numpy via
# qdrant_client). This matches the pattern in test_memgraph_memory.py.
_neo4j_mock = Mock()
patch.dict("sys.modules", {
"langchain_neo4j": _neo4j_mock,
"rank_bm25": Mock(),
}).start()
from mem0.memory.graph_memory import MemoryGraph # noqa: E402
def _create_graph_memory():
"""Create a MemoryGraph instance with mocked dependencies."""
with patch.object(MemoryGraph, "__init__", lambda self, *a, **kw: None):
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Mock()
mg.graph.query = Mock(return_value=[])
mg.embedding_model = Mock()
mg.embedding_model.embed = Mock(return_value=[0.1] * 128)
mg.llm = Mock()
mg.node_label = ":Entity"
mg.threshold = 0.7
mg.llm_provider = "openai"
return mg
class TestSoftDelete:
"""Verify _delete_entities uses SET r.valid = false, not DELETE r."""
def test_delete_entities_sends_soft_delete_cypher(self):
mg = _create_graph_memory()
mg.graph.query.return_value = [
{"source": "Alice", "target": "Bob", "relationship": "KNOWS"}
]
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
)
cypher = mg.graph.query.call_args[0][0]
assert "SET r.valid = false" in cypher
assert "r.invalidated_at = datetime()" in cypher
assert "DELETE r" not in cypher
def test_delete_entities_only_targets_valid_edges(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
def test_delete_entities_is_idempotent(self):
mg = _create_graph_memory()
item = [{"source": "alice", "destination": "bob", "relationship": "KNOWS"}]
filters = {"user_id": "user1"}
mg.graph.query.return_value = [
{"source": "Alice", "target": "Bob", "relationship": "KNOWS"}
]
mg._delete_entities(item, filters)
mg.graph.query.return_value = []
mg._delete_entities(item, filters)
# Both calls should have the same WHERE filter
for c in mg.graph.query.call_args_list:
assert "r.valid IS NULL OR r.valid = true" in c[0][0]
class TestSearchExcludesSoftDeleted:
"""Verify search and get_all filter out soft-deleted relationships."""
def test_get_all_filters_soft_deleted(self):
mg = _create_graph_memory()
mg.get_all(filters={"user_id": "user1"}, top_k=10)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
def test_search_graph_db_filters_both_directions(self):
"""_search_graph_db must filter soft-deleted edges in both outgoing and incoming queries."""
mg = _create_graph_memory()
mg.graph.query.return_value = []
mg._search_graph_db(node_list=["alice"], filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
# The UNION query has two MATCH branches — both must filter
occurrences = cypher.count("r.valid IS NULL OR r.valid = true")
assert occurrences >= 2, (
f"_search_graph_db has {occurrences} valid-filter(s) but needs >= 2 "
"(one for outgoing, one for incoming relationships)"
)
def test_delete_all_still_hard_deletes(self):
mg = _create_graph_memory()
mg.delete_all(filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
class TestMergeResetsValidFlag:
"""Verify MERGE in _add_entities sets r.valid = true.
Critical: after soft-delete, a MERGE that matches the existing
(invalidated) edge must reset valid=true, or the edge becomes
a zombie -- exists but invisible to queries.
"""
def _run_add_entities(self, source_found, dest_found):
"""Helper: call _add_entities with configurable node search results."""
mg = _create_graph_memory()
source_result = (
[{"elementId(source_candidate)": "src_id_1"}] if source_found else []
)
dest_result = (
[{"elementId(destination_candidate)": "dst_id_1"}] if dest_found else []
)
mg._search_source_node = Mock(return_value=source_result)
mg._search_destination_node = Mock(return_value=dest_result)
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
entity_type_map={},
)
cypher = mg.graph.query.call_args[0][0]
return cypher
def test_merge_sets_valid_true_when_source_found(self):
cypher = self._run_add_entities(source_found=True, dest_found=False)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_dest_found(self):
cypher = self._run_add_entities(source_found=False, dest_found=True)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_both_found(self):
cypher = self._run_add_entities(source_found=True, dest_found=True)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_neither_found(self):
cypher = self._run_add_entities(source_found=False, dest_found=False)
assert "r.valid = true" in cypher
def test_merge_clears_invalidated_at_on_resurrection(self):
"""When a soft-deleted edge is resurrected via MERGE, invalidated_at must be cleared.
Without this, a resurrected edge (valid=true) still carries stale
invalidated_at metadata, which corrupts temporal reasoning queries.
"""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._run_add_entities(source_found=src, dest_found=dst)
assert "r.invalidated_at = null" in cypher, (
f"MERGE path '{label}': ON MATCH SET does not clear r.invalidated_at. "
"Resurrected edges will have stale invalidation timestamps."
)
class TestCypherConsistency:
"""Verify all MERGE blocks use consistent property names and variable aliases."""
def _get_merge_cypher(self, source_found, dest_found):
mg = _create_graph_memory()
mg._search_source_node = Mock(
return_value=[{"elementId(source_candidate)": "id1"}] if source_found else []
)
mg._search_destination_node = Mock(
return_value=[{"elementId(destination_candidate)": "id2"}] if dest_found else []
)
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
entity_type_map={},
)
return mg.graph.query.call_args[0][0]
def test_all_blocks_use_created_at_not_created(self):
"""All MERGE blocks must use r.created_at, not r.created."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "r.created_at" in cypher, (
f"MERGE path '{label}': uses r.created instead of r.created_at"
)
def test_all_blocks_use_r_not_rel(self):
"""All MERGE blocks must use 'r' as the relationship variable, not 'rel'."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "rel." not in cypher, (
f"MERGE path '{label}': uses 'rel' variable instead of 'r'"
)
def test_all_blocks_set_updated_at_on_create(self):
"""All MERGE blocks must set r.updated_at on CREATE for consistent timestamps."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "r.updated_at = timestamp()" in cypher, (
f"MERGE path '{label}': missing r.updated_at on CREATE SET"
)
class TestSoftDeleteWithFilters:
"""Verify soft-delete works correctly with agent_id and run_id filters."""
def test_delete_entities_with_agent_id(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "agent_id": "agent1"},
)
cypher = mg.graph.query.call_args[0][0]
params = mg.graph.query.call_args[1]["params"]
assert "SET r.valid = false" in cypher
assert "agent_id: $agent_id" in cypher
assert params["agent_id"] == "agent1"
def test_delete_entities_with_run_id(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "run_id": "run1"},
)
cypher = mg.graph.query.call_args[0][0]
params = mg.graph.query.call_args[1]["params"]
assert "SET r.valid = false" in cypher
assert "run_id: $run_id" in cypher
assert params["run_id"] == "run1"
def test_get_all_with_agent_id_filters_soft_deleted(self):
mg = _create_graph_memory()
mg.get_all(filters={"user_id": "user1", "agent_id": "agent1"}, top_k=10)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
assert "agent_id: $agent_id" in cypher
def test_merge_with_agent_id_sets_valid_true(self):
mg = _create_graph_memory()
mg._search_source_node = Mock(
return_value=[{"elementId(source_candidate)": "id1"}]
)
mg._search_destination_node = Mock(return_value=[])
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "agent_id": "agent1"},
entity_type_map={},
)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid = true" in cypher
assert "agent_id: $agent_id" in cypher
class TestResetAndCleanup:
"""Verify reset and delete_all use hard-delete (DETACH DELETE)."""
def test_reset_uses_detach_delete(self):
mg = _create_graph_memory()
mg.reset()
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
assert "valid" not in cypher.lower()
def test_delete_all_does_not_soft_delete(self):
mg = _create_graph_memory()
mg.delete_all(filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
assert "r.valid = false" not in cypher
tests/memory/test_json_prompt_fix.py
-28
View File
@@ -184,31 +184,3 @@ class TestEnsureJsonInstruction:
# Integration: verify fix is wired into both sync and async paths
# -------------------------------------------------------------------
def test_fix_applied_in_sync_memory_class(self):
"""Verify the ensure_json_instruction call exists in Memory._add_to_vector_store."""
import inspect
from mem0.memory.main import Memory
source = inspect.getsource(Memory._add_to_vector_store)
assert "ensure_json_instruction" in source, (
"ensure_json_instruction not found in Memory._add_to_vector_store (sync)"
)
def test_fix_applied_in_async_memory_class(self):
"""Verify the ensure_json_instruction call exists in AsyncMemory._add_to_vector_store."""
import inspect
from mem0.memory.main import AsyncMemory
source = inspect.getsource(AsyncMemory._add_to_vector_store)
assert "ensure_json_instruction" in source, (
"ensure_json_instruction not found in AsyncMemory._add_to_vector_store (async)"
)
def test_import_exists_in_main(self):
"""Verify ensure_json_instruction is imported in main.py."""
import inspect
import mem0.memory.main as main_module
source = inspect.getsource(main_module)
assert "from mem0.memory.utils import" in source
assert "ensure_json_instruction" in source
tests/memory/test_kuzu.py
-253
View File
@@ -1,253 +0,0 @@
from unittest.mock import MagicMock, Mock, patch
import numpy as np
import pytest
from mem0.memory.kuzu_memory import MemoryGraph
class TestKuzu:
"""Test that Kuzu memory works correctly"""
# Create distinct embeddings that won't match with threshold=0.7
# Each embedding is mostly zeros with ones in different positions to ensure low similarity
alice_emb = np.zeros(384)
alice_emb[0:96] = 1.0
bob_emb = np.zeros(384)
bob_emb[96:192] = 1.0
charlie_emb = np.zeros(384)
charlie_emb[192:288] = 1.0
dave_emb = np.zeros(384)
dave_emb[288:384] = 1.0
embeddings = {
"alice": alice_emb.tolist(),
"bob": bob_emb.tolist(),
"charlie": charlie_emb.tolist(),
"dave": dave_emb.tolist(),
}
@pytest.fixture
def mock_config(self):
"""Create a mock configuration for testing"""
config = Mock()
# Mock embedder config
config.embedder.provider = "mock_embedder"
config.embedder.config = {"model": "mock_model"}
config.vector_store.config = {"dimensions": 384}
# Mock graph store config
config.graph_store.config.db = ":memory:"
config.graph_store.threshold = 0.7
# Mock LLM config
config.llm.provider = "mock_llm"
config.llm.config = {"api_key": "test_key"}
return config
@pytest.fixture
def mock_embedding_model(self):
"""Create a mock embedding model"""
mock_model = Mock()
mock_model.config.embedding_dims = 384
def mock_embed(text):
return self.embeddings[text]
mock_model.embed.side_effect = mock_embed
return mock_model
@pytest.fixture
def mock_llm(self):
"""Create a mock LLM"""
mock_llm = Mock()
mock_llm.generate_response.return_value = {
"tool_calls": [
{
"name": "extract_entities",
"arguments": {"entities": [{"entity": "test_entity", "entity_type": "test_type"}]},
}
]
}
return mock_llm
@patch("mem0.memory.kuzu_memory.EmbedderFactory")
@patch("mem0.memory.kuzu_memory.LlmFactory")
def test_kuzu_memory_initialization(
self, mock_llm_factory, mock_embedder_factory, mock_config, mock_embedding_model, mock_llm
):
"""Test that Kuzu memory initializes correctly"""
# Setup mocks
mock_embedder_factory.create.return_value = mock_embedding_model
mock_llm_factory.create.return_value = mock_llm
# Create instance
kuzu_memory = MemoryGraph(mock_config)
# Verify initialization
assert kuzu_memory.config == mock_config
assert kuzu_memory.embedding_model == mock_embedding_model
assert kuzu_memory.embedding_dims == 384
assert kuzu_memory.llm == mock_llm
assert kuzu_memory.threshold == 0.7
@pytest.mark.parametrize(
"embedding_dims",
[None, 0, -1],
)
@patch("mem0.memory.kuzu_memory.EmbedderFactory")
def test_kuzu_memory_initialization_invalid_embedding_dims(
self, mock_embedder_factory, embedding_dims, mock_config
):
"""Test that Kuzu memory raises ValuError when initialized with invalid embedding_dims"""
# Setup mocks
mock_embedding_model = Mock()
mock_embedding_model.config.embedding_dims = embedding_dims
mock_embedder_factory.create.return_value = mock_embedding_model
with pytest.raises(ValueError, match="must be a positive"):
MemoryGraph(mock_config)
@patch("mem0.memory.kuzu_memory.EmbedderFactory")
@patch("mem0.memory.kuzu_memory.LlmFactory")
def test_kuzu(self, mock_llm_factory, mock_embedder_factory, mock_config, mock_embedding_model, mock_llm):
"""Test adding memory to the graph"""
mock_embedder_factory.create.return_value = mock_embedding_model
mock_llm_factory.create.return_value = mock_llm
kuzu_memory = MemoryGraph(mock_config)
filters = {"user_id": "test_user", "agent_id": "test_agent", "run_id": "test_run"}
data1 = [
{"source": "alice", "destination": "bob", "relationship": "knows"},
{"source": "bob", "destination": "charlie", "relationship": "knows"},
{"source": "charlie", "destination": "alice", "relationship": "knows"},
]
data2 = [
{"source": "charlie", "destination": "alice", "relationship": "likes"},
]
result = kuzu_memory._add_entities(data1, filters, {})
assert result[0] == [{"source": "alice", "relationship": "knows", "target": "bob"}]
assert result[1] == [{"source": "bob", "relationship": "knows", "target": "charlie"}]
assert result[2] == [{"source": "charlie", "relationship": "knows", "target": "alice"}]
assert get_node_count(kuzu_memory) == 3
assert get_edge_count(kuzu_memory) == 3
result = kuzu_memory._add_entities(data2, filters, {})
assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
assert get_node_count(kuzu_memory) == 3
assert get_edge_count(kuzu_memory) == 4
data3 = [
{"source": "dave", "destination": "alice", "relationship": "admires"}
]
result = kuzu_memory._add_entities(data3, filters, {})
assert result[0] == [{"source": "dave", "relationship": "admires", "target": "alice"}]
assert get_node_count(kuzu_memory) == 4 # dave is new
assert get_edge_count(kuzu_memory) == 5
results = kuzu_memory.get_all(filters)
assert set([f"{result['source']}_{result['relationship']}_{result['target']}" for result in results]) == set([
"alice_knows_bob",
"bob_knows_charlie",
"charlie_likes_alice",
"charlie_knows_alice",
"dave_admires_alice"
])
results = kuzu_memory._search_graph_db(["bob"], filters, threshold=0.8)
assert set([f"{result['source']}_{result['relationship']}_{result['destination']}" for result in results]) == set([
"alice_knows_bob",
"bob_knows_charlie",
])
result = kuzu_memory._delete_entities(data2, filters)
assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
assert get_node_count(kuzu_memory) == 4
assert get_edge_count(kuzu_memory) == 4
result = kuzu_memory._delete_entities(data1, filters)
assert result[0] == [{"source": "alice", "relationship": "knows", "target": "bob"}]
assert result[1] == [{"source": "bob", "relationship": "knows", "target": "charlie"}]
assert result[2] == [{"source": "charlie", "relationship": "knows", "target": "alice"}]
assert get_node_count(kuzu_memory) == 4
assert get_edge_count(kuzu_memory) == 1
result = kuzu_memory.delete_all(filters)
assert get_node_count(kuzu_memory) == 0
assert get_edge_count(kuzu_memory) == 0
result = kuzu_memory._add_entities(data2, filters, {})
assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
assert get_node_count(kuzu_memory) == 2
assert get_edge_count(kuzu_memory) == 1
result = kuzu_memory.reset()
assert get_node_count(kuzu_memory) == 0
assert get_edge_count(kuzu_memory) == 0
def _make_kuzu_instance():
with patch.object(MemoryGraph, "__init__", return_value=None):
instance = MemoryGraph.__new__(MemoryGraph)
instance.llm_provider = "openai"
instance.llm = MagicMock()
instance.embedding_model = MagicMock()
instance.config = MagicMock()
instance.config.graph_store.custom_prompt = None
return instance
class TestRetrieveNodesFromData:
"""Tests for _retrieve_nodes_from_data in KuzuMemoryGraph."""
def test_missing_entities_key_returns_empty(self):
"""LLM returns extract_entities tool call without 'entities' key — should not crash.
Reproduces the exact scenario from issue #4238."""
instance = _make_kuzu_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
}
result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
assert result == {}
def test_normal_entities_extracted(self):
instance = _make_kuzu_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "Alice", "entity_type": "person"},
{"entity": "hiking", "entity_type": "activity"},
]}}]
}
result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
assert result == {"alice": "person", "hiking": "activity"}
def test_none_tool_calls_returns_empty(self):
instance = _make_kuzu_instance()
instance.llm.generate_response.return_value = {"tool_calls": None}
result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
assert result == {}
def get_node_count(kuzu_memory):
results = kuzu_memory.kuzu_execute(
"""
MATCH (n)
RETURN COUNT(n) as count
"""
)
return int(results[0]['count'])
def get_edge_count(kuzu_memory):
results = kuzu_memory.kuzu_execute(
"""
MATCH (n)-[e]->(m)
RETURN COUNT(e) as count
"""
)
return int(results[0]['count'])
tests/memory/test_main.py
+35 -273
View File
@@ -4,7 +4,7 @@ from unittest.mock import MagicMock, Mock
import pytest
from mem0.memory.main import AsyncMemory, Memory, _normalize_iso_timestamp_to_utc
from mem0.memory.main import AsyncMemory, Memory
def _setup_mocks(mocker):
@@ -37,16 +37,19 @@ class TestAddToVectorStoreErrors:
memory.config = mocker.MagicMock()
memory.config.custom_instructions = None
memory.config.custom_update_memory_prompt = None
memory.custom_instructions = None
memory.api_version = "v1.1"
# v3 pipeline needs db.get_last_messages to return a list
memory.db.get_last_messages = MagicMock(return_value=[])
memory.db.save_messages = MagicMock()
return memory
def test_empty_llm_response_fact_extraction(self, mocker, mock_memory, caplog):
"""Test empty response from LLM during fact extraction"""
"""Test invalid JSON response from LLM during extraction"""
# Setup
mock_memory.llm.generate_response.return_value = "invalid json" # This will trigger a JSON decode error
mock_capture_event = mocker.MagicMock()
mocker.patch("mem0.memory.main.capture_event", mock_capture_event)
mock_memory.llm.generate_response.return_value = "invalid json"
mocker.patch("mem0.memory.main.capture_event")
# Execute
with caplog.at_level(logging.ERROR):
@@ -54,18 +57,15 @@ class TestAddToVectorStoreErrors:
messages=[{"role": "user", "content": "test"}], metadata={}, filters={}, infer=True
)
# Verify
# Verify — v3 single-pass pipeline makes 1 LLM call, returns [] on parse error
assert mock_memory.llm.generate_response.call_count == 1
assert result == [] # Should return empty list when no memories processed
# Check for error message in any of the log records
assert any("Error in new_retrieved_facts" in record.msg for record in caplog.records), "Expected error message not found in logs"
assert mock_capture_event.call_count == 1
assert result == []
assert any("Error parsing extraction response" in record.message for record in caplog.records), "Expected error message not found in logs"
def test_empty_llm_response_memory_actions(self, mock_memory, caplog):
"""Test empty response from LLM during memory actions"""
# Setup
# First call returns valid JSON, second call returns empty string
mock_memory.llm.generate_response.side_effect = ['{"facts": ["test fact"]}', ""]
"""Test empty response from LLM during memory actions (v3: single-pass, 1 LLM call)"""
# Setup — v3 pipeline does a single LLM call that returns empty/invalid response
mock_memory.llm.generate_response.return_value = ""
# Execute
with caplog.at_level(logging.WARNING):
@@ -73,10 +73,9 @@ class TestAddToVectorStoreErrors:
messages=[{"role": "user", "content": "test"}], metadata={}, filters={}, infer=True
)
# Verify
assert mock_memory.llm.generate_response.call_count == 2
# Verify — v3 only makes 1 LLM call (no separate merge step)
assert mock_memory.llm.generate_response.call_count == 1
assert result == [] # Should return empty list when no memories processed
assert "Empty response from LLM, no memories to extract" in caplog.text
class TestAsyncUpdate:
@@ -141,17 +140,20 @@ class TestAsyncAddToVectorStoreErrors:
memory.config = mocker.MagicMock()
memory.config.custom_instructions = None
memory.config.custom_update_memory_prompt = None
memory.custom_instructions = None
memory.api_version = "v1.1"
# v3 pipeline needs db.get_last_messages to return a list
memory.db.get_last_messages = MagicMock(return_value=[])
memory.db.save_messages = MagicMock()
return memory
@pytest.mark.asyncio
async def test_async_empty_llm_response_fact_extraction(self, mock_async_memory, caplog, mocker):
"""Test empty response in AsyncMemory._add_to_vector_store"""
"""Test invalid JSON response from LLM during extraction (async)"""
mocker.patch("mem0.utils.factory.EmbedderFactory.create", return_value=MagicMock())
mock_async_memory.llm.generate_response.return_value = "invalid json" # This will trigger a JSON decode error
mock_capture_event = mocker.MagicMock()
mocker.patch("mem0.memory.main.capture_event", mock_capture_event)
mock_async_memory.llm.generate_response.return_value = "invalid json"
mocker.patch("mem0.memory.main.capture_event")
with caplog.at_level(logging.ERROR):
result = await mock_async_memory._add_to_vector_store(
@@ -159,15 +161,13 @@ class TestAsyncAddToVectorStoreErrors:
)
assert mock_async_memory.llm.generate_response.call_count == 1
assert result == []
# Check for error message in any of the log records
assert any("Error in new_retrieved_facts" in record.msg for record in caplog.records), "Expected error message not found in logs"
assert mock_capture_event.call_count == 1
assert any("Error parsing extraction response" in record.message for record in caplog.records), "Expected error message not found in logs"
@pytest.mark.asyncio
async def test_async_empty_llm_response_memory_actions(self, mock_async_memory, caplog, mocker):
"""Test empty response in AsyncMemory._add_to_vector_store"""
"""Test empty response in AsyncMemory._add_to_vector_store (v3: single-pass, 1 LLM call)"""
mocker.patch("mem0.utils.factory.EmbedderFactory.create", return_value=MagicMock())
mock_async_memory.llm.generate_response.side_effect = ['{"facts": ["test fact"]}', ""]
mock_async_memory.llm.generate_response.return_value = ""
mock_capture_event = mocker.MagicMock()
mocker.patch("mem0.memory.main.capture_event", mock_capture_event)
@@ -177,8 +177,7 @@ class TestAsyncAddToVectorStoreErrors:
)
assert result == []
assert "Empty response from LLM, no memories to extract" in caplog.text
assert mock_capture_event.call_count == 1
assert mock_async_memory.llm.generate_response.call_count == 1
def _build_memory_instance(mocker, memory_cls):
@@ -237,8 +236,8 @@ def test_update_memory_uses_utc_timestamps(mocker):
)
memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
payload = memory.vector_store.update.call_args.kwargs["payload"]
assert payload["created_at"] == "2026-03-18T00:00:00+00:00"
_assert_utc_timestamp(payload["updated_at"])
assert payload["created_at"] == "2026-03-17T17:00:00-07:00"
assert payload["updated_at"] is not None
@pytest.mark.asyncio
@@ -281,8 +280,8 @@ async def test_async_update_memory_uses_utc_timestamps(mocker):
)
await memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
payload = memory.vector_store.update.call_args.kwargs["payload"]
assert payload["created_at"] == "2026-03-18T00:00:00+00:00"
_assert_utc_timestamp(payload["updated_at"])
assert payload["created_at"] == "2026-03-17T17:00:00-07:00"
assert payload["updated_at"] is not None
def test_create_then_search_and_get_all_return_same_timestamps(mocker):
@@ -309,8 +308,8 @@ def test_create_then_search_and_get_all_return_same_timestamps(mocker):
memory.vector_store.list.return_value = [[mem_result]]
# Step 3: Call search and get_all, compare timestamps
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, top_k=10, threshold=None)
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, limit=10)
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, limit=100)
search_item = search_results[0]
get_all_item = get_all_results[0]
@@ -376,8 +375,8 @@ def test_search_and_get_all_consistent_after_update(mocker):
memory.vector_store.search.return_value = [mem_result]
memory.vector_store.list.return_value = [[mem_result]]
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, top_k=10, threshold=None)
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, limit=10)
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, limit=100)
assert search_results[0]["created_at"] == get_all_results[0]["created_at"]
assert search_results[0]["updated_at"] == get_all_results[0]["updated_at"]
@@ -627,240 +626,3 @@ async def test_async_update_preserves_actor_id_when_different_actor_updates(mock
assert stored["actor_id"] == "Alice"
class TestHallucinatedIdGuard:
"""Tests for temp_uuid_mapping guard against LLM-hallucinated IDs (issue #3931).
When the LLM returns an UPDATE or DELETE with an ID that doesn't exist in
temp_uuid_mapping, the code should skip gracefully instead of raising KeyError.
"""
def test_sync_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
"""Sync UPDATE with an out-of-range ID should be skipped with a warning."""
memory = _build_memory_instance(mocker, Memory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
# Simulate: 2 existing memories (IDs "0" and "1"), but LLM returns UPDATE for ID "12"
existing_mem = MagicMock()
existing_mem.id = "uuid-aaa"
existing_mem.payload = {"data": "User likes coffee"}
memory.vector_store.search.return_value = [existing_mem]
# First LLM call: fact extraction → returns one fact
# Second LLM call: memory update actions → returns UPDATE with hallucinated ID "12"
memory.llm.generate_response.side_effect = [
'{"facts": ["User likes tea"]}',
'{"memory": [{"id": "12", "text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
]
with caplog.at_level(logging.WARNING):
result = memory._add_to_vector_store(
messages=[{"role": "user", "content": "I like tea"}],
metadata={},
filters={},
infer=True,
)
# Should not crash, should return empty (the hallucinated UPDATE was skipped)
assert result == []
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
# _update_memory should NOT have been called
memory.vector_store.update.assert_not_called()
def test_sync_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
"""Sync DELETE with an out-of-range ID should be skipped with a warning."""
memory = _build_memory_instance(mocker, Memory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-aaa"
existing_mem.payload = {"data": "User likes coffee"}
memory.vector_store.search.return_value = [existing_mem]
memory.llm.generate_response.side_effect = [
'{"facts": ["Remove coffee preference"]}',
'{"memory": [{"id": "9", "text": "User likes coffee", "event": "DELETE"}]}',
]
with caplog.at_level(logging.WARNING):
result = memory._add_to_vector_store(
messages=[{"role": "user", "content": "I no longer like coffee"}],
metadata={},
filters={},
infer=True,
)
assert result == []
assert "DELETE skipped: LLM returned unknown id" in caplog.text
memory.vector_store.delete.assert_not_called()
def test_sync_valid_id_still_processes_normally(self, mocker, caplog):
"""A valid ID should still be processed — the guard must not block legitimate operations."""
memory = _build_memory_instance(mocker, Memory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-aaa"
existing_mem.payload = {"data": "User likes coffee"}
memory.vector_store.search.return_value = [existing_mem]
memory.vector_store.get.return_value = MagicMock(
payload={"data": "User likes coffee", "created_at": "2026-01-01T00:00:00+00:00"}
)
# ID "0" is valid since there's exactly 1 existing memory
memory.llm.generate_response.side_effect = [
'{"facts": ["User likes tea now"]}',
'{"memory": [{"id": "0", "text": "User likes tea now", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
]
with caplog.at_level(logging.WARNING):
result = memory._add_to_vector_store(
messages=[{"role": "user", "content": "I like tea now"}],
metadata={},
filters={},
infer=True,
)
assert len(result) == 1
assert result[0]["event"] == "UPDATE"
assert result[0]["memory"] == "User likes tea now"
assert result[0]["id"] == "uuid-aaa"
assert "skipped" not in caplog.text
@pytest.mark.asyncio
async def test_async_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
"""Async UPDATE with an out-of-range ID should be skipped with a warning."""
memory = _build_memory_instance(mocker, AsyncMemory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-bbb"
existing_mem.payload = {"data": "User works at Acme"}
memory.vector_store.search.return_value = [existing_mem]
memory.llm.generate_response.side_effect = [
'{"facts": ["User works at Globex"]}',
'{"memory": [{"id": "7", "text": "User works at Globex", "event": "UPDATE", "old_memory": "User works at Acme"}]}',
]
with caplog.at_level(logging.WARNING):
result = await memory._add_to_vector_store(
messages=[{"role": "user", "content": "I now work at Globex"}],
metadata={},
effective_filters={},
infer=True,
)
assert result == []
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
@pytest.mark.asyncio
async def test_async_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
"""Async DELETE with an out-of-range ID should be skipped with a warning."""
memory = _build_memory_instance(mocker, AsyncMemory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-ccc"
existing_mem.payload = {"data": "User lives in SF"}
memory.vector_store.search.return_value = [existing_mem]
memory.llm.generate_response.side_effect = [
'{"facts": ["Remove SF reference"]}',
'{"memory": [{"id": "16", "text": "User lives in SF", "event": "DELETE"}]}',
]
with caplog.at_level(logging.WARNING):
result = await memory._add_to_vector_store(
messages=[{"role": "user", "content": "I moved away from SF"}],
metadata={},
effective_filters={},
infer=True,
)
assert result == []
assert "DELETE skipped: LLM returned unknown id" in caplog.text
def test_sync_update_with_missing_id_key_skips_gracefully(self, mocker, caplog):
"""UPDATE where the LLM omits the 'id' field entirely should be skipped."""
memory = _build_memory_instance(mocker, Memory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-aaa"
existing_mem.payload = {"data": "User likes coffee"}
memory.vector_store.search.return_value = [existing_mem]
# LLM response has no "id" key at all
memory.llm.generate_response.side_effect = [
'{"facts": ["User likes tea"]}',
'{"memory": [{"text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
]
with caplog.at_level(logging.WARNING):
result = memory._add_to_vector_store(
messages=[{"role": "user", "content": "I like tea"}],
metadata={},
filters={},
infer=True,
)
assert result == []
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
@pytest.mark.asyncio
async def test_async_valid_id_still_processes_normally(self, mocker, caplog):
"""Async path: a valid ID should process normally — no false positives from the guard."""
memory = _build_memory_instance(mocker, AsyncMemory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.memory.main.capture_event")
existing_mem = MagicMock()
existing_mem.id = "uuid-bbb"
existing_mem.payload = {"data": "User works at Acme"}
memory.vector_store.search.return_value = [existing_mem]
memory.vector_store.get.return_value = MagicMock(
payload={"data": "User works at Acme", "created_at": "2026-01-01T00:00:00+00:00"}
)
# ID "0" is valid since there's exactly 1 existing memory
memory.llm.generate_response.side_effect = [
'{"facts": ["User works at Globex now"]}',
'{"memory": [{"id": "0", "text": "User works at Globex now", "event": "UPDATE", "old_memory": "User works at Acme"}]}',
]
with caplog.at_level(logging.WARNING):
result = await memory._add_to_vector_store(
messages=[{"role": "user", "content": "I now work at Globex"}],
metadata={},
effective_filters={},
infer=True,
)
assert len(result) == 1
assert result[0]["event"] == "UPDATE"
assert result[0]["memory"] == "User works at Globex now"
assert result[0]["id"] == "uuid-bbb"
assert "skipped" not in caplog.text
def test_normalize_iso_timestamp_to_utc_preserves_naive_values():
assert _normalize_iso_timestamp_to_utc("2026-03-18T00:00:00") == "2026-03-18T00:00:00"
def test_normalize_iso_timestamp_to_utc_converts_pacific():
result = _normalize_iso_timestamp_to_utc("2026-03-17T17:00:00-07:00")
assert result == "2026-03-18T00:00:00+00:00"
def test_normalize_iso_timestamp_to_utc_handles_none():
assert _normalize_iso_timestamp_to_utc(None) is None
def test_normalize_iso_timestamp_to_utc_handles_empty():
assert _normalize_iso_timestamp_to_utc("") == ""
tests/memory/test_memgraph_memory.py
-107
View File
@@ -1,107 +0,0 @@
from unittest.mock import MagicMock, Mock, patch
# langchain_memgraph and rank_bm25 are optional deps — mock them so tests run without install
_memgraph_mock = Mock()
patch.dict("sys.modules", {
"langchain_memgraph": _memgraph_mock,
"langchain_memgraph.graphs": _memgraph_mock,
"langchain_memgraph.graphs.memgraph": _memgraph_mock,
"rank_bm25": Mock(),
}).start()
from mem0.memory.memgraph_memory import MemoryGraph as MemgraphMemoryGraph # noqa: E402
MemoryGraph = MemgraphMemoryGraph
def _make_instance():
with patch.object(MemoryGraph, "__init__", return_value=None):
instance = MemoryGraph.__new__(MemoryGraph)
instance.llm_provider = "openai"
instance.llm = MagicMock()
instance.embedding_model = MagicMock()
instance.config = MagicMock()
instance.config.graph_store.custom_prompt = None
return instance
class TestRetrieveNodesFromData:
"""Tests for _retrieve_nodes_from_data in MemoryGraph."""
def test_normal_entities_extracted(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "Alice", "entity_type": "person"},
{"entity": "hiking", "entity_type": "activity"},
]}}]
}
result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
assert result == {"alice": "person", "hiking": "activity"}
def test_malformed_entity_missing_entity_type_is_skipped(self):
"""LLM returns entity dict without entity_type — should skip it, keep valid ones.
Reproduces the exact data from issue #4055."""
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "matrix multiplication", "entity_type": "task"},
{"entity": "task"},
{"entity": "ReLU", "entity_type": "task"},
]}}]
}
result = instance._retrieve_nodes_from_data("some text", {"user_id": "u1"})
assert "matrix_multiplication" in result
assert "relu" in result
assert "task" not in result
def test_missing_entities_key_returns_empty(self):
"""LLM returns extract_entities tool call without 'entities' key — should not crash.
Reproduces the exact scenario from issue #4238."""
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
}
result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
assert result == {}
def test_none_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": None}
result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
assert result == {}
class TestEstablishNodesRelationsFromData:
"""Tests for _establish_nodes_relations_from_data in MemoryGraph."""
def test_none_response_does_not_crash(self):
"""openai_structured returns None when no relations found — must not crash.
Exact crash from issue #4055: TypeError: 'NoneType' object is not subscriptable."""
instance = _make_instance()
instance.llm.generate_response.return_value = None
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_empty_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": []}
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_valid_entities_returned(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "add_entities", "arguments": {"entities": [
{"source": "alice", "relationship": "loves", "destination": "hiking"}
]}}]
}
result = instance._establish_nodes_relations_from_data(
"Alice loves hiking", {"user_id": "u1"}, {"alice": "person"}
)
assert len(result) == 1
assert result[0]["source"] == "alice"
tests/memory/test_neo4j_cypher_syntax.py
-318
View File
@@ -1,318 +0,0 @@
import os
from unittest.mock import Mock, patch
from mem0.memory.utils import sanitize_relationship_for_cypher
class TestSanitizeRelationshipForCypher:
"""Test that relationship names are properly sanitized for Neo4j Cypher queries."""
def test_hyphen_replaced_with_underscore(self):
"""Hyphens in relationship names cause Neo4j CypherSyntaxError and must be replaced."""
assert sanitize_relationship_for_cypher("manages_via_low-cost_models") == "manages_via_low_cost_models"
def test_multiple_hyphens(self):
assert sanitize_relationship_for_cypher("co-owns-with") == "co_owns_with"
def test_no_special_chars_unchanged(self):
assert sanitize_relationship_for_cypher("works_at") == "works_at"
def test_spaces_not_handled_here(self):
"""Spaces are replaced upstream before this function is called."""
# sanitize only handles special chars, spaces are handled by the caller
result = sanitize_relationship_for_cypher("has relationship")
assert result == "has relationship"
def test_existing_chars_still_sanitized(self):
assert "_slash_" in sanitize_relationship_for_cypher("read/write")
assert "_at_" in sanitize_relationship_for_cypher("user@company")
class TestNeo4jCypherSyntaxFix:
"""Test that Neo4j Cypher syntax fixes work correctly"""
def test_get_all_generates_valid_cypher_with_agent_id(self):
"""Test that get_all method generates valid Cypher with agent_id"""
# Mock the langchain_neo4j module to avoid import issues
with patch.dict('sys.modules', {'langchain_neo4j': Mock()}):
from mem0.memory.graph_memory import MemoryGraph
# Create instance (will fail on actual connection, but that's fine for syntax testing)
try:
_ = MemoryGraph(url="bolt://localhost:7687", username="test", password="test")
except Exception:
# Expected to fail on connection, just test the class exists
assert MemoryGraph is not None
return
def test_cypher_syntax_validation(self):
"""Test that our Cypher fixes don't contain problematic patterns"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Ensure the old buggy pattern is not present
assert "AND n.agent_id = $agent_id AND m.agent_id = $agent_id" not in content
assert "WHERE 1=1 {agent_filter}" not in content
# Ensure proper node property syntax is present
assert "node_props" in content
assert "agent_id: $agent_id" in content
# Ensure run_id follows the same pattern
# Check for absence of problematic run_id patterns
assert "AND n.run_id = $run_id AND m.run_id = $run_id" not in content
assert "WHERE 1=1 {run_id_filter}" not in content
def test_no_undefined_variables_in_cypher(self):
"""Test that we don't have undefined variable patterns"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check for patterns that would cause "Variable 'm' not defined" errors
lines = content.split('\n')
for i, line in enumerate(lines):
# Look for WHERE clauses that reference variables not in MATCH
if 'WHERE' in line and 'm.agent_id' in line:
# Check if there's a MATCH clause before this that defines 'm'
preceding_lines = lines[max(0, i-10):i]
match_found = any('MATCH' in prev_line and ' m ' in prev_line for prev_line in preceding_lines)
assert match_found, f"Line {i+1}: WHERE clause references 'm' without MATCH definition"
# Also check for run_id patterns that might have similar issues
if 'WHERE' in line and 'm.run_id' in line:
# Check if there's a MATCH clause before this that defines 'm'
preceding_lines = lines[max(0, i-10):i]
match_found = any('MATCH' in prev_line and ' m ' in prev_line for prev_line in preceding_lines)
assert match_found, f"Line {i+1}: WHERE clause references 'm.run_id' without MATCH definition"
def test_agent_id_integration_syntax(self):
"""Test that agent_id is properly integrated into MATCH clauses"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Should have node property building logic
assert 'node_props = [' in content
assert 'node_props.append("agent_id: $agent_id")' in content
assert 'node_props_str = ", ".join(node_props)' in content
# Should use the node properties in MATCH clauses
assert '{{{node_props_str}}}' in content or '{node_props_str}' in content
def test_run_id_integration_syntax(self):
"""Test that run_id is properly integrated into MATCH clauses"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Should have node property building logic for run_id
assert 'node_props = [' in content
assert 'node_props.append("run_id: $run_id")' in content
assert 'node_props_str = ", ".join(node_props)' in content
# Should use the node properties in MATCH clauses
assert '{{{node_props_str}}}' in content or '{node_props_str}' in content
def test_agent_id_filter_patterns(self):
"""Test that agent_id filtering follows the correct pattern"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that agent_id is handled in filters
assert 'if filters.get("agent_id"):' in content
assert 'params["agent_id"] = filters["agent_id"]' in content
# Check that agent_id is used in node properties
assert 'node_props.append("agent_id: $agent_id")' in content
def test_run_id_filter_patterns(self):
"""Test that run_id filtering follows the same pattern as agent_id"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that run_id is handled in filters
assert 'if filters.get("run_id"):' in content
assert 'params["run_id"] = filters["run_id"]' in content
# Check that run_id is used in node properties
assert 'node_props.append("run_id: $run_id")' in content
def test_agent_id_cypher_generation(self):
"""Test that agent_id is properly included in Cypher query generation"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that the dynamic property building pattern exists
assert 'node_props = [' in content
assert 'node_props_str = ", ".join(node_props)' in content
# Check that agent_id is handled in the pattern
assert 'if filters.get(' in content
assert 'node_props.append(' in content
# Verify the pattern is used in MATCH clauses
assert '{{{node_props_str}}}' in content or '{node_props_str}' in content
def test_run_id_cypher_generation(self):
"""Test that run_id is properly included in Cypher query generation"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that the dynamic property building pattern exists
assert 'node_props = [' in content
assert 'node_props_str = ", ".join(node_props)' in content
# Check that run_id is handled in the pattern
assert 'if filters.get(' in content
assert 'node_props.append(' in content
# Verify the pattern is used in MATCH clauses
assert '{{{node_props_str}}}' in content or '{node_props_str}' in content
def test_agent_id_implementation_pattern(self):
"""Test that the code structure supports agent_id implementation"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Verify that agent_id pattern is used consistently
assert 'node_props = [' in content
assert 'node_props_str = ", ".join(node_props)' in content
assert 'if filters.get("agent_id"):' in content
assert 'node_props.append("agent_id: $agent_id")' in content
def test_run_id_implementation_pattern(self):
"""Test that the code structure supports run_id implementation"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Verify that run_id pattern is used consistently
assert 'node_props = [' in content
assert 'node_props_str = ", ".join(node_props)' in content
assert 'if filters.get("run_id"):' in content
assert 'node_props.append("run_id: $run_id")' in content
def test_user_identity_integration(self):
"""Test that both agent_id and run_id are properly integrated into user identity"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that user_identity building includes both agent_id and run_id
assert 'user_identity = f"user_id: {filters[\'user_id\']}"' in content
assert 'user_identity += f", agent_id: {filters[\'agent_id\']}"' in content
assert 'user_identity += f", run_id: {filters[\'run_id\']}"' in content
def test_search_methods_integration(self):
"""Test that both agent_id and run_id are properly integrated into search methods"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that search methods handle both agent_id and run_id
assert 'where_conditions.append("source_candidate.agent_id = $agent_id")' in content
assert 'where_conditions.append("source_candidate.run_id = $run_id")' in content
assert 'where_conditions.append("destination_candidate.agent_id = $agent_id")' in content
assert 'where_conditions.append("destination_candidate.run_id = $run_id")' in content
def test_add_entities_integration(self):
"""Test that both agent_id and run_id are properly integrated into add_entities"""
graph_memory_path = 'mem0/memory/graph_memory.py'
# Check if file exists before reading
if not os.path.exists(graph_memory_path):
# Skip test if file doesn't exist (e.g., in CI environment)
return
with open(graph_memory_path, 'r') as f:
content = f.read()
# Check that add_entities handles both agent_id and run_id
assert 'agent_id = filters.get("agent_id", None)' in content
assert 'run_id = filters.get("run_id", None)' in content
# Check that merge properties include both
assert 'if agent_id:' in content
assert 'if run_id:' in content
assert 'merge_props.append("agent_id: $agent_id")' in content
assert 'merge_props.append("run_id: $run_id")' in content
tests/memory/test_neptune_analytics_memory.py
-338
View File
@@ -1,338 +0,0 @@
import unittest
from unittest.mock import MagicMock, patch
import pytest
from mem0.graphs.neptune.base import NeptuneBase
from mem0.graphs.neptune.neptunegraph import MemoryGraph
class TestNeptuneMemory(unittest.TestCase):
"""Test suite for the Neptune Memory implementation."""
def setUp(self):
"""Set up test fixtures before each test method."""
# Create a mock config
self.config = MagicMock()
self.config.graph_store.config.endpoint = "neptune-graph://test-graph"
self.config.graph_store.config.base_label = True
self.config.graph_store.threshold = 0.7
self.config.llm.provider = "openai_structured"
self.config.graph_store.llm = None
self.config.graph_store.custom_prompt = None
# Create mock for NeptuneAnalyticsGraph
self.mock_graph = MagicMock()
self.mock_graph.client.get_graph.return_value = {"status": "AVAILABLE"}
# Create mocks for static methods
self.mock_embedding_model = MagicMock()
self.mock_llm = MagicMock()
# Patch the necessary components
self.neptune_analytics_graph_patcher = patch("mem0.graphs.neptune.neptunegraph.NeptuneAnalyticsGraph")
self.mock_neptune_analytics_graph = self.neptune_analytics_graph_patcher.start()
self.mock_neptune_analytics_graph.return_value = self.mock_graph
# Patch the static methods
self.create_embedding_model_patcher = patch.object(NeptuneBase, "_create_embedding_model")
self.mock_create_embedding_model = self.create_embedding_model_patcher.start()
self.mock_create_embedding_model.return_value = self.mock_embedding_model
self.create_llm_patcher = patch.object(NeptuneBase, "_create_llm")
self.mock_create_llm = self.create_llm_patcher.start()
self.mock_create_llm.return_value = self.mock_llm
# Create the MemoryGraph instance
self.memory_graph = MemoryGraph(self.config)
# Set up common test data
self.user_id = "test_user"
self.test_filters = {"user_id": self.user_id}
def tearDown(self):
"""Tear down test fixtures after each test method."""
self.neptune_analytics_graph_patcher.stop()
self.create_embedding_model_patcher.stop()
self.create_llm_patcher.stop()
def test_initialization(self):
"""Test that the MemoryGraph is initialized correctly."""
self.assertEqual(self.memory_graph.graph, self.mock_graph)
self.assertEqual(self.memory_graph.embedding_model, self.mock_embedding_model)
self.assertEqual(self.memory_graph.llm, self.mock_llm)
self.assertEqual(self.memory_graph.llm_provider, "openai_structured")
self.assertEqual(self.memory_graph.node_label, ":`__Entity__`")
self.assertEqual(self.memory_graph.threshold, 0.7)
def test_init(self):
"""Test the class init functions"""
# Create a mock config with bad endpoint
config_no_endpoint = MagicMock()
config_no_endpoint.graph_store.config.endpoint = None
# Create the MemoryGraph instance
with pytest.raises(ValueError):
MemoryGraph(config_no_endpoint)
# Create a mock config with bad endpoint
config_ndb_endpoint = MagicMock()
config_ndb_endpoint.graph_store.config.endpoint = "neptune-db://test-graph"
with pytest.raises(ValueError):
MemoryGraph(config_ndb_endpoint)
def test_add_method(self):
"""Test the add method with mocked components."""
# Mock the necessary methods that add() calls
self.memory_graph._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person", "bob": "person"})
self.memory_graph._establish_nodes_relations_from_data = MagicMock(
return_value=[{"source": "alice", "relationship": "knows", "destination": "bob"}]
)
self.memory_graph._search_graph_db = MagicMock(return_value=[])
self.memory_graph._get_delete_entities_from_search_output = MagicMock(return_value=[])
self.memory_graph._delete_entities = MagicMock(return_value=[])
self.memory_graph._add_entities = MagicMock(
return_value=[{"source": "alice", "relationship": "knows", "target": "bob"}]
)
# Call the add method
result = self.memory_graph.add("Alice knows Bob", self.test_filters)
# Verify the method calls
self.memory_graph._retrieve_nodes_from_data.assert_called_once_with("Alice knows Bob", self.test_filters)
self.memory_graph._establish_nodes_relations_from_data.assert_called_once()
self.memory_graph._search_graph_db.assert_called_once()
self.memory_graph._get_delete_entities_from_search_output.assert_called_once()
self.memory_graph._delete_entities.assert_called_once_with([], self.user_id)
self.memory_graph._add_entities.assert_called_once()
# Check the result structure
self.assertIn("deleted_entities", result)
self.assertIn("added_entities", result)
def test_search_method(self):
"""Test the search method with mocked components."""
# Mock the necessary methods that search() calls
self.memory_graph._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
# Mock search results
mock_search_results = [
{"source": "alice", "relationship": "knows", "destination": "bob"},
{"source": "alice", "relationship": "works_with", "destination": "charlie"},
]
self.memory_graph._search_graph_db = MagicMock(return_value=mock_search_results)
# Mock BM25Okapi
with patch("mem0.graphs.neptune.base.BM25Okapi") as mock_bm25:
mock_bm25_instance = MagicMock()
mock_bm25.return_value = mock_bm25_instance
# Mock get_top_n to return reranked results
reranked_results = [["alice", "knows", "bob"], ["alice", "works_with", "charlie"]]
mock_bm25_instance.get_top_n.return_value = reranked_results
# Call the search method
result = self.memory_graph.search("Find Alice", self.test_filters, top_k=5)
# Verify the method calls
self.memory_graph._retrieve_nodes_from_data.assert_called_once_with("Find Alice", self.test_filters)
self.memory_graph._search_graph_db.assert_called_once_with(node_list=["alice"], filters=self.test_filters)
# Check the result structure
self.assertEqual(len(result), 2)
self.assertEqual(result[0]["source"], "alice")
self.assertEqual(result[0]["relationship"], "knows")
self.assertEqual(result[0]["destination"], "bob")
def test_get_all_method(self):
"""Test the get_all method."""
# Mock the _get_all_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"user_id": self.user_id, "limit": 10}
self.memory_graph._get_all_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [
{"source": "alice", "relationship": "knows", "target": "bob"},
{"source": "bob", "relationship": "works_with", "target": "charlie"},
]
self.mock_graph.query.return_value = mock_query_result
# Call the get_all method
result = self.memory_graph.get_all(self.test_filters, top_k=10)
# Verify the method calls
self.memory_graph._get_all_cypher.assert_called_once_with(self.test_filters, 10)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result structure
self.assertEqual(len(result), 2)
self.assertEqual(result[0]["source"], "alice")
self.assertEqual(result[0]["relationship"], "knows")
self.assertEqual(result[0]["target"], "bob")
def test_delete_all_method(self):
"""Test the delete_all method."""
# Mock the _delete_all_cypher method
mock_cypher = "MATCH (n) DETACH DELETE n"
mock_params = {"user_id": self.user_id}
self.memory_graph._delete_all_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Call the delete_all method
self.memory_graph.delete_all(self.test_filters)
# Verify the method calls
self.memory_graph._delete_all_cypher.assert_called_once_with(self.test_filters)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
def test_search_source_node(self):
"""Test the _search_source_node method."""
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
# Mock the _search_source_node_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.9}
self.memory_graph._search_source_node_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"id(source_candidate)": 123, "cosine_similarity": 0.95}]
self.mock_graph.query.return_value = mock_query_result
# Call the _search_source_node method
result = self.memory_graph._search_source_node(mock_embedding, self.user_id, threshold=0.9)
# Verify the method calls
self.memory_graph._search_source_node_cypher.assert_called_once_with(mock_embedding, self.user_id, 0.9)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, mock_query_result)
def test_search_destination_node(self):
"""Test the _search_destination_node method."""
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
# Mock the _search_destination_node_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"destination_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.9}
self.memory_graph._search_destination_node_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"id(destination_candidate)": 456, "cosine_similarity": 0.92}]
self.mock_graph.query.return_value = mock_query_result
# Call the _search_destination_node method
result = self.memory_graph._search_destination_node(mock_embedding, self.user_id, threshold=0.9)
# Verify the method calls
self.memory_graph._search_destination_node_cypher.assert_called_once_with(mock_embedding, self.user_id, 0.9)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, mock_query_result)
def test_search_graph_db(self):
"""Test the _search_graph_db method."""
# Mock node list
node_list = ["alice", "bob"]
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
self.mock_embedding_model.embed.return_value = mock_embedding
# Mock the _search_graph_db_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"n_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.7, "limit": 10}
self.memory_graph._search_graph_db_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query results
mock_query_result1 = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
mock_query_result2 = [{"source": "bob", "relationship": "works_with", "destination": "charlie"}]
self.mock_graph.query.side_effect = [mock_query_result1, mock_query_result2]
# Call the _search_graph_db method
result = self.memory_graph._search_graph_db(node_list, self.test_filters, top_k=10)
# Verify the method calls
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
self.assertEqual(self.memory_graph._search_graph_db_cypher.call_count, 2)
self.assertEqual(self.mock_graph.query.call_count, 2)
# Check the result
expected_result = mock_query_result1 + mock_query_result2
self.assertEqual(result, expected_result)
def test_add_entities(self):
"""Test the _add_entities method."""
# Mock data
to_be_added = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
entity_type_map = {"alice": "person", "bob": "person"}
# Mock embeddings
mock_embedding = [0.1, 0.2, 0.3]
self.mock_embedding_model.embed.return_value = mock_embedding
# Mock search results
mock_source_search = [{"id(source_candidate)": 123, "cosine_similarity": 0.95}]
mock_dest_search = [{"id(destination_candidate)": 456, "cosine_similarity": 0.92}]
# Mock the search methods
self.memory_graph._search_source_node = MagicMock(return_value=mock_source_search)
self.memory_graph._search_destination_node = MagicMock(return_value=mock_dest_search)
# Mock the _add_entities_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_id": 123, "destination_id": 456}
self.memory_graph._add_entities_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"source": "alice", "relationship": "knows", "target": "bob"}]
self.mock_graph.query.return_value = mock_query_result
# Call the _add_entities method
result = self.memory_graph._add_entities(to_be_added, self.user_id, entity_type_map)
# Verify the method calls
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
self.memory_graph._add_entities_cypher.assert_called_once()
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, [mock_query_result])
def test_delete_entities(self):
"""Test the _delete_entities method."""
# Mock data
to_be_deleted = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
# Mock the _delete_entities_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_name": "alice", "dest_name": "bob", "user_id": self.user_id}
self.memory_graph._delete_entities_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"source": "alice", "relationship": "knows", "target": "bob"}]
self.mock_graph.query.return_value = mock_query_result
# Call the _delete_entities method
result = self.memory_graph._delete_entities(to_be_deleted, self.user_id)
# Verify the method calls
self.memory_graph._delete_entities_cypher.assert_called_once_with("alice", "bob", "knows", self.user_id)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, [mock_query_result])
if __name__ == "__main__":
unittest.main()
tests/memory/test_neptune_memory.py
-411
View File
@@ -1,411 +0,0 @@
import unittest
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
import pytest
from mem0.graphs.neptune.base import NeptuneBase
from mem0.graphs.neptune.neptunedb import MemoryGraph
class TestNeptuneMemory(unittest.TestCase):
"""Test suite for the Neptune Memory implementation."""
def setUp(self):
"""Set up test fixtures before each test method."""
# Create a mock config
self.config = MagicMock()
self.config.graph_store.config.endpoint = "neptune-db://test-graph"
self.config.graph_store.config.base_label = True
self.config.graph_store.threshold = 0.7
self.config.llm.provider = "openai_structured"
self.config.graph_store.llm = None
self.config.graph_store.custom_prompt = None
self.config.vector_store.provider = "qdrant"
self.config.vector_store.config = MagicMock()
# Create mock for NeptuneGraph
self.mock_graph = MagicMock()
# Create mocks for static methods
self.mock_embedding_model = MagicMock()
self.mock_llm = MagicMock()
self.mock_vector_store = MagicMock()
# Patch the necessary components
self.neptune_graph_patcher = patch("mem0.graphs.neptune.neptunedb.NeptuneGraph")
self.mock_neptune_graph = self.neptune_graph_patcher.start()
self.mock_neptune_graph.return_value = self.mock_graph
# Patch the static methods
self.create_embedding_model_patcher = patch.object(NeptuneBase, "_create_embedding_model")
self.mock_create_embedding_model = self.create_embedding_model_patcher.start()
self.mock_create_embedding_model.return_value = self.mock_embedding_model
self.create_llm_patcher = patch.object(NeptuneBase, "_create_llm")
self.mock_create_llm = self.create_llm_patcher.start()
self.mock_create_llm.return_value = self.mock_llm
self.create_vector_store_patcher = patch.object(NeptuneBase, "_create_vector_store")
self.mock_create_vector_store = self.create_vector_store_patcher.start()
self.mock_create_vector_store.return_value = self.mock_vector_store
# Create the MemoryGraph instance
self.memory_graph = MemoryGraph(self.config)
# Set up common test data
self.user_id = "test_user"
self.test_filters = {"user_id": self.user_id}
def tearDown(self):
"""Tear down test fixtures after each test method."""
self.neptune_graph_patcher.stop()
self.create_embedding_model_patcher.stop()
self.create_llm_patcher.stop()
self.create_vector_store_patcher.stop()
def test_initialization(self):
"""Test that the MemoryGraph is initialized correctly."""
self.assertEqual(self.memory_graph.graph, self.mock_graph)
self.assertEqual(self.memory_graph.embedding_model, self.mock_embedding_model)
self.assertEqual(self.memory_graph.llm, self.mock_llm)
self.assertEqual(self.memory_graph.vector_store, self.mock_vector_store)
self.assertEqual(self.memory_graph.llm_provider, "openai_structured")
self.assertEqual(self.memory_graph.node_label, ":`__Entity__`")
self.assertEqual(self.memory_graph.threshold, 0.7)
self.assertEqual(self.memory_graph.vector_store_limit, 5)
def test_collection_name_variants(self):
"""Test all collection_name configuration variants."""
# Test 1: graph_store.config.collection_name is set
config1 = MagicMock()
config1.graph_store.config.endpoint = "neptune-db://test-graph"
config1.graph_store.config.base_label = True
config1.graph_store.config.collection_name = "custom_collection"
config1.llm.provider = "openai"
config1.graph_store.llm = None
config1.vector_store.provider = "qdrant"
config1.vector_store.config = MagicMock()
MemoryGraph(config1)
self.assertEqual(config1.vector_store.config.collection_name, "custom_collection")
# Test 2: vector_store.config.collection_name exists, graph_store.config.collection_name is None
config2 = MagicMock()
config2.graph_store.config.endpoint = "neptune-db://test-graph"
config2.graph_store.config.base_label = True
config2.graph_store.config.collection_name = None
config2.llm.provider = "openai"
config2.graph_store.llm = None
config2.vector_store.provider = "qdrant"
config2.vector_store.config = MagicMock()
config2.vector_store.config.collection_name = "existing_collection"
MemoryGraph(config2)
self.assertEqual(config2.vector_store.config.collection_name, "existing_collection_neptune_vector_store")
# Test 3: Neither collection_name is set (default case)
config3 = MagicMock()
config3.graph_store.config.endpoint = "neptune-db://test-graph"
config3.graph_store.config.base_label = True
config3.graph_store.config.collection_name = None
config3.llm.provider = "openai"
config3.graph_store.llm = None
config3.vector_store.provider = "qdrant"
config3.vector_store.config = MagicMock()
config3.vector_store.config.collection_name = None
MemoryGraph(config3)
self.assertEqual(config3.vector_store.config.collection_name, "mem0_neptune_vector_store")
def test_init(self):
"""Test the class init functions"""
# Create a mock config with bad endpoint
config_no_endpoint = MagicMock()
config_no_endpoint.graph_store.config.endpoint = None
# Create the MemoryGraph instance
with pytest.raises(ValueError):
MemoryGraph(config_no_endpoint)
# Create a mock config with wrong endpoint type
config_wrong_endpoint = MagicMock()
config_wrong_endpoint.graph_store.config.endpoint = "neptune-graph://test-graph"
with pytest.raises(ValueError):
MemoryGraph(config_wrong_endpoint)
def test_add_method(self):
"""Test the add method with mocked components."""
# Mock the necessary methods that add() calls
self.memory_graph._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person", "bob": "person"})
self.memory_graph._establish_nodes_relations_from_data = MagicMock(
return_value=[{"source": "alice", "relationship": "knows", "destination": "bob"}]
)
self.memory_graph._search_graph_db = MagicMock(return_value=[])
self.memory_graph._get_delete_entities_from_search_output = MagicMock(return_value=[])
self.memory_graph._delete_entities = MagicMock(return_value=[])
self.memory_graph._add_entities = MagicMock(
return_value=[{"source": "alice", "relationship": "knows", "target": "bob"}]
)
# Call the add method
result = self.memory_graph.add("Alice knows Bob", self.test_filters)
# Verify the method calls
self.memory_graph._retrieve_nodes_from_data.assert_called_once_with("Alice knows Bob", self.test_filters)
self.memory_graph._establish_nodes_relations_from_data.assert_called_once()
self.memory_graph._search_graph_db.assert_called_once()
self.memory_graph._get_delete_entities_from_search_output.assert_called_once()
self.memory_graph._delete_entities.assert_called_once_with([], self.user_id)
self.memory_graph._add_entities.assert_called_once()
# Check the result structure
self.assertIn("deleted_entities", result)
self.assertIn("added_entities", result)
def test_search_method(self):
"""Test the search method with mocked components."""
# Mock the necessary methods that search() calls
self.memory_graph._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
# Mock search results
mock_search_results = [
{"source": "alice", "relationship": "knows", "destination": "bob"},
{"source": "alice", "relationship": "works_with", "destination": "charlie"},
]
self.memory_graph._search_graph_db = MagicMock(return_value=mock_search_results)
# Mock BM25Okapi
with patch("mem0.graphs.neptune.base.BM25Okapi") as mock_bm25:
mock_bm25_instance = MagicMock()
mock_bm25.return_value = mock_bm25_instance
# Mock get_top_n to return reranked results
reranked_results = [["alice", "knows", "bob"], ["alice", "works_with", "charlie"]]
mock_bm25_instance.get_top_n.return_value = reranked_results
# Call the search method
result = self.memory_graph.search("Find Alice", self.test_filters, top_k=5)
# Verify the method calls
self.memory_graph._retrieve_nodes_from_data.assert_called_once_with("Find Alice", self.test_filters)
self.memory_graph._search_graph_db.assert_called_once_with(node_list=["alice"], filters=self.test_filters)
# Check the result structure
self.assertEqual(len(result), 2)
self.assertEqual(result[0]["source"], "alice")
self.assertEqual(result[0]["relationship"], "knows")
self.assertEqual(result[0]["destination"], "bob")
def test_get_all_method(self):
"""Test the get_all method."""
# Mock the _get_all_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"user_id": self.user_id, "limit": 10}
self.memory_graph._get_all_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [
{"source": "alice", "relationship": "knows", "target": "bob"},
{"source": "bob", "relationship": "works_with", "target": "charlie"},
]
self.mock_graph.query.return_value = mock_query_result
# Call the get_all method
result = self.memory_graph.get_all(self.test_filters, top_k=10)
# Verify the method calls
self.memory_graph._get_all_cypher.assert_called_once_with(self.test_filters, 10)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result structure
self.assertEqual(len(result), 2)
self.assertEqual(result[0]["source"], "alice")
self.assertEqual(result[0]["relationship"], "knows")
self.assertEqual(result[0]["target"], "bob")
def test_delete_all_method(self):
"""Test the delete_all method."""
# Mock the _delete_all_cypher method
mock_cypher = "MATCH (n) DETACH DELETE n"
mock_params = {"user_id": self.user_id}
self.memory_graph._delete_all_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Call the delete_all method
self.memory_graph.delete_all(self.test_filters)
# Verify the method calls
self.memory_graph._delete_all_cypher.assert_called_once_with(self.test_filters)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
def test_search_source_node(self):
"""Test the _search_source_node method."""
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
# Mock the _search_source_node_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.9}
self.memory_graph._search_source_node_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"id(source_candidate)": 123, "cosine_similarity": 0.95}]
self.mock_graph.query.return_value = mock_query_result
# Call the _search_source_node method
result = self.memory_graph._search_source_node(mock_embedding, self.user_id, threshold=0.9)
# Verify the method calls
self.memory_graph._search_source_node_cypher.assert_called_once_with(mock_embedding, self.user_id, 0.9)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, mock_query_result)
def test_search_destination_node(self):
"""Test the _search_destination_node method."""
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
# Mock the _search_destination_node_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"destination_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.9}
self.memory_graph._search_destination_node_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"id(destination_candidate)": 456, "cosine_similarity": 0.92}]
self.mock_graph.query.return_value = mock_query_result
# Call the _search_destination_node method
result = self.memory_graph._search_destination_node(mock_embedding, self.user_id, threshold=0.9)
# Verify the method calls
self.memory_graph._search_destination_node_cypher.assert_called_once_with(mock_embedding, self.user_id, 0.9)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, mock_query_result)
def test_add_new_entities_payloads_use_utc_timestamps(self):
"""Test that Neptune vector-store payloads use UTC timestamps."""
self.memory_graph._add_new_entities_cypher(
source="alice",
source_embedding=[0.1, 0.2],
source_type="person",
destination="bob",
dest_embedding=[0.3, 0.4],
destination_type="person",
relationship="KNOWS",
user_id=self.user_id,
)
_, kwargs = self.mock_vector_store.insert.call_args
for payload in kwargs["payloads"]:
parsed = datetime.fromisoformat(payload["created_at"])
self.assertEqual(parsed.tzinfo, timezone.utc)
self.assertEqual(parsed.utcoffset().total_seconds(), 0)
def test_search_graph_db(self):
"""Test the _search_graph_db method."""
# Mock node list
node_list = ["alice", "bob"]
# Mock embedding
mock_embedding = [0.1, 0.2, 0.3]
self.mock_embedding_model.embed.return_value = mock_embedding
# Mock the _search_graph_db_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"n_embedding": mock_embedding, "user_id": self.user_id, "threshold": 0.7, "limit": 10}
self.memory_graph._search_graph_db_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query results
mock_query_result1 = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
mock_query_result2 = [{"source": "bob", "relationship": "works_with", "destination": "charlie"}]
self.mock_graph.query.side_effect = [mock_query_result1, mock_query_result2]
# Call the _search_graph_db method
result = self.memory_graph._search_graph_db(node_list, self.test_filters, top_k=10)
# Verify the method calls
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
self.assertEqual(self.memory_graph._search_graph_db_cypher.call_count, 2)
self.assertEqual(self.mock_graph.query.call_count, 2)
# Check the result
expected_result = mock_query_result1 + mock_query_result2
self.assertEqual(result, expected_result)
def test_add_entities(self):
"""Test the _add_entities method."""
# Mock data
to_be_added = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
entity_type_map = {"alice": "person", "bob": "person"}
# Mock embeddings
mock_embedding = [0.1, 0.2, 0.3]
self.mock_embedding_model.embed.return_value = mock_embedding
# Mock search results
mock_source_search = [{"id(source_candidate)": 123, "cosine_similarity": 0.95}]
mock_dest_search = [{"id(destination_candidate)": 456, "cosine_similarity": 0.92}]
# Mock the search methods
self.memory_graph._search_source_node = MagicMock(return_value=mock_source_search)
self.memory_graph._search_destination_node = MagicMock(return_value=mock_dest_search)
# Mock the _add_entities_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_id": 123, "destination_id": 456}
self.memory_graph._add_entities_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"source": "alice", "relationship": "knows", "target": "bob"}]
self.mock_graph.query.return_value = mock_query_result
# Call the _add_entities method
result = self.memory_graph._add_entities(to_be_added, self.user_id, entity_type_map)
# Verify the method calls
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
self.memory_graph._add_entities_cypher.assert_called_once()
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, [mock_query_result])
def test_delete_entities(self):
"""Test the _delete_entities method."""
# Mock data
to_be_deleted = [{"source": "alice", "relationship": "knows", "destination": "bob"}]
# Mock the _delete_entities_cypher method
mock_cypher = "MATCH (n) RETURN n"
mock_params = {"source_name": "alice", "dest_name": "bob", "user_id": self.user_id}
self.memory_graph._delete_entities_cypher = MagicMock(return_value=(mock_cypher, mock_params))
# Mock the graph.query result
mock_query_result = [{"source": "alice", "relationship": "knows", "target": "bob"}]
self.mock_graph.query.return_value = mock_query_result
# Call the _delete_entities method
result = self.memory_graph._delete_entities(to_be_deleted, self.user_id)
# Verify the method calls
self.memory_graph._delete_entities_cypher.assert_called_once_with("alice", "bob", "knows", self.user_id)
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
# Check the result
self.assertEqual(result, [mock_query_result])
if __name__ == "__main__":
unittest.main()
tests/memory/test_safe_deepcopy_config.py
+1 -1
View File
@@ -166,7 +166,7 @@ class TestRealWorldFieldCoverage:
# AWS
("aws_session_token", True),
# Azure MySQL
("use_azure_credential", False),
("use_azure_credential", True),
# General non-sensitive
("collection_name", False),
("embedding_model_dims", False),
tests/test_client.py
+149
View File
@@ -0,0 +1,149 @@
"""Tests for MemoryClient entity parameter rejection."""
from unittest.mock import MagicMock, patch
import pytest
@pytest.fixture
def mock_memory_client():
"""Create a mock MemoryClient for testing entity param rejection."""
with patch("mem0.client.main.httpx.Client") as mock_httpx:
# Create a mock client instance
mock_http_client = MagicMock()
mock_http_client.get.return_value = MagicMock(
json=lambda: {"org_id": "org1", "project_id": "proj1", "user_email": "test@test.com"},
raise_for_status=lambda: None
)
mock_httpx.return_value = mock_http_client
with patch("mem0.client.main.capture_client_event"):
from mem0.client.main import MemoryClient
client = MemoryClient(api_key="test-api-key")
yield client
class TestSearchEntityParamRejection:
"""Tests that top-level entity params are rejected in search()."""
def test_search_rejects_user_id_kwarg(self, mock_memory_client):
"""search() should reject user_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"user_id"):
mock_memory_client.search("test query", user_id="u1")
def test_search_rejects_agent_id_kwarg(self, mock_memory_client):
"""search() should reject agent_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"agent_id"):
mock_memory_client.search("test query", agent_id="a1")
def test_search_rejects_app_id_kwarg(self, mock_memory_client):
"""search() should reject app_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"app_id"):
mock_memory_client.search("test query", app_id="app1")
def test_search_rejects_run_id_kwarg(self, mock_memory_client):
"""search() should reject run_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"run_id"):
mock_memory_client.search("test query", run_id="r1")
def test_search_rejects_multiple_entity_params(self, mock_memory_client):
"""search() should reject multiple top-level entity params."""
with pytest.raises(ValueError, match=r"user_id|agent_id"):
mock_memory_client.search("test query", user_id="u1", agent_id="a1")
class TestGetAllEntityParamRejection:
"""Tests that top-level entity params are rejected in get_all()."""
def test_get_all_rejects_user_id_kwarg(self, mock_memory_client):
"""get_all() should reject user_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"user_id"):
mock_memory_client.get_all(user_id="u1")
def test_get_all_rejects_agent_id_kwarg(self, mock_memory_client):
"""get_all() should reject agent_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"agent_id"):
mock_memory_client.get_all(agent_id="a1")
def test_get_all_rejects_app_id_kwarg(self, mock_memory_client):
"""get_all() should reject app_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"app_id"):
mock_memory_client.get_all(app_id="app1")
def test_get_all_rejects_run_id_kwarg(self, mock_memory_client):
"""get_all() should reject run_id as top-level kwarg."""
with pytest.raises(ValueError, match=r"run_id"):
mock_memory_client.get_all(run_id="r1")
class TestFilterOperatorPassthrough:
"""Tests that AND/OR/NOT filter operators are passed through to the API."""
def test_search_passes_and_filters(self, mock_memory_client):
"""search() should pass AND filters to the API."""
mock_memory_client.client.post.return_value = MagicMock(
json=lambda: {"results": []},
raise_for_status=lambda: None
)
mock_memory_client.search(
"test query",
filters={"AND": [{"user_id": "u1"}, {"created_at": {"gte": "2024-01-01"}}]}
)
# Verify the POST was called with filters intact
call_args = mock_memory_client.client.post.call_args
payload = call_args.kwargs.get("json", call_args.args[1] if len(call_args.args) > 1 else {})
assert payload["filters"] == {"AND": [{"user_id": "u1"}, {"created_at": {"gte": "2024-01-01"}}]}
def test_search_passes_or_filters(self, mock_memory_client):
"""search() should pass OR filters to the API."""
mock_memory_client.client.post.return_value = MagicMock(
json=lambda: {"results": []},
raise_for_status=lambda: None
)
mock_memory_client.search(
"test query",
filters={"OR": [{"user_id": "u1"}, {"agent_id": "a1"}]}
)
call_args = mock_memory_client.client.post.call_args
payload = call_args.kwargs.get("json", call_args.args[1] if len(call_args.args) > 1 else {})
assert payload["filters"] == {"OR": [{"user_id": "u1"}, {"agent_id": "a1"}]}
def test_search_passes_not_filters(self, mock_memory_client):
"""search() should pass NOT filters to the API."""
mock_memory_client.client.post.return_value = MagicMock(
json=lambda: {"results": []},
raise_for_status=lambda: None
)
mock_memory_client.search(
"test query",
filters={"AND": [{"user_id": "u1"}, {"NOT": {"categories": {"in": ["spam"]}}}]}
)
call_args = mock_memory_client.client.post.call_args
payload = call_args.kwargs.get("json", call_args.args[1] if len(call_args.args) > 1 else {})
assert payload["filters"] == {"AND": [{"user_id": "u1"}, {"NOT": {"categories": {"in": ["spam"]}}}]}
def test_search_passes_complex_nested_filters(self, mock_memory_client):
"""search() should pass complex nested AND/OR/NOT filters to the API."""
mock_memory_client.client.post.return_value = MagicMock(
json=lambda: {"results": []},
raise_for_status=lambda: None
)
complex_filter = {
"AND": [
{"user_id": "u1"},
{"created_at": {"gte": "2024-01-01"}},
{"NOT": {"OR": [{"categories": {"in": ["spam"]}}, {"categories": {"in": ["test"]}}]}}
]
}
mock_memory_client.search("test query", filters=complex_filter)
call_args = mock_memory_client.client.post.call_args
payload = call_args.kwargs.get("json", call_args.args[1] if len(call_args.args) > 1 else {})
assert payload["filters"] == complex_filter
tests/test_graph_delete_docker.py
-821
View File
@@ -1,821 +0,0 @@
"""
End-to-end tests for graph cleanup on memory deletion against real
Neo4j, Memgraph, and Apache AGE instances running in Docker.
Requires:
docker run -d --name mem0-neo4j-test -p 7687:7687 -e NEO4J_AUTH=neo4j/testpassword neo4j:5.23
docker run -d --name mem0-memgraph-test -p 7688:7687 memgraph/memgraph:latest
docker run -d --name mem0-age-test -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=testpassword -e POSTGRES_DB=testdb apache/age:latest
Tests are skipped automatically if the databases or required Python
packages are not available.
"""
import hashlib
import sys
import warnings
from unittest.mock import MagicMock
import pytest
warnings.filterwarnings("ignore")
EMBEDDING_DIMS = 64
# ---------------------------------------------------------------------------
# Deterministic embedding helper (shared across backends)
# ---------------------------------------------------------------------------
def _make_deterministic_embedder():
cache = {}
counter = [0]
def embed(text, *args, **kwargs):
t = text.lower().strip()
if t not in cache:
vec = [0.0] * EMBEDDING_DIMS
idx = counter[0] % EMBEDDING_DIMS
vec[idx] = 1.0
h = hashlib.sha256(t.encode()).digest()
for i in range(EMBEDDING_DIMS):
vec[i] += float(h[i % len(h)]) / 25500.0
norm = sum(v * v for v in vec) ** 0.5
cache[t] = [v / norm for v in vec]
counter[0] += 1
return cache[t]
mock = MagicMock()
mock.embed.side_effect = embed
mock.config.embedding_dims = EMBEDDING_DIMS
return mock
def _make_mock_llm(entities, relations):
"""Create an LLM mock that returns specific entities and relations."""
mock = MagicMock()
def generate_response(messages, tools):
tool_names = []
for t in tools:
if isinstance(t, dict):
fn = t.get("function", t)
tool_names.append(fn.get("name", ""))
else:
tool_names.append(getattr(t, "name", str(t)))
if any("extract_entities" in n for n in tool_names):
return {
"tool_calls": [
{"name": "extract_entities", "arguments": {"entities": entities}}
]
}
elif any("establish" in n or "relation" in n for n in tool_names):
return {
"tool_calls": [
{"name": "establish_nodes_relations", "arguments": {"entities": relations}}
]
}
elif any("delete" in n for n in tool_names):
return {"tool_calls": []}
return {"tool_calls": []}
mock.generate_response.side_effect = generate_response
return mock
# ===========================================================================
# NEO4J
# ===========================================================================
def _port_open(host, port, timeout=1):
"""Quick TCP check — avoids slow driver-level timeouts."""
import socket
try:
with socket.create_connection((host, port), timeout=timeout):
return True
except OSError:
return False
def _neo4j_available():
if not _port_open("localhost", 7687):
return False
try:
from langchain_neo4j import Neo4jGraph
g = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
g.query("RETURN 1")
return True
except Exception:
return False
requires_neo4j = pytest.mark.skipif(not _neo4j_available(), reason="Neo4j not available")
@pytest.fixture
def neo4j_graph():
"""Create a Neo4j-backed MemoryGraph with mocked LLM/embedder."""
from langchain_neo4j import Neo4jGraph
from mem0.memory.graph_memory import MemoryGraph
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
mg.graph.query("MATCH (n) DETACH DELETE n")
mg.node_label = ":`__Entity__`"
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
mg.config.graph_store.config.base_label = True
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
@requires_neo4j
class TestNeo4jDeleteE2E:
def _node_count(self, mg):
return mg.graph.query("MATCH (n) RETURN count(n) AS cnt")[0]["cnt"]
def _valid_edge_count(self, mg):
return mg.graph.query(
"MATCH ()-[r]->() WHERE r.valid IS NULL OR r.valid = true RETURN count(r) AS cnt"
)[0]["cnt"]
def _invalid_edge_count(self, mg):
return mg.graph.query(
"MATCH ()-[r]->() WHERE r.valid = false RETURN count(r) AS cnt"
)[0]["cnt"]
def test_add_creates_graph_data(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
assert self._valid_edge_count(mg) == 1
def test_delete_soft_deletes_relationships(self, neo4j_graph):
"""Neo4j delete() should set r.valid=false (soft-delete), not hard-delete."""
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
assert self._invalid_edge_count(mg) == 0
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 0
assert self._invalid_edge_count(mg) == 1 # soft-deleted, not removed
assert self._node_count(mg) == 2 # nodes preserved
def test_delete_preserves_other_relationships(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
assert self._invalid_edge_count(mg) == 1
def test_delete_user_isolation(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert self._valid_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
def test_delete_all_hard_deletes(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert self._node_count(mg) == 0
def test_add_delete_add_cycle(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
# ===========================================================================
# MEMGRAPH
# ===========================================================================
def _memgraph_available():
if not _port_open("localhost", 7688):
return False
try:
from langchain_memgraph.graphs.memgraph import Memgraph
g = Memgraph("bolt://localhost:7688", "memgraph", "memgraph")
g.query("RETURN 1")
return True
except Exception:
return False
requires_memgraph = pytest.mark.skipif(
not _memgraph_available(), reason="Memgraph not available"
)
@pytest.fixture
def memgraph_graph():
"""Create a Memgraph-backed MemoryGraph with mocked LLM/embedder."""
from langchain_memgraph.graphs.memgraph import Memgraph
from mem0.memory.memgraph_memory import MemoryGraph
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Memgraph("bolt://localhost:7688", "memgraph", "memgraph")
mg.graph.query("MATCH (n) DETACH DELETE n")
try:
mg.graph.query("DROP VECTOR INDEX memzero;")
except Exception:
pass
mg.graph.query(
f"CREATE VECTOR INDEX memzero ON :Entity(embedding) "
f"WITH CONFIG {{'dimension': {EMBEDDING_DIMS}, 'capacity': 1000, 'metric': 'cos'}};"
)
try:
mg.graph.query("CREATE INDEX ON :Entity(user_id);")
except Exception:
pass
try:
mg.graph.query("CREATE INDEX ON :Entity;")
except Exception:
pass
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
mg.config.embedder.config = {"embedding_dims": EMBEDDING_DIMS}
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
@requires_memgraph
class TestMemgraphDeleteE2E:
def _node_count(self, mg):
return mg.graph.query("MATCH (n:Entity) RETURN count(n) AS cnt")[0]["cnt"]
def _edge_count(self, mg):
return mg.graph.query("MATCH ()-[r]->() RETURN count(r) AS cnt")[0]["cnt"]
def test_add_creates_graph_data(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
assert self._edge_count(mg) == 1
def test_delete_hard_deletes_relationships(self, memgraph_graph):
"""Memgraph delete() should hard-delete the relationship (DELETE r)."""
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 0
assert self._node_count(mg) == 2
def test_delete_preserves_other_relationships(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert self._edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
def test_delete_user_isolation(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert self._edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
def test_delete_all_hard_deletes(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert self._node_count(mg) == 0
assert self._edge_count(mg) == 0
def test_add_delete_add_cycle(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
# ===========================================================================
# APACHE AGE
# ===========================================================================
def _age_available():
if not _port_open("localhost", 5432):
return False
try:
import age
ag = age.connect(
host="localhost",
port=5432,
dbname="testdb",
user="postgres",
password="testpassword",
)
with ag.connection.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS age;")
cur.execute("SET search_path = ag_catalog, '$user', public;")
ag.connection.commit()
ag.close()
return True
except Exception:
return False
requires_age = pytest.mark.skipif(not _age_available(), reason="Apache AGE not available")
@pytest.fixture
def age_graph():
"""Create an Apache AGE-backed MemoryGraph with mocked LLM/embedder."""
import age
from mem0.memory.apache_age_memory import MemoryGraph
graph_name = "mem0_test_delete"
ag = age.connect(
graph=graph_name,
host="localhost",
port=5432,
dbname="testdb",
user="postgres",
password="testpassword",
)
with ag.connection.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS age;")
cur.execute("SET search_path = ag_catalog, '$user', public;")
ag.connection.commit()
age.setUpAge(ag.connection, graph_name)
ag.connection.commit()
mg = MemoryGraph.__new__(MemoryGraph)
mg.ag = ag
mg.graph_name = graph_name
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
try:
ag.execCypher("MATCH (n) DETACH DELETE n")
ag.commit()
except Exception:
ag.rollback()
yield mg
try:
ag.execCypher("MATCH (n) DETACH DELETE n")
ag.commit()
except Exception:
ag.rollback()
ag.close()
def _age_node_count(mg):
cursor = mg.ag.execCypher("MATCH (n) RETURN count(n)", cols=["cnt"])
rows = cursor.fetchall()
return rows[0][0] if rows else 0
def _age_edge_count(mg):
cursor = mg.ag.execCypher("MATCH ()-[r]->() RETURN count(r)", cols=["cnt"])
rows = cursor.fetchall()
return rows[0][0] if rows else 0
@requires_age
class TestApacheAgeDeleteE2E:
def test_add_creates_graph_data(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_node_count(mg) == 2
assert _age_edge_count(mg) == 1
def test_delete_hard_deletes_relationships(self, age_graph):
"""Apache AGE delete() should hard-delete the relationship."""
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 0
assert _age_node_count(mg) == 2
def test_delete_preserves_other_relationships(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert _age_edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
def test_delete_user_isolation(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert _age_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
def test_delete_all_hard_deletes(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert _age_node_count(mg) == 0
assert _age_edge_count(mg) == 0
def test_add_delete_add_cycle(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
# ===========================================================================
# NEPTUNE (tested via Neo4j OpenCypher — same query language)
# ===========================================================================
def _neptune_test_available():
"""Neptune uses OpenCypher — we test NeptuneBase.delete() against Neo4j."""
if not _port_open("localhost", 7687):
return False
try:
# Mock langchain_aws so NeptuneBase can be imported without AWS deps
sys.modules.setdefault("langchain_aws", MagicMock())
sys.modules.setdefault("botocore", MagicMock())
sys.modules.setdefault("botocore.config", MagicMock())
from mem0.graphs.neptune.base import NeptuneBase # noqa: F401
from langchain_neo4j import Neo4jGraph
g = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
g.query("RETURN 1")
return True
except Exception:
return False
requires_neptune_test = pytest.mark.skipif(
not _neptune_test_available(),
reason="Neo4j not available (used as OpenCypher backend for Neptune tests)",
)
def _make_concrete_neptune_subclass():
"""Create a concrete NeptuneBase subclass for testing, backed by Neo4j."""
# Ensure mocks are in place for import
sys.modules.setdefault("langchain_aws", MagicMock())
sys.modules.setdefault("botocore", MagicMock())
sys.modules.setdefault("botocore.config", MagicMock())
from mem0.graphs.neptune.base import NeptuneBase
class TestableNeptune(NeptuneBase):
def __init__(self):
pass
def _delete_entities_cypher(self, source, destination, relationship, user_id):
cypher = f"""
MATCH (n:`__Entity__` {{name: $source_name, user_id: $user_id}})
-[r:{relationship}]->
(m:`__Entity__` {{name: $dest_name, user_id: $user_id}})
DELETE r
RETURN n.name AS source, m.name AS target, type(r) AS relationship
"""
return cypher, {"source_name": source, "dest_name": destination, "user_id": user_id}
def _delete_all_cypher(self, filters):
return (
"MATCH (n:`__Entity__` {user_id: $user_id}) DETACH DELETE n",
{"user_id": filters["user_id"]},
)
# Stubs for abstract methods not used in delete path
def _add_entities_by_source_cypher(self, *a, **kw): pass
def _add_entities_by_destination_cypher(self, *a, **kw): pass
def _add_relationship_entities_cypher(self, *a, **kw): pass
def _add_new_entities_cypher(self, *a, **kw): pass
def _search_source_node_cypher(self, *a, **kw): pass
def _search_destination_node_cypher(self, *a, **kw): pass
def _get_all_cypher(self, *a, **kw): pass
def _search_graph_db_cypher(self, *a, **kw): pass
return TestableNeptune
@pytest.fixture
def neptune_graph():
"""NeptuneBase subclass backed by a real Neo4j container."""
from langchain_neo4j import Neo4jGraph
cls = _make_concrete_neptune_subclass()
mg = cls()
mg.graph = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
mg.graph.query("MATCH (n) DETACH DELETE n")
mg.node_label = ":`__Entity__`"
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
def _neptune_node_count(mg):
return mg.graph.query("MATCH (n) RETURN count(n) AS cnt")[0]["cnt"]
def _neptune_edge_count(mg):
return mg.graph.query("MATCH ()-[r]->() RETURN count(r) AS cnt")[0]["cnt"]
def _neptune_create_entities(mg, user_id):
"""Create test entities directly via Cypher."""
mg.graph.query(f"""
CREATE (a:`__Entity__` {{name: 'alice', user_id: '{user_id}'}})
CREATE (b:`__Entity__` {{name: 'bob', user_id: '{user_id}'}})
CREATE (a)-[:likes]->(b)
""")
@requires_neptune_test
class TestNeptuneDeleteE2E:
"""Test NeptuneBase.delete() using Neo4j as the OpenCypher backend.
Neptune uses standard OpenCypher, the same query language as Neo4j.
This validates that:
- NeptuneBase.delete() correctly calls _delete_entities(to_be_deleted, user_id) with a string
- The generated Cypher from _delete_entities_cypher runs correctly
- User isolation works
"""
def test_delete_removes_relationship(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 0
assert _neptune_node_count(mg) == 2
def test_delete_user_isolation(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
_neptune_create_entities(mg, "u2")
assert _neptune_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 1
def test_delete_passes_user_id_string_not_dict(self, neptune_graph):
"""Verify NeptuneBase.delete() passes filters['user_id'] (string) to _delete_entities."""
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
original = mg._delete_entities
call_args = []
def spy(to_be_deleted, user_id):
call_args.append(("to_be_deleted", to_be_deleted, "user_id", user_id))
return original(to_be_deleted, user_id)
mg._delete_entities = spy
_neptune_create_entities(mg, "u1")
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert len(call_args) == 1
assert call_args[0][3] == "u1"
assert isinstance(call_args[0][3], str)
def test_delete_all(self, neptune_graph):
mg = neptune_graph
_neptune_create_entities(mg, "u1")
assert _neptune_node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert _neptune_node_count(mg) == 0
def test_add_delete_add_cycle(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 0
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
tests/test_main.py
+25 -154
View File
@@ -22,20 +22,13 @@ def memory_instance():
patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
patch("mem0.utils.factory.LlmFactory") as mock_llm,
patch("mem0.memory.telemetry.capture_event"),
patch("mem0.memory.graph_memory.MemoryGraph"),
patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
):
mock_embedder.create.return_value = Mock()
mock_vector_store.create.return_value = Mock()
mock_vector_store.create.return_value.search.return_value = []
mock_llm.create.return_value = Mock()
# Create a mock instance that won't try to access config attributes
mock_graph_instance = Mock()
mock_graph_store.create.return_value = mock_graph_instance
config = MemoryConfig(version="v1.1")
config.graph_store.config = {"some_config": "value"}
return Memory(config)
@@ -46,55 +39,31 @@ def memory_custom_instance():
patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
patch("mem0.utils.factory.LlmFactory") as mock_llm,
patch("mem0.memory.telemetry.capture_event"),
patch("mem0.memory.graph_memory.MemoryGraph"),
patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
):
mock_embedder.create.return_value = Mock()
mock_vector_store.create.return_value = Mock()
mock_vector_store.create.return_value.search.return_value = []
mock_llm.create.return_value = Mock()
# Create a mock instance that won't try to access config attributes
mock_graph_instance = Mock()
mock_graph_store.create.return_value = mock_graph_instance
config = MemoryConfig(
version="v1.1",
custom_instructions="custom prompt extracting memory in json format",
custom_update_memory_prompt="custom prompt determining memory update",
)
config.graph_store.config = {"some_config": "value"}
return Memory(config)
@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
def test_add(memory_instance, version, enable_graph):
memory_instance.config.version = version
if not enable_graph:
memory_instance.graph = None
def test_add(memory_instance):
memory_instance._add_to_vector_store = Mock(return_value=[{"memory": "Test memory", "event": "ADD"}])
memory_instance._add_to_graph = Mock(return_value=[])
result = memory_instance.add(messages=[{"role": "user", "content": "Test message"}], user_id="test_user")
if enable_graph:
assert "results" in result
assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
assert "relations" in result
assert result["relations"] == []
else:
assert "results" in result
assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
assert "results" in result
assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
memory_instance._add_to_vector_store.assert_called_once_with(
[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}, {"user_id": "test_user"}, True
)
# Remove the conditional assertion for _add_to_graph
memory_instance._add_to_graph.assert_called_once_with(
[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}
)
def test_get(memory_instance):
mock_memory = Mock(
@@ -121,50 +90,32 @@ def test_get(memory_instance):
assert result["metadata"] == {"extra_field": "extra_value"}
@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
def test_search(memory_instance, version, enable_graph):
memory_instance.config.version = version
if not enable_graph:
memory_instance.graph = None
def test_search(memory_instance):
mock_memories = [
Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"}, score=0.9),
Mock(id="2", payload={"data": "Memory 2", "user_id": "test_user"}, score=0.8),
]
memory_instance.vector_store.search = Mock(return_value=mock_memories)
memory_instance.vector_store.keyword_search = Mock(return_value=None) # No BM25
memory_instance.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3])
if memory_instance.graph:
memory_instance.graph.search = Mock(return_value=[{"relation": "test_relation"}])
result = memory_instance.search("test query", user_id="test_user")
with patch("mem0.memory.main.lemmatize_for_bm25", return_value="test query"), \
patch("mem0.memory.main.extract_entities", return_value=[]):
result = memory_instance.search("test query", filters={"user_id": "test_user"})
if version == "v1.1":
assert "results" in result
assert len(result["results"]) == 2
assert result["results"][0]["id"] == "1"
assert result["results"][0]["memory"] == "Memory 1"
assert result["results"][0]["user_id"] == "test_user"
assert result["results"][0]["score"] == 0.9
if enable_graph:
assert "relations" in result
assert result["relations"] == [{"relation": "test_relation"}]
else:
assert "relations" not in result
else:
assert isinstance(result, dict)
assert "results" in result
assert len(result["results"]) == 2
assert result["results"][0]["id"] == "1"
assert result["results"][0]["memory"] == "Memory 1"
assert result["results"][0]["user_id"] == "test_user"
assert result["results"][0]["score"] == 0.9
assert "results" in result
assert len(result["results"]) == 2
assert result["results"][0]["id"] == "1"
assert result["results"][0]["memory"] == "Memory 1"
assert result["results"][0]["user_id"] == "test_user"
# Score is now combined score (semantic only since no BM25/entity), still 0.9
assert result["results"][0]["score"] == pytest.approx(0.9)
assert "score_breakdown" in result["results"][0]
# Hybrid pipeline over-fetches: max(100*4, 60) = 400
memory_instance.vector_store.search.assert_called_once_with(
query="test query", vectors=[0.1, 0.2, 0.3], top_k=100, filters={"user_id": "test_user"}
query="test query", vectors=[0.1, 0.2, 0.3], top_k=400, filters={"user_id": "test_user"}
)
memory_instance.embedding_model.embed.assert_called_once_with("test query", "search")
if enable_graph:
memory_instance.graph.search.assert_called_once_with("test query", {"user_id": "test_user"}, 100)
def test_update(memory_instance):
@@ -222,17 +173,11 @@ def test_delete(memory_instance):
assert result["message"] == "Memory deleted successfully!"
@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
def test_delete_all(memory_instance, version, enable_graph):
memory_instance.config.version = version
if not enable_graph:
memory_instance.graph = None
def test_delete_all(memory_instance):
mock_memories = [Mock(id="1"), Mock(id="2")]
memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
memory_instance.vector_store.reset = Mock()
memory_instance._delete_memory = Mock()
if memory_instance.graph:
memory_instance.graph.delete_all = Mock()
result = memory_instance.delete_all(user_id="test_user")
@@ -240,98 +185,24 @@ def test_delete_all(memory_instance, version, enable_graph):
# Ensure the collection is NOT dropped — only matched memories should be removed
memory_instance.vector_store.reset.assert_not_called()
if enable_graph:
memory_instance.graph.delete_all.assert_called_once_with({"user_id": "test_user"})
assert result["message"] == "Memories deleted successfully!"
@pytest.mark.parametrize(
"version, enable_graph, expected_result",
[
("v1.0", False, {"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}]}),
("v1.1", False, {"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}]}),
(
"v1.1",
True,
{
"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}],
"relations": [{"source": "entity1", "relationship": "rel", "target": "entity2"}],
},
),
],
)
def test_get_all(memory_instance, version, enable_graph, expected_result):
memory_instance.config.version = version
if not enable_graph:
memory_instance.graph = None
def test_get_all(memory_instance):
mock_memories = [Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"})]
memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
if memory_instance.graph:
memory_instance.graph.get_all = Mock(
return_value=[{"source": "entity1", "relationship": "rel", "target": "entity2"}]
)
result = memory_instance.get_all(user_id="test_user")
result = memory_instance.get_all(filters={"user_id": "test_user"})
assert isinstance(result, dict)
assert "results" in result
assert len(result["results"]) == len(expected_result["results"])
for expected_item, result_item in zip(expected_result["results"], result["results"]):
assert all(key in result_item for key in expected_item)
assert result_item["id"] == expected_item["id"]
assert result_item["memory"] == expected_item["memory"]
assert result_item["user_id"] == expected_item["user_id"]
if enable_graph:
assert "relations" in result
assert result["relations"] == expected_result["relations"]
else:
assert "relations" not in result
assert len(result["results"]) == 1
assert result["results"][0]["id"] == "1"
assert result["results"][0]["memory"] == "Memory 1"
assert result["results"][0]["user_id"] == "test_user"
memory_instance.vector_store.list.assert_called_once_with(filters={"user_id": "test_user"}, top_k=100)
if enable_graph:
memory_instance.graph.get_all.assert_called_once_with({"user_id": "test_user"}, 100)
def test_custom_prompts(memory_custom_instance):
messages = [{"role": "user", "content": "Test message"}]
from mem0.embeddings.mock import MockEmbeddings
memory_custom_instance.llm.generate_response = Mock()
memory_custom_instance.llm.generate_response.return_value = '{"facts": ["fact1", "fact2"]}'
memory_custom_instance.embedding_model = MockEmbeddings()
with patch("mem0.memory.main.parse_messages", return_value="Test message") as mock_parse_messages:
with patch(
"mem0.memory.main.get_update_memory_messages", return_value="custom update memory prompt"
) as mock_get_update_memory_messages:
memory_custom_instance.add(messages=messages, user_id="test_user")
## custom prompt
##
mock_parse_messages.assert_called_once_with(messages)
memory_custom_instance.llm.generate_response.assert_any_call(
messages=[
{"role": "system", "content": memory_custom_instance.config.custom_instructions},
{"role": "user", "content": f"Input:\n{mock_parse_messages.return_value}"},
],
response_format={"type": "json_object"},
)
## custom update memory prompt
##
mock_get_update_memory_messages.assert_called_once_with(
[], ["fact1", "fact2"], memory_custom_instance.config.custom_update_memory_prompt
)
memory_custom_instance.llm.generate_response.assert_any_call(
messages=[{"role": "user", "content": mock_get_update_memory_messages.return_value}],
response_format={"type": "json_object"},
)
def test_no_telemetry_vector_store_when_disabled():
"""VectorStoreFactory should only be called once (for user data) when telemetry is disabled."""
tests/test_memory.py
+54 -64
View File
@@ -73,7 +73,7 @@ def test_list_memories(memory_client):
data2 = "Name is John Doe. I like to code in Python."
memory_client.add([{"role": "user", "content": data1}], user_id="test_user")
memory_client.add([{"role": "user", "content": data2}], user_id="test_user")
memories = memory_client.get_all(user_id="test_user")
memories = memory_client.get_all(filters={"user_id": "test_user"})
assert data1 in memories
assert data2 in memories
@@ -137,11 +137,10 @@ def test_search_handles_incomplete_payloads(mock_sqlite, mock_llm_factory, mock_
result = memory._search_vector_store("test", {"user_id": "test"}, 10)
assert len(result) == 2
memories_by_id = {mem["id"]: mem for mem in result}
assert memories_by_id["mem_1"]["memory"] == ""
assert memories_by_id["mem_2"]["memory"] == "content"
# v3 search pipeline skips entries where payload has no "data" key
assert len(result) == 1
assert result[0]["id"] == "mem_2"
assert result[0]["memory"] == "content"
@patch('mem0.utils.factory.EmbedderFactory.create')
@@ -497,14 +496,16 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
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"
# V3 single-call extraction: LLM returns extracted memories directly
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.generate_response.return_value = json.dumps(
{"memory": [{"text": "The user enjoys Python"}]}
)
mock_llm_factory.return_value = mock_llm
# embed_batch is used in Phase 3 for all extracted memories
embedder.embed_batch.return_value = [[0.4, 0.5, 0.6]]
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
@@ -512,11 +513,10 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
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
# V3 pipeline: embed called once for search query (Phase 1),
# embed_batch called once for extracted memories (Phase 3)
assert embedder.embed.call_count == 1
assert embedder.embed_batch.call_count == 1
mock_vector_store.insert.assert_called_once()
@@ -526,15 +526,15 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
@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.
Regression test for issue #3723 (infer=True path): V3 is ADD-only, so this test verifies
that the single-call extraction pipeline embeds via embed_batch, not individual embed calls.
"""
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 that will be returned from search
existing_memory = MockVectorMemory(
memory_id="existing-mem-id",
payload={
@@ -549,17 +549,20 @@ def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm
mock_vector_store.get.return_value = existing_memory
mock_vector_store.insert.return_value = None
mock_vector_store.update.return_value = None
mock_vector_store.keyword_search.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"
# V3 single-call extraction: LLM returns extracted memories directly
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.generate_response.return_value = json.dumps(
{"memory": [{"text": "The user loves Python"}]}
)
mock_llm_factory.return_value = mock_llm
# embed_batch is used in Phase 3 for all extracted memories
embedder.embed_batch.return_value = [[0.4, 0.5, 0.6]]
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
@@ -567,12 +570,11 @@ def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm
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()
# V3 pipeline: embed called once for search query (Phase 1),
# embed_batch called once for extracted memories (Phase 3)
assert embedder.embed.call_count == 1
assert embedder.embed_batch.call_count == 1
mock_vector_store.insert.assert_called_once()
@patch('mem0.utils.factory.EmbedderFactory.create')
@@ -761,30 +763,6 @@ class TestProcessMetadataFiltersMerge:
# --- Issue #3040: reset() should clean up graph database ---
@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_reset_calls_graph_reset_when_graph_enabled(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""Test that reset() calls graph.reset() when graph is enabled (issue #3040)."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
config = MemoryConfig()
memory = Memory(config)
# Simulate graph being enabled by setting graph instance
mock_graph = MagicMock()
memory.graph = mock_graph
memory.reset()
mock_graph.reset.assert_called_once()
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@@ -809,26 +787,38 @@ def test_reset_skips_graph_when_graph_disabled(mock_sqlite, mock_llm_factory, mo
assert memory.graph is None
# ─── Entity Param Rejection Tests ─────────────────────────────────────────────
@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_reset_continues_if_graph_reset_fails(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""Test that reset() doesn't crash if graph.reset() raises an exception (issue #3040)."""
def test_search_rejects_user_id_kwarg(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""search() should reject user_id as top-level kwarg."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_vector_factory.return_value = MagicMock()
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
config = MemoryConfig()
memory = Memory(config)
mock_graph = MagicMock()
mock_graph.reset.side_effect = Exception("Neo4j connection failed")
memory.graph = mock_graph
with pytest.raises(ValueError, match=r"user_id.*filters"):
memory.search("test query", user_id="u1")
# Should NOT raise — graph failure is logged but reset continues
memory.reset()
mock_graph.reset.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_get_all_rejects_user_id_kwarg(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""get_all() should reject user_id as top-level kwarg."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_factory.return_value = MagicMock()
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
config = MemoryConfig()
memory = Memory(config)
with pytest.raises(ValueError, match=r"user_id.*filters"):
memory.get_all(user_id="u1")
tests/test_telemetry.py
-49
View File
@@ -1,4 +1,3 @@
import asyncio
import threading
from unittest.mock import MagicMock, patch
@@ -291,11 +290,6 @@ class TestMemoryLifecycle:
m.db = MagicMock()
return m
def test_close_calls_db_close(self):
m = self._make_mock_memory()
m.close()
m.db.close.assert_called_once()
def test_double_close_is_safe(self):
"""close() should be safe to call twice (SQLiteManager.close sets connection=None)."""
m = self._make_mock_memory()
@@ -319,21 +313,6 @@ class TestMemoryLifecycle:
# db attribute not set at all
m.close() # should not raise due to hasattr guard
def test_context_manager(self):
"""Memory should support `with` statement and close on exit."""
m = self._make_mock_memory()
with m:
pass
m.db.close.assert_called_once()
def test_context_manager_closes_on_exception(self):
"""Memory should close even if the with-block raises."""
m = self._make_mock_memory()
with pytest.raises(ValueError):
with m:
raise ValueError("boom")
m.db.close.assert_called_once()
class TestAsyncMemoryLifecycle:
"""Verify AsyncMemory.close() and async context manager support."""
@@ -346,11 +325,6 @@ class TestAsyncMemoryLifecycle:
m.db = MagicMock()
return m
def test_close_calls_db_close(self):
m = self._make_mock_async_memory()
m.close()
m.db.close.assert_called_once()
def test_close_when_db_is_none(self):
m = self._make_mock_async_memory()
m.db = None
@@ -363,29 +337,6 @@ class TestAsyncMemoryLifecycle:
m = AsyncMemory.__new__(AsyncMemory)
m.close() # should not raise
def test_async_context_manager(self):
"""AsyncMemory should support `async with` and close on exit."""
m = self._make_mock_async_memory()
async def run():
async with m:
pass
asyncio.run(run())
m.db.close.assert_called_once()
def test_async_context_manager_closes_on_exception(self):
"""AsyncMemory should close even if the async with-block raises."""
m = self._make_mock_async_memory()
async def run():
async with m:
raise ValueError("boom")
with pytest.raises(ValueError):
asyncio.run(run())
m.db.close.assert_called_once()
class TestTelemetryEnvVar:
"""Verify the MEM0_TELEMETRY env var parsing logic."""
tests/utils/__init__.py
View File
tests/utils/test_entity_extraction.py
+101
View File
@@ -0,0 +1,101 @@
import pytest
@pytest.fixture(autouse=True)
def _ensure_spacy():
"""Skip tests if spaCy model is not available."""
try:
import spacy
spacy.load("en_core_web_sm")
except Exception:
pytest.skip("spaCy en_core_web_sm model not available")
class TestExtractEntities:
def test_proper_nouns(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("John Smith works at Google on machine learning projects")
entity_texts = [e[1] for e in entities]
# Should extract proper nouns
found_proper = any("John" in t or "Google" in t for t in entity_texts)
assert found_proper, f"Expected proper nouns, got {entities}"
def test_quoted_text(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities('She is reading "The Great Gatsby" this week')
entity_texts = [e[1] for e in entities]
assert any("Great Gatsby" in t for t in entity_texts), f"Expected quoted text, got {entities}"
def test_compound_nouns(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("The machine learning engineer built a neural network")
entity_texts = [e[1].lower() for e in entities]
has_compound = any("machine" in t and "learning" in t for t in entity_texts) or \
any("neural" in t and "network" in t for t in entity_texts)
assert has_compound, f"Expected compound nouns, got {entities}"
def test_empty_string(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("")
assert entities == []
def test_no_entities(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("I like things and stuff")
# Generic words should be filtered out
entity_texts = [e[1].lower() for e in entities]
assert "things" not in entity_texts
assert "stuff" not in entity_texts
def test_deduplication(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("Google is great. I love working at Google.")
google_count = sum(1 for _, t in entities if "Google" in t)
assert google_count <= 1, f"Expected dedup, got {entities}"
def test_returns_tuples(self):
from mem0.utils.entity_extraction import extract_entities
entities = extract_entities("John Smith lives in New York City")
for entity in entities:
assert isinstance(entity, tuple)
assert len(entity) == 2
assert entity[0] in ("PROPER", "QUOTED", "COMPOUND", "NOUN")
assert isinstance(entity[1], str)
class TestExtractEntitiesBatch:
def test_batch_processing(self):
from mem0.utils.entity_extraction import extract_entities_batch
texts = [
"John works at Google",
"Mary lives in Paris",
"The cat sat on the mat",
]
results = extract_entities_batch(texts)
assert len(results) == 3
assert isinstance(results[0], list)
assert isinstance(results[1], list)
assert isinstance(results[2], list)
def test_empty_input(self):
from mem0.utils.entity_extraction import extract_entities_batch
assert extract_entities_batch([]) == []
def test_consistency_with_single(self):
from mem0.utils.entity_extraction import extract_entities, extract_entities_batch
text = "John Smith works at Google headquarters"
single = extract_entities(text)
batch = extract_entities_batch([text])
assert len(batch) == 1
# Both should extract the same entities
assert set(t for _, t in single) == set(t for _, t in batch[0])
tests/utils/test_lemmatization.py
+67
View File
@@ -0,0 +1,67 @@
import pytest
@pytest.fixture(autouse=True)
def _ensure_spacy():
"""Skip tests if spaCy model is not available."""
try:
import spacy
spacy.load("en_core_web_sm")
except Exception:
pytest.skip("spaCy en_core_web_sm model not available")
class TestLemmatizeForBm25:
def test_basic_lemmatization(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("The cats are running quickly")
assert "cat" in result
assert "run" in result or "running" in result
# Stop words and punctuation should be removed
assert "the" not in result.split()
def test_verb_forms_normalized(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("she attended multiple meetings yesterday")
assert "attend" in result or "attended" in result
assert "meeting" in result # -ing form preserved alongside lemma
# "multiple" is kept (not a spaCy stop word)
def test_ing_preservation(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("attending the morning meeting")
tokens = result.split()
# Should have both the lemma and the -ing form
assert "attending" in tokens or "attend" in tokens
def test_empty_string(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("")
assert result == ""
def test_punctuation_removed(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("Hello, world! How are you?")
assert "," not in result
assert "!" not in result
assert "?" not in result
def test_lowercased(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("PYTHON Programming LANGUAGE")
for token in result.split():
assert token == token.lower()
def test_stop_words_removed(self):
from mem0.utils.lemmatization import lemmatize_for_bm25
result = lemmatize_for_bm25("this is a very simple test of the system")
tokens = result.split()
for stop in ["this", "is", "a", "very", "of", "the"]:
assert stop not in tokens
tests/utils/test_scoring.py
+142
View File
@@ -0,0 +1,142 @@
import pytest
from mem0.utils.scoring import (
get_bm25_params,
normalize_bm25,
score_and_rank,
ENTITY_BOOST_WEIGHT,
)
class TestGetBm25Params:
def test_short_query(self):
midpoint, steepness = get_bm25_params("hello world", lemmatized="hello world")
assert midpoint == 5.0
assert steepness == 0.7
def test_medium_query(self):
midpoint, steepness = get_bm25_params("x", lemmatized="one two three four five")
assert midpoint == 7.0
assert steepness == 0.6
def test_long_query(self):
words = " ".join(f"word{i}" for i in range(20))
midpoint, steepness = get_bm25_params("x", lemmatized=words)
assert midpoint == 12.0
assert steepness == 0.5
def test_empty_lemmatized(self):
midpoint, steepness = get_bm25_params("test", lemmatized="")
# Empty string -> 1 term -> short query params
assert midpoint == 5.0
class TestNormalizeBm25:
def test_at_midpoint(self):
score = normalize_bm25(5.0, 5.0, 0.7)
assert abs(score - 0.5) < 0.01 # Should be ~0.5 at midpoint
def test_high_score(self):
score = normalize_bm25(20.0, 5.0, 0.7)
assert score > 0.99 # Well above midpoint
def test_low_score(self):
score = normalize_bm25(0.0, 5.0, 0.7)
assert score < 0.05 # Well below midpoint
def test_range(self):
for raw in [0, 1, 5, 10, 20, 50]:
score = normalize_bm25(float(raw), 5.0, 0.7)
assert 0.0 <= score <= 1.0
class TestScoreAndRank:
def test_semantic_only(self):
results = [
{"id": "a", "score": 0.9, "payload": {"data": "mem a"}},
{"id": "b", "score": 0.5, "payload": {"data": "mem b"}},
]
scored = score_and_rank(results, {}, {}, threshold=0.1, top_k=10)
assert len(scored) == 2
# With no BM25/entity, max_possible=1.0, so scores stay the same
assert scored[0]["score"] == pytest.approx(0.9)
assert scored[1]["score"] == pytest.approx(0.5)
def test_semantic_plus_bm25(self):
results = [
{"id": "a", "score": 0.8, "payload": {"data": "mem a"}},
{"id": "b", "score": 0.6, "payload": {"data": "mem b"}},
]
bm25 = {"a": 0.3, "b": 0.9}
scored = score_and_rank(results, bm25, {}, threshold=0.1, top_k=10)
# max_possible = 2.0 (semantic + bm25)
# a: (0.8 + 0.3) / 2.0 = 0.55
# b: (0.6 + 0.9) / 2.0 = 0.75
assert scored[0]["id"] == "b" # b should rank higher due to BM25
assert scored[0]["score"] == pytest.approx(0.75)
assert scored[1]["id"] == "a"
assert scored[1]["score"] == pytest.approx(0.55)
def test_all_three_signals(self):
results = [{"id": "a", "score": 0.8, "payload": {"data": "mem a"}}]
bm25 = {"a": 0.6}
entity = {"a": 0.3}
scored = score_and_rank(results, bm25, entity, threshold=0.1, top_k=10)
# max_possible = 2.5
expected = (0.8 + 0.6 + 0.3) / 2.5
assert scored[0]["score"] == pytest.approx(expected)
def test_threshold_gates_on_semantic(self):
results = [
{"id": "a", "score": 0.05, "payload": {"data": "mem a"}}, # Below threshold
{"id": "b", "score": 0.5, "payload": {"data": "mem b"}},
]
bm25 = {"a": 0.99} # High BM25 shouldn't save it
scored = score_and_rank(results, bm25, {}, threshold=0.1, top_k=10)
assert len(scored) == 1
assert scored[0]["id"] == "b"
def test_top_k_limit(self):
results = [{"id": str(i), "score": 0.5, "payload": {}} for i in range(20)]
scored = score_and_rank(results, {}, {}, threshold=0.1, top_k=5)
assert len(scored) == 5
def test_score_breakdown_present(self):
results = [{"id": "a", "score": 0.8, "payload": {"data": "x"}}]
bm25 = {"a": 0.4}
entity = {"a": 0.2}
scored = score_and_rank(results, bm25, entity, threshold=0.1, top_k=10)
breakdown = scored[0]["score_breakdown"]
assert breakdown["semantic"] == 0.8
assert breakdown["bm25"] == 0.4
assert breakdown["entity_boost"] == 0.2
def test_adaptive_divisor_semantic_only(self):
results = [{"id": "a", "score": 0.8, "payload": {}}]
scored = score_and_rank(results, {}, {}, threshold=0.1, top_k=10)
# max_possible = 1.0 (no bm25, no entity)
assert scored[0]["score"] == pytest.approx(0.8)
def test_adaptive_divisor_semantic_plus_entity(self):
results = [{"id": "a", "score": 0.8, "payload": {}}]
entity = {"a": 0.3}
scored = score_and_rank(results, {}, entity, threshold=0.1, top_k=10)
# max_possible = 1.5 (semantic + entity)
expected = (0.8 + 0.3) / 1.5
assert scored[0]["score"] == pytest.approx(expected)
def test_empty_results(self):
scored = score_and_rank([], {}, {}, threshold=0.1, top_k=10)
assert scored == []
def test_score_clamped_to_1(self):
results = [{"id": "a", "score": 1.0, "payload": {}}]
bm25 = {"a": 1.0}
entity = {"a": 0.5}
scored = score_and_rank(results, bm25, entity, threshold=0.1, top_k=10)
assert scored[0]["score"] <= 1.0
class TestEntityBoostWeight:
def test_weight_value(self):
assert ENTITY_BOOST_WEIGHT == 0.5
tests/vector_stores/test_mongodb.py
+6 -3
View File
@@ -42,9 +42,12 @@ def test_initalize_create_col(mongo_vector_fixture):
mock_collection.insert_one.assert_called_once_with({"_id": 0, "placeholder": True})
mock_collection.delete_one.assert_called_once_with({"_id": 0})
assert mongo_vector.index_name == "test_collection_vector_index"
mock_collection.list_search_indexes.assert_called_once_with(name="test_collection_vector_index")
mock_collection.create_search_index.assert_called_once()
args, _ = mock_collection.create_search_index.call_args
mock_collection.list_search_indexes.assert_any_call(name="test_collection_vector_index")
mock_collection.list_search_indexes.assert_any_call(name="test_collection_text_search_index")
assert mock_collection.list_search_indexes.call_count == 2
# Two indexes created: vector search + text search
assert mock_collection.create_search_index.call_count == 2
args, _ = mock_collection.create_search_index.call_args_list[0]
search_index_model = args[0].document
assert search_index_model == {
"name": "test_collection_vector_index",
tests/vector_stores/test_qdrant.py
+12 -3
View File
@@ -4,7 +4,7 @@ import unittest
import uuid
from unittest.mock import MagicMock, patch
from qdrant_client import QdrantClient
from qdrant_client import QdrantClient, models
from qdrant_client.models import (
DatetimeRange,
Distance,
@@ -18,6 +18,7 @@ from qdrant_client.models import (
PointStruct,
PointVectors,
Range,
SparseVectorParams,
VectorParams,
)
@@ -59,8 +60,14 @@ class TestQdrant(unittest.TestCase):
expected_config = VectorParams(size=128, distance=Distance.COSINE, on_disk=True)
expected_sparse_config = {
"bm25": SparseVectorParams(modifier=models.Modifier.IDF),
}
self.client_mock.create_collection.assert_called_with(
collection_name="test_collection", vectors_config=expected_config
collection_name="test_collection",
vectors_config=expected_config,
sparse_vectors_config=expected_sparse_config,
)
def test_insert(self):
@@ -231,7 +238,9 @@ class TestQdrant(unittest.TestCase):
self.client_mock.upsert.assert_called_once()
point = self.client_mock.upsert.call_args[1]["points"][0]
self.assertEqual(point.id, vector_id)
self.assertEqual(point.vector, updated_vector)
# v3 uses named vectors: dense vector stored under "" key
self.assertIn("", point.vector)
self.assertEqual(point.vector[""], updated_vector)
self.assertEqual(point.payload, updated_payload)
def test_update_with_none_vector_uses_set_payload(self):