From b4543561f6f0fe9db025ef88798ba297774bbc9c Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Mon, 24 Feb 2025 03:32:33 +0800
Subject: [PATCH] Limit subgraph size to 500 nodes in NetworkXStorage

- Add max_graph_nodes check
- Reduce subgraph by degree
- Log graph size reduction
---
 lightrag/kg/networkx_impl.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py
index a6958704..e68d2887 100644
--- a/lightrag/kg/networkx_impl.py
+++ b/lightrag/kg/networkx_impl.py
@@ -219,10 +219,21 @@ class NetworkXStorage(BaseGraphStorage):
             if not nodes_to_explore:
                 logger.warning(f"No nodes found with label {node_label}")
                 return result
-                
+
             # Get subgraph using ego_graph
             subgraph = nx.ego_graph(self._graph, nodes_to_explore[0], radius=max_depth)
         
+        # Check if number of nodes exceeds max_graph_nodes
+        max_graph_nodes=500
+        if len(subgraph.nodes()) > max_graph_nodes:
+            origin_nodes=len(subgraph.nodes())
+            node_degrees = dict(subgraph.degree())
+            top_nodes = sorted(node_degrees.items(), key=lambda x: x[1], reverse=True)[:max_graph_nodes]
+            top_node_ids = [node[0] for node in top_nodes]
+            # Create new subgraph with only top nodes
+            subgraph = subgraph.subgraph(top_node_ids)
+            logger.info(f"Reduced graph from {origin_nodes} nodes to {max_graph_nodes} nodes by degree (depth={max_depth})")
+
         # Add nodes to result
         for node in subgraph.nodes():
             if str(node) in seen_nodes: