From 002948d3422fff442ac1552cb7d6041933aabf0f Mon Sep 17 00:00:00 2001 From: yangdx Date: Tue, 4 Mar 2025 16:08:05 +0800 Subject: [PATCH] Added search mode and min degree filtering for NetworkX - Implemented exact and inclusive search modes - Added min degree filtering for nodes - Updated API to parse label for search options --- lightrag/api/routers/graph_routes.py | 38 +++++++++++++++++++++++++++- lightrag/kg/networkx_impl.py | 21 ++++++++++++--- lightrag/lightrag.py | 18 +++++++++++-- 3 files changed, 70 insertions(+), 7 deletions(-) diff --git a/lightrag/api/routers/graph_routes.py b/lightrag/api/routers/graph_routes.py index e6f894a2..ddaac676 100644 --- a/lightrag/api/routers/graph_routes.py +++ b/lightrag/api/routers/graph_routes.py @@ -34,6 +34,11 @@ def create_graph_routes(rag, api_key: Optional[str] = None): 2. Followed by nodes directly connected to the matching nodes 3. Finally, the degree of the nodes Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000) + Control search mode by label content: + 1. only label-name : exact search with the label name (selecting from the label list return previously) + 2. label-name follow by '>n' : exact search of nodes with degree more than n + 3. label-name follow by* : inclusive search of nodes with degree more than n + 4. label-name follow by '>n*' : inclusive search Args: label (str): Label to get knowledge graph for @@ -42,6 +47,37 @@ def create_graph_routes(rag, api_key: Optional[str] = None): Returns: Dict[str, List[str]]: Knowledge graph for label """ - return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth) + # Parse label to extract search mode and min degree if specified + search_mode = "exact" # Default search mode + min_degree = 0 # Default minimum degree + original_label = label + + # First check if label ends with * + if label.endswith("*"): + search_mode = "inclusive" # Always set to inclusive if ends with * + label = label[:-1].strip() # Remove trailing * + + # Try to parse >n if it exists + if ">" in label: + try: + degree_pos = label.rfind(">") + degree_str = label[degree_pos + 1:].strip() + min_degree = int(degree_str) + 1 + label = label[:degree_pos].strip() + except ValueError: + # If degree parsing fails, just remove * and keep the rest as label + label = original_label[:-1].strip() + # If no *, check for >n pattern + elif ">" in label: + try: + degree_pos = label.rfind(">") + degree_str = label[degree_pos + 1:].strip() + min_degree = int(degree_str) + 1 + label = label[:degree_pos].strip() + except ValueError: + # If degree parsing fails, treat the whole string as label + label = original_label + + return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth, search_mode=search_mode, min_degree=min_degree) return router diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index 563fc554..ae46e18d 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -232,7 +232,7 @@ class NetworkXStorage(BaseGraphStorage): return sorted(list(labels)) async def get_knowledge_graph( - self, node_label: str, max_depth: int = 5 + self, node_label: str, max_depth: int = 5, search_mode: str = "exact", min_degree: int = 0 ) -> KnowledgeGraph: """ Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. @@ -245,6 +245,8 @@ class NetworkXStorage(BaseGraphStorage): Args: node_label: Label of the starting node max_depth: Maximum depth of the subgraph + search_mode (str, optional): Search mode, either "exact" or "inclusive". Defaults to "exact". + min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0. Returns: KnowledgeGraph object containing nodes and edges @@ -262,11 +264,16 @@ class NetworkXStorage(BaseGraphStorage): graph.copy() ) # Create a copy to avoid modifying the original graph else: - # Find nodes with matching node id (partial match) + # Find nodes with matching node id based on search_mode nodes_to_explore = [] for n, attr in graph.nodes(data=True): - if node_label in str(n): # Use partial matching - nodes_to_explore.append(n) + node_str = str(n) + if search_mode == "exact": + if node_label == node_str: # Use exact matching + nodes_to_explore.append(n) + else: # inclusive mode + if node_label in node_str: # Use partial matching + nodes_to_explore.append(n) if not nodes_to_explore: logger.warning(f"No nodes found with label {node_label}") @@ -277,6 +284,12 @@ class NetworkXStorage(BaseGraphStorage): for start_node in nodes_to_explore: node_subgraph = nx.ego_graph(graph, start_node, radius=max_depth) combined_subgraph = nx.compose(combined_subgraph, node_subgraph) + + # Filter nodes based on min_degree + if min_degree > 0: + nodes_to_keep = [node for node, degree in combined_subgraph.degree() if degree >= min_degree] + combined_subgraph = combined_subgraph.subgraph(nodes_to_keep) + subgraph = combined_subgraph # Check if number of nodes exceeds max_graph_nodes diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 74121cf1..2dfd95ef 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -504,10 +504,24 @@ class LightRAG: return text async def get_knowledge_graph( - self, node_label: str, max_depth: int + self, node_label: str, max_depth: int, search_mode: str = "exact", min_degree: int = 0 ) -> KnowledgeGraph: + """Get knowledge graph for a given label + + Args: + node_label (str): Label to get knowledge graph for + max_depth (int): Maximum depth of graph + search_mode (str, optional): Search mode, either "exact" or "inclusive". Defaults to "exact". + min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0. + + Returns: + KnowledgeGraph: Knowledge graph containing nodes and edges + """ return await self.chunk_entity_relation_graph.get_knowledge_graph( - node_label=node_label, max_depth=max_depth + node_label=node_label, + max_depth=max_depth, + search_mode=search_mode, + min_degree=min_degree ) def _get_storage_class(self, storage_name: str) -> Callable[..., Any]: