gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

wip fix Memgraph get_knowledge_graph issues

Browse source
This commit is contained in:
DavIvek 2025-07-11 11:55:53 +02:00
parent 63d554bad8
commit 8d295bd294
1 changed files with 109 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

189
lightrag/kg/memgraph_impl.py
View file

@ -734,7 +734,7 @@ class MemgraphStorage(BaseGraphStorage):
self, self,
node_label: str, node_label: str,
max_depth: int = 3, max_depth: int = 3,
max_nodes: int = MAX_GRAPH_NODES, max_nodes: int = None,
) -> KnowledgeGraph: ) -> KnowledgeGraph:
""" """
Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
@ -747,137 +747,166 @@ class MemgraphStorage(BaseGraphStorage):
Returns: Returns:
KnowledgeGraph object containing nodes and edges, with an is_truncated flag KnowledgeGraph object containing nodes and edges, with an is_truncated flag
indicating whether the graph was truncated due to max_nodes limit indicating whether the graph was truncated due to max_nodes limit
Raises:
Exception: If there is an error executing the query
""" """
if self._driver is None: # Get max_nodes from global_config if not provided
raise RuntimeError( if max_nodes is None:
"Memgraph driver is not initialized. Call 'await initialize()' first." max_nodes = self.global_config.get("max_graph_nodes", 1000)
) else:
# Limit max_nodes to not exceed global_config max_graph_nodes
max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000))
workspace_label = self._get_workspace_label()
result = KnowledgeGraph() result = KnowledgeGraph()
seen_nodes = set() seen_nodes = set()
seen_edges = set() seen_edges = set()
workspace_label = self._get_workspace_label()
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:
if node_label == "*": if node_label == "*":
# First check if database has any nodes # First check total node count to determine if graph is truncated
count_query = "MATCH (n) RETURN count(n) as total" count_query = (
f"MATCH (n:`{workspace_label}`) RETURN count(n) as total"
)
count_result = None count_result = None
total_count = 0
try: try:
count_result = await session.run(count_query) count_result = await session.run(count_query)
count_record = await count_result.single() count_record = await count_result.single()
if count_record:
total_count = count_record["total"] if count_record and count_record["total"] > max_nodes:
if total_count == 0: result.is_truncated = True
logger.debug("No nodes found in database") logger.info(
return result f"Graph truncated: {count_record['total']} nodes found, limited to {max_nodes}"
if total_count > max_nodes: )
result.is_truncated = True
logger.info(
f"Graph truncated: {total_count} nodes found, limited to {max_nodes}"
)
finally: finally:
if count_result: if count_result:
await count_result.consume() await count_result.consume()
# Run the main query to get nodes with highest degree # Run main query to get nodes with highest degree
main_query = f""" main_query = f"""
MATCH (n:`{workspace_label}`) MATCH (n:`{workspace_label}`)
OPTIONAL MATCH (n)-[r]-() OPTIONAL MATCH (n)-[r]-()
WITH n, COALESCE(count(r), 0) AS degree WITH n, COALESCE(count(r), 0) AS degree
ORDER BY degree DESC ORDER BY degree DESC
LIMIT $max_nodes LIMIT $max_nodes
WITH collect(n) AS kept_nodes WITH collect({{node: n}}) AS filtered_nodes
MATCH (a)-[r]-(b) UNWIND filtered_nodes AS node_info
WITH collect(node_info.node) AS kept_nodes, filtered_nodes
OPTIONAL MATCH (a)-[r]-(b)
WHERE a IN kept_nodes AND b IN kept_nodes WHERE a IN kept_nodes AND b IN kept_nodes
RETURN [node IN kept_nodes | {{node: node}}] AS node_info, RETURN filtered_nodes AS node_info,
collect(DISTINCT r) AS relationships collect(DISTINCT r) AS relationships
""" """
result_set = None result_set = None
try: try:
result_set = await session.run( result_set = await session.run(
main_query, {"max_nodes": max_nodes} main_query,
{"max_nodes": max_nodes},
) )
record = await result_set.single() record = await result_set.single()
if not record:
logger.debug("No record returned from main query")
return result
finally: finally:
if result_set: if result_set:
await result_set.consume() await result_set.consume()
else: else:
bfs_query = f""" # return await self._robust_fallback(node_label, max_depth, max_nodes)
# First try without limit to check if we need to truncate
full_query = f"""
MATCH (start:`{workspace_label}`) MATCH (start:`{workspace_label}`)
WHERE start.entity_id = $entity_id WHERE start.entity_id = $entity_id
WITH start WITH start
CALL {{ MATCH path = (start)-[*BFS ..{max_depth}]-(node:`{workspace_label}`)
WITH start WITH nodes(path) AS path_nodes, relationships(path) AS path_rels
MATCH path = (start)-[*BFS 0..{max_depth}]-(node) UNWIND path_nodes AS n
WITH nodes(path) AS path_nodes, relationships(path) AS path_rels WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
UNWIND path_nodes AS n WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
RETURN all_nodes, all_rels
}}
WITH all_nodes AS nodes, all_rels AS relationships, size(all_nodes) AS total_nodes_found
WITH
CASE
WHEN total_nodes_found <= {max_nodes} THEN nodes
ELSE nodes[0..{max_nodes}]
END AS limited_nodes,
relationships,
total_nodes_found,
total_nodes_found > {max_nodes} AS is_truncated
UNWIND relationships AS rel
WITH limited_nodes, rel, total_nodes_found, is_truncated
WHERE startNode(rel) IN limited_nodes AND endNode(rel) IN limited_nodes
WITH limited_nodes, collect(DISTINCT rel) AS limited_relationships, total_nodes_found, is_truncated
RETURN RETURN
[node IN limited_nodes | {{node: node}}] AS node_info, [node IN all_nodes | {{node: node}}] AS node_info,
limited_relationships AS relationships, all_rels AS relationships,
total_nodes_found, size(all_nodes) AS total_nodes
is_truncated
""" """
result_set = None # Try to get full result
full_result = None
try: try:
result_set = await session.run( full_result = await session.run(
bfs_query, full_query,
{ {
"entity_id": node_label, "entity_id": node_label,
}, },
) )
record = await result_set.single() full_record = await full_result.single()
if not record:
# If no record found, return empty KnowledgeGraph
if not full_record:
logger.debug(f"No nodes found for entity_id: {node_label}") logger.debug(f"No nodes found for entity_id: {node_label}")
return result return result
# Check if the query indicates truncation # If record found, check node count
if "is_truncated" in record and record["is_truncated"]: total_nodes = full_record["total_nodes"]
if total_nodes <= max_nodes:
# If node count is within limit, use full result directly
logger.debug(
f"Using full result with {total_nodes} nodes (no truncation needed)"
)
record = full_record
else:
# If node count exceeds limit, set truncated flag and run limited query
result.is_truncated = True result.is_truncated = True
logger.info( logger.info(
f"Graph truncated: breadth-first search limited to {max_nodes} nodes" f"Graph truncated: {total_nodes} nodes found, breadth-first search limited to {max_nodes}"
) )
finally: # Run limited query
if result_set: limited_query = f"""
await result_set.consume() MATCH (start:`{workspace_label}`)
WHERE start.entity_id = $entity_id
WITH start
MATCH path = (start)-[*BFS ..{max_depth}]-(node:`{workspace_label}`)
WITH nodes(path) AS path_nodes, relationships(path) AS path_rels
UNWIND path_nodes AS n
WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
WITH
CASE
WHEN size(all_nodes) <= $max_nodes THEN all_nodes
ELSE all_nodes[0..$max_nodes]
END AS limited_nodes,
all_rels
UNWIND all_rels AS rel
WITH limited_nodes, rel
WHERE startNode(rel) IN limited_nodes AND endNode(rel) IN limited_nodes
WITH limited_nodes, collect(DISTINCT rel) AS limited_relationships
RETURN
[node IN limited_nodes | {{node: node}}] AS node_info,
limited_relationships AS relationships
"""
# Process the record if it exists result_set = None
if record and record["node_info"]: try:
result_set = await session.run(
limited_query,
{
"entity_id": node_label,
"max_nodes": max_nodes,
},
)
record = await result_set.single()
finally:
if result_set:
await result_set.consume()
finally:
if full_result:
await full_result.consume()
if record:
# Handle nodes (compatible with multi-label cases)
for node_info in record["node_info"]: for node_info in record["node_info"]:
node = node_info["node"] node = node_info["node"]
node_id = node.id node_id = node.id
if node_id not in seen_nodes: if node_id not in seen_nodes:
seen_nodes.add(node_id)
result.nodes.append( result.nodes.append(
KnowledgeGraphNode( KnowledgeGraphNode(
id=f"{node_id}", id=f"{node_id}",
@ -885,11 +914,12 @@ class MemgraphStorage(BaseGraphStorage):
properties=dict(node), properties=dict(node),
) )
) )
seen_nodes.add(node_id)
# Handle relationships (including direction information)
for rel in record["relationships"]: for rel in record["relationships"]:
edge_id = rel.id edge_id = rel.id
if edge_id not in seen_edges: if edge_id not in seen_edges:
seen_edges.add(edge_id)
start = rel.start_node start = rel.start_node
end = rel.end_node end = rel.end_node
result.edges.append( result.edges.append(
@ -901,14 +931,13 @@ class MemgraphStorage(BaseGraphStorage):
properties=dict(rel), properties=dict(rel),
) )
) )
seen_edges.add(edge_id)
logger.info( logger.info(
f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}" f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
) )
except Exception as e: except Exception as e:
logger.error(f"Error getting knowledge graph: {str(e)}") logger.error(f"Error during subgraph query for {node_label}: {str(e)}")
# Return empty but properly initialized KnowledgeGraph on error
return KnowledgeGraph()
return result return result