diff --git a/mem0/memory/kuzu_memory.py b/mem0/memory/kuzu_memory.py index acb5bb8dc..828e2e174 100644 --- a/mem0/memory/kuzu_memory.py +++ b/mem0/memory/kuzu_memory.py @@ -522,12 +522,12 @@ class MemoryGraph: WITH destination MERGE (source {source_label} {{{merge_props_str}}}) ON CREATE SET - source.created = current_timestamp(), - source.mentions = 1 - source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]') + source.created = current_timestamp(), + source.mentions = 1, + source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]') ON MATCH SET - source.mentions = coalesce(source.mentions, 0) + 1 - source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]') + source.mentions = coalesce(source.mentions, 0) + 1, + source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]') WITH source, destination MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination) ON CREATE SET diff --git a/tests/memory/test_kuzu.py b/tests/memory/test_kuzu.py index f2e429527..334912add 100644 --- a/tests/memory/test_kuzu.py +++ b/tests/memory/test_kuzu.py @@ -11,6 +11,7 @@ class TestKuzu: "alice": np.random.uniform(0.0, 0.9, 384).tolist(), "bob": np.random.uniform(0.0, 0.9, 384).tolist(), "charlie": np.random.uniform(0.0, 0.9, 384).tolist(), + "dave": np.random.uniform(0.0, 0.9, 384).tolist(), } @pytest.fixture @@ -78,6 +79,7 @@ class TestKuzu: assert kuzu_memory.llm == mock_llm assert kuzu_memory.threshold == 0.7 + @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): @@ -109,12 +111,21 @@ class TestKuzu: 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" + "charlie_knows_alice", + "dave_admires_alice" ]) results = kuzu_memory._search_graph_db(["bob"], filters, threshold=0.8) @@ -125,15 +136,15 @@ class TestKuzu: result = kuzu_memory._delete_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) == 3 + 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) == 3 - assert get_edge_count(kuzu_memory) == 0 + 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