fix: delete graph test files that import removed graph modules
These test files import mem0.memory.graph_memory, kuzu_memory, memgraph_memory, apache_age_memory, and mem0.graphs.neptune which were deleted in the graph store removal commit. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Soumil Rathi
File diff suppressed because it is too large
Load Diff
@@ -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 == []
|
||||
@@ -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
|
||||
@@ -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'])
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user