From b13c947360091b7110497033b19c5c17b1ffd508 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Mon, 24 Feb 2025 02:37:42 +0800
Subject: [PATCH] Implement get_all_labels and get_knowledge_graph methods in
 NetworkXStorage

---
 lightrag/kg/networkx_impl.py | 107 ++++++++++++++++++++++++++++++++++-
 1 file changed, 104 insertions(+), 3 deletions(-)

diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py
index 9850b8c4..a6958704 100644
--- a/lightrag/kg/networkx_impl.py
+++ b/lightrag/kg/networkx_impl.py
@@ -5,7 +5,7 @@ from typing import Any, final
 import numpy as np
 
 
-from lightrag.types import KnowledgeGraph
+from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
 from lightrag.utils import (
     logger,
 )
@@ -169,9 +169,110 @@ class NetworkXStorage(BaseGraphStorage):
                 self._graph.remove_edge(source, target)
 
     async def get_all_labels(self) -> list[str]:
-        raise NotImplementedError
+        """
+        Get all node labels in the graph
+        Returns:
+            [label1, label2, ...]  # Alphabetically sorted label list
+        """
+        # Get all labels from nodes
+        labels = set()
+        for node in self._graph.nodes():
+            # node_data = dict(self._graph.nodes[node])
+            # if "entity_type" in node_data:
+            #     if isinstance(node_data["entity_type"], list):
+            #         labels.update(node_data["entity_type"])
+            #     else:
+            #         labels.add(node_data["entity_type"])
+            labels.add(str(node))  # Add node id as a label
+        
+        # Return sorted list
+        return sorted(list(labels))
 
     async def get_knowledge_graph(
         self, node_label: str, max_depth: int = 5
     ) -> KnowledgeGraph:
-        raise NotImplementedError
+        """
+        Get complete connected subgraph for specified node (including the starting node itself)
+        
+        Args:
+            node_label: Label of the starting node
+            max_depth: Maximum depth of the subgraph
+            
+        Returns:
+            KnowledgeGraph object containing nodes and edges
+        """
+        result = KnowledgeGraph()
+        seen_nodes = set()
+        seen_edges = set()
+        
+        # Handle special case for "*" label
+        if node_label == "*":
+            # For "*", return the entire graph including all nodes and edges
+            subgraph = self._graph.copy()  # Create a copy to avoid modifying the original graph
+        else:
+            # Find nodes with matching node id (partial match)
+            nodes_to_explore = []
+            for n, attr in self._graph.nodes(data=True):
+                if node_label in str(n):  # Use partial matching
+                    nodes_to_explore.append(n)
+            
+            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)
+        
+        # Add nodes to result
+        for node in subgraph.nodes():
+            if str(node) in seen_nodes:
+                continue
+                
+            node_data = dict(subgraph.nodes[node])
+            # Get entity_type as labels
+            labels = []
+            if "entity_type" in node_data:
+                if isinstance(node_data["entity_type"], list):
+                    labels.extend(node_data["entity_type"])
+                else:
+                    labels.append(node_data["entity_type"])
+            
+            # Create node with properties
+            node_properties = {k: v for k, v in node_data.items()}
+
+            result.nodes.append(
+                KnowledgeGraphNode(
+                    id=str(node),
+                    labels=[str(node)],
+                    properties=node_properties
+                )
+            )
+            seen_nodes.add(str(node))
+        
+        # Add edges to result
+        for edge in subgraph.edges():
+            source, target = edge
+            edge_id = f"{source}-{target}"
+            if edge_id in seen_edges:
+                continue
+                
+            edge_data = dict(subgraph.edges[edge])
+            
+            # Create edge with complete information
+            result.edges.append(
+                KnowledgeGraphEdge(
+                    id=edge_id,
+                    type="DIRECTED",
+                    source=str(source),
+                    target=str(target),
+                    properties=edge_data,
+                )
+            )
+            seen_edges.add(edge_id)
+        
+        # logger.info(result.edges)
+
+        logger.info(
+            f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
+        )
+        return result