Add weights to Neo4j model (#2657)

This commit is contained in:
Tomaz Bratanic
2025-05-10 23:51:34 +02:00
committed by GitHub
parent d7b8497b24
commit caeae60dda
2 changed files with 35 additions and 9 deletions

View File

@@ -338,15 +338,23 @@ class MemoryGraph:
cypher = f""" cypher = f"""
MATCH (source) MATCH (source)
WHERE elementId(source) = $source_id 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}}) MERGE (destination:{destination_type} {{name: $destination_name, user_id: $user_id}})
ON CREATE SET 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 WITH source, destination
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding) CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
WITH source, destination WITH source, destination
MERGE (source)-[r:{relationship}]->(destination) MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET 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 RETURN source.name AS source, type(r) AS relationship, destination.name AS target
""" """
@@ -360,15 +368,23 @@ class MemoryGraph:
cypher = f""" cypher = f"""
MATCH (destination) MATCH (destination)
WHERE elementId(destination) = $destination_id 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}}) MERGE (source:{source_type} {{name: $source_name, user_id: $user_id}})
ON CREATE SET 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 WITH source, destination
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding) CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
WITH source, destination WITH source, destination
MERGE (source)-[r:{relationship}]->(destination) MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET 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 RETURN source.name AS source, type(r) AS relationship, destination.name AS target
""" """
@@ -382,12 +398,17 @@ class MemoryGraph:
cypher = f""" cypher = f"""
MATCH (source) MATCH (source)
WHERE elementId(source) = $source_id WHERE elementId(source) = $source_id
SET source.mentions = coalesce(source.mentions, 0) + 1
WITH source
MATCH (destination) MATCH (destination)
WHERE elementId(destination) = $destination_id WHERE elementId(destination) = $destination_id
SET destination.mentions = coalesce(destination.mentions) + 1
MERGE (source)-[r:{relationship}]->(destination) MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET ON CREATE SET
r.created_at = timestamp(), 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 RETURN source.name AS source, type(r) AS relationship, destination.name AS target
@@ -400,17 +421,22 @@ class MemoryGraph:
else: else:
cypher = f""" cypher = f"""
MERGE (n:{source_type} {{name: $source_name, user_id: $user_id}}) 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 WITH n
CALL db.create.setNodeVectorProperty(n, 'embedding', $source_embedding) CALL db.create.setNodeVectorProperty(n, 'embedding', $source_embedding)
WITH n WITH n
MERGE (m:{destination_type} {{name: $dest_name, user_id: $user_id}}) 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 WITH n, m
CALL db.create.setNodeVectorProperty(m, 'embedding', $source_embedding) CALL db.create.setNodeVectorProperty(m, 'embedding', $source_embedding)
WITH n, m WITH n, m
MERGE (n)-[rel:{relationship}]->(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 RETURN n.name AS source, type(rel) AS relationship, m.name AS target
""" """
params = { params = {

View File

@@ -25,7 +25,7 @@ def format_entities(entities):
formatted_lines = [] formatted_lines = []
for entity in entities: for entity in entities:
simplified = f"{entity['source']} -- {entity['relatationship']} -- {entity['destination']}" simplified = f"{entity['source']} -- {entity['relationship']} -- {entity['destination']}"
formatted_lines.append(simplified) formatted_lines.append(simplified)
return "\n".join(formatted_lines) return "\n".join(formatted_lines)