gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor Neo4JStorage to use entity_id for node identification, use entity_type for node label

Browse Source
This commit is contained in:
yangdx
2025-03-11 09:59:42 +08:00
parent 7ed7cb4ecf
commit aefd596990
1 changed files with 95 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files

287
lightrag/kg/neo4j_impl.py
View File

@@ -176,23 +176,6 @@ class Neo4JStorage(BaseGraphStorage):
# Noe4J handles persistence automatically # Noe4J handles persistence automatically
pass pass
def _ensure_label(self, label: str) -> str:
"""Ensure a label is valid
Args:
label: The label to validate
Returns:
str: The cleaned label
Raises:
ValueError: If label is empty after cleaning
"""
clean_label = label.strip('"')
if not clean_label:
raise ValueError("Neo4j: Label cannot be empty")
return clean_label
async def has_node(self, node_id: str) -> bool: async def has_node(self, node_id: str) -> bool:
""" """
Check if a node with the given label exists in the database Check if a node with the given label exists in the database
@@ -207,19 +190,18 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists" query = "MATCH (n:base {entity_id: $entity_id}) RETURN count(n) > 0 AS node_exists"
result = await session.run(query) result = await session.run(query, entity_id = node_id)
single_result = await result.single() single_result = await result.single()
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
return single_result["node_exists"] return single_result["node_exists"]
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error checking node existence for {entity_name_label}: {str(e)}" f"Error checking node existence for {node_id}: {str(e)}"
) )
await result.consume() # Ensure results are consumed even on error await result.consume() # Ensure results are consumed even on error
raise raise
@@ -239,24 +221,21 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If either node_id is invalid ValueError: If either node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label_source = self._ensure_label(source_node_id)
entity_name_label_target = self._ensure_label(target_node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = ( query = (
f"MATCH (a:`{entity_name_label_source}`)-[r]-(b:`{entity_name_label_target}`) " "MATCH (a:base {entity_id: $source_entity_id})-[r]-(b:base {entity_id: $target_entity_id}) "
"RETURN COUNT(r) > 0 AS edgeExists" "RETURN COUNT(r) > 0 AS edgeExists"
) )
result = await session.run(query) result = await session.run(query, source_entity_id = source_node_id, target_entity_id = target_node_id)
single_result = await result.single() single_result = await result.single()
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
return single_result["edgeExists"] return single_result["edgeExists"]
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error checking edge existence between {entity_name_label_source} and {entity_name_label_target}: {str(e)}" f"Error checking edge existence between {source_node_id} and {target_node_id}: {str(e)}"
) )
await result.consume() # Ensure results are consumed even on error await result.consume() # Ensure results are consumed even on error
raise raise
@@ -275,13 +254,12 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f"MATCH (n:`{entity_name_label}` {{entity_id: $entity_id}}) RETURN n" query = "MATCH (n:base {entity_id: $entity_id}) RETURN n"
result = await session.run(query, entity_id=entity_name_label) result = await session.run(query, entity_id=node_id)
try: try:
records = await result.fetch( records = await result.fetch(
2 2
@@ -289,20 +267,21 @@ class Neo4JStorage(BaseGraphStorage):
if len(records) > 1: if len(records) > 1:
logger.warning( logger.warning(
f"Multiple nodes found with label '{entity_name_label}'. Using first node." f"Multiple nodes found with label '{node_id}'. Using first node."
) )
if records: if records:
node = records[0]["n"] node = records[0]["n"]
node_dict = dict(node) node_dict = dict(node)
logger.debug( # Remove base label from labels list if it exists
f"{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}" if "labels" in node_dict:
) node_dict["labels"] = [label for label in node_dict["labels"] if label != "base"]
logger.debug(f"Neo4j query node {query} return: {node_dict}")
return node_dict return node_dict
return None return None
finally: finally:
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
except Exception as e: except Exception as e:
logger.error(f"Error getting node for {entity_name_label}: {str(e)}") logger.error(f"Error getting node for {node_id}: {str(e)}")
raise raise
async def node_degree(self, node_id: str) -> int: async def node_degree(self, node_id: str) -> int:
@@ -320,42 +299,33 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f""" query = """
MATCH (n:`{entity_name_label}`) MATCH (n:base {entity_id: $entity_id})
OPTIONAL MATCH (n)-[r]-() OPTIONAL MATCH (n)-[r]-()
RETURN n, COUNT(r) AS degree RETURN COUNT(r) AS degree
""" """
result = await session.run(query) result = await session.run(query, entity_id = node_id)
try: try:
records = await result.fetch(100) record = await result.single()
if not records: if not record:
logger.warning( logger.warning(
f"No node found with label '{entity_name_label}'" f"No node found with label '{node_id}'"
) )
return 0 return 0
if len(records) > 1: degree = record["degree"]
logger.warning( logger.debug("Neo4j query node degree for {node_id} return: {degree}")
f"Multiple nodes ({len(records)}) found with label '{entity_name_label}', using first node's degree"
)
degree = records[0]["degree"]
logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{degree}"
)
return degree return degree
finally: finally:
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error getting node degree for {entity_name_label}: {str(e)}" f"Error getting node degree for {node_id}: {str(e)}"
) )
raise raise
@@ -369,11 +339,8 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
int: Sum of the degrees of both nodes int: Sum of the degrees of both nodes
""" """
entity_name_label_source = self._ensure_label(src_id) src_degree = await self.node_degree(src_id)
entity_name_label_target = self._ensure_label(tgt_id) trg_degree = await self.node_degree(tgt_id)
src_degree = await self.node_degree(entity_name_label_source)
trg_degree = await self.node_degree(entity_name_label_target)
# Convert None to 0 for addition # Convert None to 0 for addition
src_degree = 0 if src_degree is None else src_degree src_degree = 0 if src_degree is None else src_degree
@@ -399,24 +366,20 @@ class Neo4JStorage(BaseGraphStorage):
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
try: try:
entity_name_label_source = self._ensure_label(source_node_id)
entity_name_label_target = self._ensure_label(target_node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = f""" query = """
MATCH (start:`{entity_name_label_source}`)-[r]-(end:`{entity_name_label_target}`) MATCH (start:base {entity_id: $source_entity_id})-[r]-(end:base {entity_id: $target_entity_id})
RETURN properties(r) as edge_properties RETURN properties(r) as edge_properties
""" """
result = await session.run(query, source_entity_id=source_node_id, target_entity_id=target_node_id)
result = await session.run(query)
try: try:
records = await result.fetch(2) records = await result.fetch(2)
if len(records) > 1: if len(records) > 1:
logger.warning( logger.warning(
f"Multiple edges found between '{entity_name_label_source}' and '{entity_name_label_target}'. Using first edge." f"Multiple edges found between '{source_node_id}' and '{target_node_id}'. Using first edge."
) )
if records: if records:
try: try:
@@ -433,7 +396,7 @@ class Neo4JStorage(BaseGraphStorage):
if key not in edge_result: if key not in edge_result:
edge_result[key] = default_value edge_result[key] = default_value
logger.warning( logger.warning(
f"Edge between {entity_name_label_source} and {entity_name_label_target} " f"Edge between {source_node_id} and {target_node_id} "
f"missing {key}, using default: {default_value}" f"missing {key}, using default: {default_value}"
) )
@@ -443,8 +406,8 @@ class Neo4JStorage(BaseGraphStorage):
return edge_result return edge_result
except (KeyError, TypeError, ValueError) as e: except (KeyError, TypeError, ValueError) as e:
logger.error( logger.error(
f"Error processing edge properties between {entity_name_label_source} " f"Error processing edge properties between {source_node_id} "
f"and {entity_name_label_target}: {str(e)}" f"and {target_node_id}: {str(e)}"
) )
# Return default edge properties on error # Return default edge properties on error
return { return {
@@ -455,7 +418,7 @@ class Neo4JStorage(BaseGraphStorage):
} }
logger.debug( logger.debug(
f"{inspect.currentframe().f_code.co_name}: No edge found between {entity_name_label_source} and {entity_name_label_target}" f"{inspect.currentframe().f_code.co_name}: No edge found between {source_node_id} and {target_node_id}"
) )
# Return default edge properties when no edge found # Return default edge properties when no edge found
return { return {
@@ -488,30 +451,30 @@ class Neo4JStorage(BaseGraphStorage):
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
try: try:
node_label = self._ensure_label(source_node_id)
query = f"""MATCH (n:`{node_label}`)
OPTIONAL MATCH (n)-[r]-(connected)
RETURN n, r, connected"""
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
results = await session.run(query) query = """MATCH (n:base {entity_id: $entity_id})
edges = [] OPTIONAL MATCH (n)-[r]-(connected:base)
WHERE connected.entity_id IS NOT NULL
RETURN n, r, connected"""
results = await session.run(query, entity_id=source_node_id)
edges = []
async for record in results: async for record in results:
source_node = record["n"] source_node = record["n"]
connected_node = record["connected"] connected_node = record["connected"]
# Skip if either node is None
if not source_node or not connected_node:
continue
source_label = ( source_label = (
list(source_node.labels)[0] if source_node.labels else None source_node.get("entity_id") if source_node.get("entity_id") else None
) )
target_label = ( target_label = (
list(connected_node.labels)[0] connected_node.get("entity_id") if connected_node.get("entity_id") else None
if connected_node and connected_node.labels
else None
) )
if source_label and target_label: if source_label and target_label:
@@ -520,7 +483,7 @@ class Neo4JStorage(BaseGraphStorage):
await results.consume() # Ensure results are consumed await results.consume() # Ensure results are consumed
return edges return edges
except Exception as e: except Exception as e:
logger.error(f"Error getting edges for node {node_label}: {str(e)}") logger.error(f"Error getting edges for node {source_node_id}: {str(e)}")
await results.consume() # Ensure results are consumed even on error await results.consume() # Ensure results are consumed even on error
raise raise
except Exception as e: except Exception as e:
@@ -547,8 +510,9 @@ class Neo4JStorage(BaseGraphStorage):
node_id: The unique identifier for the node (used as label) node_id: The unique identifier for the node (used as label)
node_data: Dictionary of node properties node_data: Dictionary of node properties
""" """
label = self._ensure_label(node_id)
properties = node_data properties = node_data
entity_type = properties["entity_type"]
entity_id = properties["entity_id"]
if "entity_id" not in properties: if "entity_id" not in properties:
raise ValueError("Neo4j: node properties must contain an 'entity_id' field") raise ValueError("Neo4j: node properties must contain an 'entity_id' field")
@@ -556,13 +520,14 @@ class Neo4JStorage(BaseGraphStorage):
async with self._driver.session(database=self._DATABASE) as session: async with self._driver.session(database=self._DATABASE) as session:
async def execute_upsert(tx: AsyncManagedTransaction): async def execute_upsert(tx: AsyncManagedTransaction):
query = f""" query = """
MERGE (n:`{label}` {{entity_id: $properties.entity_id}}) MERGE (n:base {entity_id: $properties.entity_id})
SET n += $properties SET n += $properties
""" SET n:`%s`
""" % entity_type
result = await tx.run(query, properties=properties) result = await tx.run(query, properties=properties)
logger.debug( logger.debug(
f"Upserted node with label '{label}' and properties: {properties}" f"Upserted node with entity_id '{entity_id}' and properties: {properties}"
) )
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
@@ -583,51 +548,6 @@ class Neo4JStorage(BaseGraphStorage):
) )
), ),
) )
async def _get_unique_node_entity_id(self, node_label: str) -> str:
"""
Get the entity_id of a node with the given label, ensuring the node is unique.
Args:
node_label (str): Label of the node to check
Returns:
str: The entity_id of the unique node
Raises:
ValueError: If no node with the given label exists or if multiple nodes have the same label
"""
async with self._driver.session(
database=self._DATABASE, default_access_mode="READ"
) as session:
query = f"""
MATCH (n:`{node_label}`)
RETURN n, count(n) as node_count
"""
result = await session.run(query)
try:
records = await result.fetch(
2
) # We only need to know if there are 0, 1, or >1 nodes
if not records or records[0]["node_count"] == 0:
raise ValueError(
f"Neo4j: node with label '{node_label}' does not exist"
)
if records[0]["node_count"] > 1:
raise ValueError(
f"Neo4j: multiple nodes found with label '{node_label}', cannot determine unique node"
)
node = records[0]["n"]
if "entity_id" not in node:
raise ValueError(
f"Neo4j: node with label '{node_label}' does not have an entity_id property"
)
return node["entity_id"]
finally:
await result.consume() # Ensure result is fully consumed
@retry( @retry(
stop=stop_after_attempt(3), stop=stop_after_attempt(3),
@@ -657,38 +577,30 @@ class Neo4JStorage(BaseGraphStorage):
Raises: Raises:
ValueError: If either source or target node does not exist or is not unique ValueError: If either source or target node does not exist or is not unique
""" """
source_label = self._ensure_label(source_node_id)
target_label = self._ensure_label(target_node_id)
edge_properties = edge_data
# Get entity_ids for source and target nodes, ensuring they are unique
source_entity_id = await self._get_unique_node_entity_id(source_label)
target_entity_id = await self._get_unique_node_entity_id(target_label)
try: try:
edge_properties = edge_data
async with self._driver.session(database=self._DATABASE) as session: async with self._driver.session(database=self._DATABASE) as session:
async def execute_upsert(tx: AsyncManagedTransaction): async def execute_upsert(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (source:`{source_label}` {{entity_id: $source_entity_id}}) MATCH (source:base {entity_id: $source_entity_id})
WITH source WITH source
MATCH (target:`{target_label}` {{entity_id: $target_entity_id}}) MATCH (target:base {entity_id: $target_entity_id})
MERGE (source)-[r:DIRECTED]-(target) MERGE (source)-[r:DIRECTED]-(target)
SET r += $properties SET r += $properties
RETURN r, source, target RETURN r, source, target
""" """
result = await tx.run( result = await tx.run(
query, query,
source_entity_id=source_entity_id, source_entity_id=source_node_id,
target_entity_id=target_entity_id, target_entity_id=target_node_id,
properties=edge_properties, properties=edge_properties,
) )
try: try:
records = await result.fetch(100) records = await result.fetch(2)
if records: if records:
logger.debug( logger.debug(
f"Upserted edge from '{source_label}' (entity_id: {source_entity_id}) " f"Upserted edge from '{source_node_id}' to '{target_node_id}'"
f"to '{target_label}' (entity_id: {target_entity_id}) "
f"with properties: {edge_properties}" f"with properties: {edge_properties}"
) )
finally: finally:
@@ -726,7 +638,6 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
KnowledgeGraph: Complete connected subgraph for specified node KnowledgeGraph: Complete connected subgraph for specified node
""" """
label = node_label.strip('"')
result = KnowledgeGraph() result = KnowledgeGraph()
seen_nodes = set() seen_nodes = set()
seen_edges = set() seen_edges = set()
@@ -735,7 +646,7 @@ class Neo4JStorage(BaseGraphStorage):
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
if label == "*": if node_label == "*":
main_query = """ main_query = """
MATCH (n) MATCH (n)
OPTIONAL MATCH (n)-[r]-() OPTIONAL MATCH (n)-[r]-()
@@ -760,12 +671,11 @@ class Neo4JStorage(BaseGraphStorage):
# Main query uses partial matching # Main query uses partial matching
main_query = """ main_query = """
MATCH (start) MATCH (start)
WHERE any(label IN labels(start) WHERE WHERE
CASE CASE
WHEN $inclusive THEN label CONTAINS $label WHEN $inclusive THEN start.entity_id CONTAINS $entity_id
ELSE label = $label ELSE start.entity_id = $entity_id
END END
)
WITH start WITH start
CALL apoc.path.subgraphAll(start, { CALL apoc.path.subgraphAll(start, {
relationshipFilter: '', relationshipFilter: '',
@@ -799,7 +709,7 @@ class Neo4JStorage(BaseGraphStorage):
main_query, main_query,
{ {
"max_nodes": MAX_GRAPH_NODES, "max_nodes": MAX_GRAPH_NODES,
"label": label, "entity_id": node_label,
"inclusive": inclusive, "inclusive": inclusive,
"max_depth": max_depth, "max_depth": max_depth,
"min_degree": min_degree, "min_degree": min_degree,
@@ -818,7 +728,7 @@ class Neo4JStorage(BaseGraphStorage):
result.nodes.append( result.nodes.append(
KnowledgeGraphNode( KnowledgeGraphNode(
id=f"{node_id}", id=f"{node_id}",
labels=list(node.labels), labels=[label for label in node.labels if label != "base"],
properties=dict(node), properties=dict(node),
) )
) )
@@ -849,7 +759,7 @@ class Neo4JStorage(BaseGraphStorage):
except neo4jExceptions.ClientError as e: except neo4jExceptions.ClientError as e:
logger.warning(f"APOC plugin error: {str(e)}") logger.warning(f"APOC plugin error: {str(e)}")
if label != "*": if node_label != "*":
logger.warning( logger.warning(
"Neo4j: falling back to basic Cypher recursive search..." "Neo4j: falling back to basic Cypher recursive search..."
) )
@@ -857,12 +767,12 @@ class Neo4JStorage(BaseGraphStorage):
logger.warning( logger.warning(
"Neo4j: inclusive search mode is not supported in recursive query, using exact matching" "Neo4j: inclusive search mode is not supported in recursive query, using exact matching"
) )
return await self._robust_fallback(label, max_depth, min_degree) return await self._robust_fallback(node_label, max_depth, min_degree)
return result return result
async def _robust_fallback( async def _robust_fallback(
self, label: str, max_depth: int, min_degree: int = 0 self, node_label: str, max_depth: int, min_degree: int = 0
) -> KnowledgeGraph: ) -> KnowledgeGraph:
""" """
Fallback implementation when APOC plugin is not available or incompatible. Fallback implementation when APOC plugin is not available or incompatible.
@@ -895,12 +805,11 @@ class Neo4JStorage(BaseGraphStorage):
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = """ query = """
MATCH (a)-[r]-(b) MATCH (a:base {entity_id: $entity_id})-[r]-(b)
WHERE id(a) = toInteger($node_id)
WITH r, b, id(r) as edge_id, id(b) as target_id WITH r, b, id(r) as edge_id, id(b) as target_id
RETURN r, b, edge_id, target_id RETURN r, b, edge_id, target_id
""" """
results = await session.run(query, {"node_id": node.id}) results = await session.run(query, entity_id=node.id)
# Get all records and release database connection # Get all records and release database connection
records = await results.fetch( records = await results.fetch(
@@ -928,14 +837,14 @@ class Neo4JStorage(BaseGraphStorage):
edge_id = str(record["edge_id"]) edge_id = str(record["edge_id"])
if edge_id not in visited_edges: if edge_id not in visited_edges:
b_node = record["b"] b_node = record["b"]
target_id = str(record["target_id"]) target_id = b_node.get("entity_id")
if b_node.labels: # Only process if target node has labels if target_id: # Only process if target node has entity_id
# Create KnowledgeGraphNode for target # Create KnowledgeGraphNode for target
target_node = KnowledgeGraphNode( target_node = KnowledgeGraphNode(
id=f"{target_id}", id=f"{target_id}",
labels=list(b_node.labels), labels=[label for label in b_node.labels if label != "base"],
properties=dict(b_node), properties=dict(b_node.properties),
) )
# Create KnowledgeGraphEdge # Create KnowledgeGraphEdge
@@ -961,11 +870,11 @@ class Neo4JStorage(BaseGraphStorage):
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = f""" query = """
MATCH (n:`{label}`) MATCH (n:base {entity_id: $entity_id})
RETURN id(n) as node_id, n RETURN id(n) as node_id, n
""" """
node_result = await session.run(query) node_result = await session.run(query, entity_id=node_label)
try: try:
node_record = await node_result.single() node_record = await node_result.single()
if not node_record: if not node_record:
@@ -973,9 +882,9 @@ class Neo4JStorage(BaseGraphStorage):
# Create initial KnowledgeGraphNode # Create initial KnowledgeGraphNode
start_node = KnowledgeGraphNode( start_node = KnowledgeGraphNode(
id=f"{node_record['node_id']}", id=f"{node_record['n'].get('entity_id')}",
labels=list(node_record["n"].labels), labels=[label for label in node_record["n"].labels if label != "base"],
properties=dict(node_record["n"]), properties=dict(node_record["n"].properties),
) )
finally: finally:
await node_result.consume() # Ensure results are consumed await node_result.consume() # Ensure results are consumed
@@ -999,11 +908,10 @@ class Neo4JStorage(BaseGraphStorage):
# Method 2: Query compatible with older versions # Method 2: Query compatible with older versions
query = """ query = """
MATCH (n) MATCH (n)
WITH DISTINCT labels(n) AS node_labels WHERE n.entity_id IS NOT NULL
UNWIND node_labels AS label RETURN DISTINCT n.entity_id AS label
RETURN DISTINCT label ORDER BY label
ORDER BY label
""" """
result = await session.run(query) result = await session.run(query)
labels = [] labels = []
@@ -1034,15 +942,13 @@ class Neo4JStorage(BaseGraphStorage):
Args: Args:
node_id: The label of the node to delete node_id: The label of the node to delete
""" """
label = self._ensure_label(node_id)
async def _do_delete(tx: AsyncManagedTransaction): async def _do_delete(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (n:`{label}`) MATCH (n:base {entity_id: $entity_id})
DETACH DELETE n DETACH DELETE n
""" """
result = await tx.run(query) result = await tx.run(query, entity_id=node_id)
logger.debug(f"Deleted node with label '{label}'") logger.debug(f"Deleted node with label '{node_id}'")
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
try: try:
@@ -1092,16 +998,13 @@ class Neo4JStorage(BaseGraphStorage):
edges: List of edges to be deleted, each edge is a (source, target) tuple edges: List of edges to be deleted, each edge is a (source, target) tuple
""" """
for source, target in edges: for source, target in edges:
source_label = self._ensure_label(source)
target_label = self._ensure_label(target)
async def _do_delete_edge(tx: AsyncManagedTransaction): async def _do_delete_edge(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (source:`{source_label}`)-[r]-(target:`{target_label}`) MATCH (source:base {entity_id: $source_entity_id})-[r]-(target:base {entity_id: $target_entity_id})
DELETE r DELETE r
""" """
result = await tx.run(query) result = await tx.run(query, source_entity_id=source, target_entity_id=target)
logger.debug(f"Deleted edge from '{source_label}' to '{target_label}'") logger.debug(f"Deleted edge from '{source}' to '{target}'")
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
try: try: