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

Refactor Neo4JStorage methods for robustness and clarity.

Browse Source 
- Add error handling and resource cleanup
- Improve method documentation
- Optimize result consumption
This commit is contained in:
yangdx
2025-03-08 10:19:20 +08:00
parent fcb04e47e5
commit 84222b8b76
1 changed files with 255 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

412
lightrag/kg/neo4j_impl.py
View File

@@ -163,13 +163,14 @@ class Neo4JStorage(BaseGraphStorage):
} }
async def close(self): async def close(self):
"""Close the Neo4j driver and release all resources"""
if self._driver: if self._driver:
await self._driver.close() await self._driver.close()
self._driver = None self._driver = None
async def __aexit__(self, exc_type, exc, tb): async def __aexit__(self, exc_type, exc, tb):
if self._driver: """Ensure driver is closed when context manager exits"""
await self._driver.close() await self.close()
async def index_done_callback(self) -> None: async def index_done_callback(self) -> None:
# Noe4J handles persistence automatically # Noe4J handles persistence automatically
@@ -187,33 +188,72 @@ class Neo4JStorage(BaseGraphStorage):
return clean_label 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
Args:
node_id: Label of the node to check
Returns:
bool: True if node exists, False otherwise
Raises:
ValueError: If node_id is invalid
Exception: If there is an error executing the query
"""
entity_name_label = await self._ensure_label(node_id) entity_name_label = await 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:
query = ( try:
f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists" query = f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists"
) result = await session.run(query)
result = await session.run(query) 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:
logger.error(
f"Error checking node existence for {entity_name_label}: {str(e)}"
)
await result.consume() # Ensure results are consumed even on error
raise
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool: async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
entity_name_label_source = source_node_id.strip('"') """
entity_name_label_target = target_node_id.strip('"') Check if an edge exists between two nodes
Args:
source_node_id: Label of the source node
target_node_id: Label of the target node
Returns:
bool: True if edge exists, False otherwise
Raises:
ValueError: If either node_id is invalid
Exception: If there is an error executing the query
"""
entity_name_label_source = await self._ensure_label(source_node_id)
entity_name_label_target = await 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 = ( try:
f"MATCH (a:`{entity_name_label_source}`)-[r]-(b:`{entity_name_label_target}`) " query = (
"RETURN COUNT(r) > 0 AS edgeExists" f"MATCH (a:`{entity_name_label_source}`)-[r]-(b:`{entity_name_label_target}`) "
) "RETURN COUNT(r) > 0 AS edgeExists"
result = await session.run(query) )
single_result = await result.single() result = await session.run(query)
await result.consume() # Ensure result is fully consumed single_result = await result.single()
return single_result["edgeExists"] await result.consume() # Ensure result is fully consumed
return single_result["edgeExists"]
except Exception as e:
logger.error(
f"Error checking edge existence between {entity_name_label_source} and {entity_name_label_target}: {str(e)}"
)
await result.consume() # Ensure results are consumed even on error
raise
async def get_node(self, node_id: str) -> dict[str, str] | None: async def get_node(self, node_id: str) -> dict[str, str] | None:
"""Get node by its label identifier. """Get node by its label identifier.
@@ -224,29 +264,40 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
dict: Node properties if found dict: Node properties if found
None: If node not found None: If node not found
Raises:
ValueError: If node_id is invalid
Exception: If there is an error executing the query
""" """
entity_name_label = await 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:
entity_name_label = await self._ensure_label(node_id) try:
query = f"MATCH (n:`{entity_name_label}`) RETURN n" query = f"MATCH (n:`{entity_name_label}`) RETURN n"
result = await session.run(query) result = await session.run(query)
records = await result.fetch( try:
2 records = await result.fetch(
) # Get up to 2 records to check for duplicates 2
await result.consume() # Ensure result is fully consumed ) # Get up to 2 records to check for duplicates
if len(records) > 1:
logger.warning( if len(records) > 1:
f"Multiple nodes found with label '{entity_name_label}'. Using first node." logger.warning(
) f"Multiple nodes found with label '{entity_name_label}'. Using first node."
if records: )
node = records[0]["n"] if records:
node_dict = dict(node) node = records[0]["n"]
logger.debug( node_dict = dict(node)
f"{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}" logger.debug(
) f"{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}"
return node_dict )
return None return node_dict
return None
finally:
await result.consume() # Ensure result is fully consumed
except Exception as e:
logger.error(f"Error getting node for {entity_name_label}: {str(e)}")
raise
async def node_degree(self, node_id: str) -> int: async def node_degree(self, node_id: str) -> int:
"""Get the degree (number of relationships) of a node with the given label. """Get the degree (number of relationships) of a node with the given label.
@@ -258,39 +309,63 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
int: The number of relationships the node has, or 0 if no node found int: The number of relationships the node has, or 0 if no node found
Raises:
ValueError: If node_id is invalid
Exception: If there is an error executing the query
""" """
entity_name_label = node_id.strip('"') entity_name_label = await 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:
query = f""" try:
MATCH (n:`{entity_name_label}`) query = f"""
OPTIONAL MATCH (n)-[r]-() MATCH (n:`{entity_name_label}`)
RETURN n, COUNT(r) AS degree OPTIONAL MATCH (n)-[r]-()
""" RETURN n, COUNT(r) AS degree
result = await session.run(query) """
records = await result.fetch(100) result = await session.run(query)
await result.consume() # Ensure result is fully consumed try:
records = await result.fetch(100)
if not records: if not records:
logger.warning(f"No node found with label '{entity_name_label}'") logger.warning(
return 0 f"No node found with label '{entity_name_label}'"
)
return 0
if len(records) > 1: if len(records) > 1:
logger.warning( logger.warning(
f"Multiple nodes ({len(records)}) found with label '{entity_name_label}', using first node's 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
finally:
await result.consume() # Ensure result is fully consumed
except Exception as e:
logger.error(
f"Error getting node degree for {entity_name_label}: {str(e)}"
) )
raise
degree = records[0]["degree"]
logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{degree}"
)
return degree
async def edge_degree(self, src_id: str, tgt_id: str) -> int: async def edge_degree(self, src_id: str, tgt_id: str) -> int:
entity_name_label_source = src_id.strip('"') """Get the total degree (sum of relationships) of two nodes.
entity_name_label_target = tgt_id.strip('"')
Args:
src_id: Label of the source node
tgt_id: Label of the target node
Returns:
int: Sum of the degrees of both nodes
"""
entity_name_label_source = await self._ensure_label(src_id)
entity_name_label_target = await self._ensure_label(tgt_id)
src_degree = await self.node_degree(entity_name_label_source) src_degree = await self.node_degree(entity_name_label_source)
trg_degree = await self.node_degree(entity_name_label_target) trg_degree = await self.node_degree(entity_name_label_target)
@@ -299,17 +374,27 @@ class Neo4JStorage(BaseGraphStorage):
trg_degree = 0 if trg_degree is None else trg_degree trg_degree = 0 if trg_degree is None else trg_degree
degrees = int(src_degree) + int(trg_degree) degrees = int(src_degree) + int(trg_degree)
logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:src_Degree+trg_degree:result:{degrees}"
)
return degrees return degrees
async def get_edge( async def get_edge(
self, source_node_id: str, target_node_id: str self, source_node_id: str, target_node_id: str
) -> dict[str, str] | None: ) -> dict[str, str] | None:
"""Get edge properties between two nodes.
Args:
source_node_id: Label of the source node
target_node_id: Label of the target node
Returns:
dict: Edge properties if found, default properties if not found or on error
Raises:
ValueError: If either node_id is invalid
Exception: If there is an error executing the query
"""
try: try:
entity_name_label_source = source_node_id.strip('"') entity_name_label_source = await self._ensure_label(source_node_id)
entity_name_label_target = target_node_id.strip('"') entity_name_label_target = await 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"
@@ -320,109 +405,123 @@ class Neo4JStorage(BaseGraphStorage):
""" """
result = await session.run(query) result = await session.run(query)
records = await result.fetch(2) # Get up to 2 records to check for duplicates try:
await result.consume() # Ensure result is fully consumed before processing records records = await result.fetch(
2
if len(records) > 1: ) # Get up to 2 records to check for duplicates
logger.warning(
f"Multiple edges found between '{entity_name_label_source}' and '{entity_name_label_target}'. Using first edge." if len(records) > 1:
logger.warning(
f"Multiple edges found between '{entity_name_label_source}' and '{entity_name_label_target}'. Using first edge."
)
if records:
try:
edge_result = dict(records[0]["edge_properties"])
logger.debug(f"Result: {edge_result}")
# Ensure required keys exist with defaults
required_keys = {
"weight": 0.0,
"source_id": None,
"description": None,
"keywords": None,
}
for key, default_value in required_keys.items():
if key not in edge_result:
edge_result[key] = default_value
logger.warning(
f"Edge between {entity_name_label_source} and {entity_name_label_target} "
f"missing {key}, using default: {default_value}"
)
logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{edge_result}"
)
return edge_result
except (KeyError, TypeError, ValueError) as e:
logger.error(
f"Error processing edge properties between {entity_name_label_source} "
f"and {entity_name_label_target}: {str(e)}"
)
# Return default edge properties on error
return {
"weight": 0.0,
"source_id": None,
"description": None,
"keywords": None,
}
logger.debug(
f"{inspect.currentframe().f_code.co_name}: No edge found between {entity_name_label_source} and {entity_name_label_target}"
) )
if records: # Return default edge properties when no edge found
try: return {
edge_result = dict(records[0]["edge_properties"]) "weight": 0.0,
logger.debug(f"Result: {edge_result}") "source_id": None,
# Ensure required keys exist with defaults "description": None,
required_keys = { "keywords": None,
"weight": 0.0, }
"source_id": None, finally:
"description": None, await result.consume() # Ensure result is fully consumed
"keywords": None,
}
for key, default_value in required_keys.items():
if key not in edge_result:
edge_result[key] = default_value
logger.warning(
f"Edge between {entity_name_label_source} and {entity_name_label_target} "
f"missing {key}, using default: {default_value}"
)
logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{edge_result}"
)
return edge_result
except (KeyError, TypeError, ValueError) as e:
logger.error(
f"Error processing edge properties between {entity_name_label_source} "
f"and {entity_name_label_target}: {str(e)}"
)
# Return default edge properties on error
return {
"weight": 0.0,
"description": None,
"keywords": None,
"source_id": None,
}
logger.debug(
f"{inspect.currentframe().f_code.co_name}: No edge found between {entity_name_label_source} and {entity_name_label_target}"
)
# Return default edge properties when no edge found
return {
"weight": 0.0,
"description": None,
"keywords": None,
"source_id": None,
}
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error in get_edge between {source_node_id} and {target_node_id}: {str(e)}" f"Error in get_edge between {source_node_id} and {target_node_id}: {str(e)}"
) )
# Return default edge properties on error raise
return {
"weight": 0.0,
"description": None,
"keywords": None,
"source_id": None,
}
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None: async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
node_label = source_node_id.strip('"') """Retrieves all edges (relationships) for a particular node identified by its label.
Args:
source_node_id: Label of the node to get edges for
Returns:
list[tuple[str, str]]: List of (source_label, target_label) tuples representing edges
None: If no edges found
Raises:
ValueError: If source_node_id is invalid
Exception: If there is an error executing the query
""" """
Retrieves all edges (relationships) for a particular node identified by its label. try:
:return: List of dictionaries containing edge information node_label = await 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(
database=self._DATABASE, default_access_mode="READ"
) as session:
results = await session.run(query)
edges = []
try:
async for record in results:
source_node = record["n"]
connected_node = record["connected"]
source_label = ( query = f"""MATCH (n:`{node_label}`)
list(source_node.labels)[0] if source_node.labels else None OPTIONAL MATCH (n)-[r]-(connected)
) RETURN n, r, connected"""
target_label = (
list(connected_node.labels)[0]
if connected_node and connected_node.labels
else None
)
if source_label and target_label: async with self._driver.session(
edges.append((source_label, target_label)) database=self._DATABASE, default_access_mode="READ"
finally: ) as session:
await ( try:
results.consume() results = await session.run(query)
) # Ensure results are consumed even if processing fails edges = []
return edges async for record in results:
source_node = record["n"]
connected_node = record["connected"]
source_label = (
list(source_node.labels)[0] if source_node.labels else None
)
target_label = (
list(connected_node.labels)[0]
if connected_node and connected_node.labels
else None
)
if source_label and target_label:
edges.append((source_label, target_label))
await results.consume() # Ensure results are consumed
return edges if edges else None
except Exception as e:
logger.error(f"Error getting edges for node {node_label}: {str(e)}")
await results.consume() # Ensure results are consumed even on error
raise
except Exception as e:
logger.error(f"Error in get_node_edges for {source_node_id}: {str(e)}")
raise
@retry( @retry(
stop=stop_after_attempt(3), stop=stop_after_attempt(3),
@@ -838,7 +937,6 @@ class Neo4JStorage(BaseGraphStorage):
RETURN DISTINCT label RETURN DISTINCT label
ORDER BY label ORDER BY label
""" """
result = await session.run(query) result = await session.run(query)
labels = [] labels = []
try: try: