From 8d295bd2941b8350d0bd4b0b4ad2fef5dad1134b Mon Sep 17 00:00:00 2001 From: DavIvek Date: Fri, 11 Jul 2025 11:55:53 +0200 Subject: [PATCH] wip fix Memgraph get_knowledge_graph issues --- lightrag/kg/memgraph_impl.py | 189 ++++++++++++++++++++--------------- 1 file changed, 109 insertions(+), 80 deletions(-) diff --git a/lightrag/kg/memgraph_impl.py b/lightrag/kg/memgraph_impl.py index dc78ca19..d0044499 100644 --- a/lightrag/kg/memgraph_impl.py +++ b/lightrag/kg/memgraph_impl.py @@ -734,7 +734,7 @@ class MemgraphStorage(BaseGraphStorage): self, node_label: str, max_depth: int = 3, - max_nodes: int = MAX_GRAPH_NODES, + max_nodes: int = None, ) -> KnowledgeGraph: """ Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. @@ -747,137 +747,166 @@ class MemgraphStorage(BaseGraphStorage): Returns: KnowledgeGraph object containing nodes and edges, with an is_truncated flag 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: - raise RuntimeError( - "Memgraph driver is not initialized. Call 'await initialize()' first." - ) + # Get max_nodes from global_config if not provided + if max_nodes is None: + 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() seen_nodes = set() seen_edges = set() - workspace_label = self._get_workspace_label() + async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: try: if node_label == "*": - # First check if database has any nodes - count_query = "MATCH (n) RETURN count(n) as total" + # First check total node count to determine if graph is truncated + count_query = ( + f"MATCH (n:`{workspace_label}`) RETURN count(n) as total" + ) count_result = None - total_count = 0 try: count_result = await session.run(count_query) count_record = await count_result.single() - if count_record: - total_count = count_record["total"] - if total_count == 0: - logger.debug("No nodes found in database") - return result - if total_count > max_nodes: - result.is_truncated = True - logger.info( - f"Graph truncated: {total_count} nodes found, limited to {max_nodes}" - ) + + if count_record and count_record["total"] > max_nodes: + result.is_truncated = True + logger.info( + f"Graph truncated: {count_record['total']} nodes found, limited to {max_nodes}" + ) finally: if count_result: 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""" MATCH (n:`{workspace_label}`) OPTIONAL MATCH (n)-[r]-() WITH n, COALESCE(count(r), 0) AS degree ORDER BY degree DESC LIMIT $max_nodes - WITH collect(n) AS kept_nodes - MATCH (a)-[r]-(b) + WITH collect({{node: n}}) AS filtered_nodes + 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 - RETURN [node IN kept_nodes | {{node: node}}] AS node_info, - collect(DISTINCT r) AS relationships + RETURN filtered_nodes AS node_info, + collect(DISTINCT r) AS relationships """ result_set = None try: result_set = await session.run( - main_query, {"max_nodes": max_nodes} + main_query, + {"max_nodes": max_nodes}, ) record = await result_set.single() - if not record: - logger.debug("No record returned from main query") - return result finally: if result_set: await result_set.consume() 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}`) WHERE start.entity_id = $entity_id WITH start - CALL {{ - WITH start - MATCH path = (start)-[*BFS 0..{max_depth}]-(node) - 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 - 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 + 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 RETURN - [node IN limited_nodes | {{node: node}}] AS node_info, - limited_relationships AS relationships, - total_nodes_found, - is_truncated + [node IN all_nodes | {{node: node}}] AS node_info, + all_rels AS relationships, + size(all_nodes) AS total_nodes """ - result_set = None + # Try to get full result + full_result = None try: - result_set = await session.run( - bfs_query, + full_result = await session.run( + full_query, { "entity_id": node_label, }, ) - record = await result_set.single() - if not record: + full_record = await full_result.single() + + # If no record found, return empty KnowledgeGraph + if not full_record: logger.debug(f"No nodes found for entity_id: {node_label}") return result - # Check if the query indicates truncation - if "is_truncated" in record and record["is_truncated"]: + # If record found, check node count + 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 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: - if result_set: - await result_set.consume() + # Run limited query + limited_query = f""" + 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 - if record and record["node_info"]: + result_set = None + 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"]: node = node_info["node"] node_id = node.id if node_id not in seen_nodes: - seen_nodes.add(node_id) result.nodes.append( KnowledgeGraphNode( id=f"{node_id}", @@ -885,11 +914,12 @@ class MemgraphStorage(BaseGraphStorage): properties=dict(node), ) ) + seen_nodes.add(node_id) + # Handle relationships (including direction information) for rel in record["relationships"]: edge_id = rel.id if edge_id not in seen_edges: - seen_edges.add(edge_id) start = rel.start_node end = rel.end_node result.edges.append( @@ -901,14 +931,13 @@ class MemgraphStorage(BaseGraphStorage): properties=dict(rel), ) ) + seen_edges.add(edge_id) - logger.info( - f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}" - ) + logger.info( + f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}" + ) except Exception as e: - logger.error(f"Error getting knowledge graph: {str(e)}") - # Return empty but properly initialized KnowledgeGraph on error - return KnowledgeGraph() + logger.error(f"Error during subgraph query for {node_label}: {str(e)}") return result