From caeae60dda4bdd9ff431e73fcf1a8fc0292a5ff0 Mon Sep 17 00:00:00 2001
From: Tomaz Bratanic <bratanic.tomaz@gmail.com>
Date: Sat, 10 May 2025 23:51:34 +0200
Subject: [PATCH] Add weights to Neo4j model (#2657)

---
 mem0/memory/graph_memory.py | 42 ++++++++++++++++++++++++++++++-------
 mem0/memory/utils.py        |  2 +-
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/mem0/memory/graph_memory.py b/mem0/memory/graph_memory.py
index 08df0724..8167bddf 100644
--- a/mem0/memory/graph_memory.py
+++ b/mem0/memory/graph_memory.py
@@ -338,15 +338,23 @@ class MemoryGraph:
                 cypher = f"""
                     MATCH (source)
                     WHERE elementId(source) = $source_id
+                    SET source.mentions = coalesce(source.mentions, 0) + 1
+                    WITH source
                     MERGE (destination:{destination_type} {{name: $destination_name, user_id: $user_id}})
                     ON CREATE SET
-                        destination.created = timestamp()
+                        destination.created = timestamp(),
+                        destination.mentions = 1
+                    ON MATCH SET
+                        destination.mentions = coalesce(destination.mentions, 0) + 1
                     WITH source, destination
                     CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
                     WITH source, destination
                     MERGE (source)-[r:{relationship}]->(destination)
                     ON CREATE SET 
-                        r.created = timestamp()
+                        r.created = timestamp(),
+                        r.mentions = 1
+                    ON MATCH SET
+                        r.mentions = coalesce(r.mentions, 0) + 1
                     RETURN source.name AS source, type(r) AS relationship, destination.name AS target
                     """
 
@@ -360,15 +368,23 @@ class MemoryGraph:
                 cypher = f"""
                     MATCH (destination)
                     WHERE elementId(destination) = $destination_id
+                    SET destination.mentions = coalesce(destination.mentions, 0) + 1
+                    WITH destination
                     MERGE (source:{source_type} {{name: $source_name, user_id: $user_id}})
                     ON CREATE SET
-                        source.created = timestamp()
+                        source.created = timestamp(),
+                        source.mentions = 1
+                    ON MATCH SET
+                        source.mentions = coalesce(source.mentions, 0) + 1
                     WITH source, destination
                     CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
                     WITH source, destination
                     MERGE (source)-[r:{relationship}]->(destination)
                     ON CREATE SET 
-                        r.created = timestamp()
+                        r.created = timestamp(),
+                        r.mentions = 1
+                    ON MATCH SET
+                        r.mentions = coalesce(r.mentions, 0) + 1
                     RETURN source.name AS source, type(r) AS relationship, destination.name AS target
                     """
 
@@ -382,12 +398,17 @@ class MemoryGraph:
                 cypher = f"""
                     MATCH (source)
                     WHERE elementId(source) = $source_id
+                    SET source.mentions = coalesce(source.mentions, 0) + 1
+                    WITH source
                     MATCH (destination)
                     WHERE elementId(destination) = $destination_id
+                    SET destination.mentions = coalesce(destination.mentions) + 1
                     MERGE (source)-[r:{relationship}]->(destination)
                     ON CREATE SET 
                         r.created_at = timestamp(),
-                        r.updated_at = timestamp()
+                        r.updated_at = timestamp(),
+                        r.mentions = 1
+                    ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
                     
                     
                     RETURN source.name AS source, type(r) AS relationship, destination.name AS target
@@ -400,17 +421,22 @@ class MemoryGraph:
             else:
                 cypher = f"""
                     MERGE (n:{source_type} {{name: $source_name, user_id: $user_id}})
-                    ON CREATE SET n.created = timestamp()
+                    ON CREATE SET n.created = timestamp(),
+                                  n.mentions = 1
+                    ON MATCH SET n.mentions = coalesce(n.mentions, 0) + 1
                     WITH n
                     CALL db.create.setNodeVectorProperty(n, 'embedding', $source_embedding)
                     WITH n
                     MERGE (m:{destination_type} {{name: $dest_name, user_id: $user_id}})
-                    ON CREATE SET m.created = timestamp()
+                    ON CREATE SET m.created = timestamp(),
+                                  m.mentions = 1
+                    ON MATCH SET m.mentions = coalesce(m.mentions, 0) + 1
                     WITH n, m
                     CALL db.create.setNodeVectorProperty(m, 'embedding', $source_embedding)
                     WITH n, m
                     MERGE (n)-[rel:{relationship}]->(m)
-                    ON CREATE SET rel.created = timestamp()
+                    ON CREATE SET rel.created = timestamp(), rel.mentions = 1
+                    ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
                     RETURN n.name AS source, type(rel) AS relationship, m.name AS target
                     """
                 params = {
diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py
index ad14bf46..ddfa2cb6 100644
--- a/mem0/memory/utils.py
+++ b/mem0/memory/utils.py
@@ -25,7 +25,7 @@ def format_entities(entities):
 
     formatted_lines = []
     for entity in entities:
-        simplified = f"{entity['source']} -- {entity['relatationship']} -- {entity['destination']}"
+        simplified = f"{entity['source']} -- {entity['relationship']} -- {entity['destination']}"
         formatted_lines.append(simplified)
 
     return "\n".join(formatted_lines)