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"""
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 = {

View File

@@ -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)