yangdx commited on
Commit
4aa7946
·
1 Parent(s): 6e653a1

Added search mode and min degree filtering for NetworkX

Browse files

- 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 CHANGED
@@ -34,6 +34,11 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
34
  2. Followed by nodes directly connected to the matching nodes
35
  3. Finally, the degree of the nodes
36
  Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000)
 
 
 
 
 
37
 
38
  Args:
39
  label (str): Label to get knowledge graph for
@@ -42,6 +47,37 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
42
  Returns:
43
  Dict[str, List[str]]: Knowledge graph for label
44
  """
45
- return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
46
 
47
  return router
 
34
  2. Followed by nodes directly connected to the matching nodes
35
  3. Finally, the degree of the nodes
36
  Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000)
37
+ Control search mode by label content:
38
+ 1. only label-name : exact search with the label name (selecting from the label list return previously)
39
+ 2. label-name follow by '>n' : exact search of nodes with degree more than n
40
+ 3. label-name follow by* : inclusive search of nodes with degree more than n
41
+ 4. label-name follow by '>n*' : inclusive search
42
 
43
  Args:
44
  label (str): Label to get knowledge graph for
 
47
  Returns:
48
  Dict[str, List[str]]: Knowledge graph for label
49
  """
50
+ # Parse label to extract search mode and min degree if specified
51
+ search_mode = "exact" # Default search mode
52
+ min_degree = 0 # Default minimum degree
53
+ original_label = label
54
+
55
+ # First check if label ends with *
56
+ if label.endswith("*"):
57
+ search_mode = "inclusive" # Always set to inclusive if ends with *
58
+ label = label[:-1].strip() # Remove trailing *
59
+
60
+ # Try to parse >n if it exists
61
+ if ">" in label:
62
+ try:
63
+ degree_pos = label.rfind(">")
64
+ degree_str = label[degree_pos + 1:].strip()
65
+ min_degree = int(degree_str) + 1
66
+ label = label[:degree_pos].strip()
67
+ except ValueError:
68
+ # If degree parsing fails, just remove * and keep the rest as label
69
+ label = original_label[:-1].strip()
70
+ # If no *, check for >n pattern
71
+ elif ">" in label:
72
+ try:
73
+ degree_pos = label.rfind(">")
74
+ degree_str = label[degree_pos + 1:].strip()
75
+ min_degree = int(degree_str) + 1
76
+ label = label[:degree_pos].strip()
77
+ except ValueError:
78
+ # If degree parsing fails, treat the whole string as label
79
+ label = original_label
80
+
81
+ return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth, search_mode=search_mode, min_degree=min_degree)
82
 
83
  return router
lightrag/kg/networkx_impl.py CHANGED
@@ -232,7 +232,7 @@ class NetworkXStorage(BaseGraphStorage):
232
  return sorted(list(labels))
233
 
234
  async def get_knowledge_graph(
235
- self, node_label: str, max_depth: int = 5
236
  ) -> KnowledgeGraph:
237
  """
238
  Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
@@ -245,6 +245,8 @@ class NetworkXStorage(BaseGraphStorage):
245
  Args:
246
  node_label: Label of the starting node
247
  max_depth: Maximum depth of the subgraph
 
 
248
 
249
  Returns:
250
  KnowledgeGraph object containing nodes and edges
@@ -262,11 +264,16 @@ class NetworkXStorage(BaseGraphStorage):
262
  graph.copy()
263
  ) # Create a copy to avoid modifying the original graph
264
  else:
265
- # Find nodes with matching node id (partial match)
266
  nodes_to_explore = []
267
  for n, attr in graph.nodes(data=True):
268
- if node_label in str(n): # Use partial matching
269
- nodes_to_explore.append(n)
 
 
 
 
 
270
 
271
  if not nodes_to_explore:
272
  logger.warning(f"No nodes found with label {node_label}")
@@ -277,6 +284,12 @@ class NetworkXStorage(BaseGraphStorage):
277
  for start_node in nodes_to_explore:
278
  node_subgraph = nx.ego_graph(graph, start_node, radius=max_depth)
279
  combined_subgraph = nx.compose(combined_subgraph, node_subgraph)
 
 
 
 
 
 
280
  subgraph = combined_subgraph
281
 
282
  # Check if number of nodes exceeds max_graph_nodes
 
232
  return sorted(list(labels))
233
 
234
  async def get_knowledge_graph(
235
+ self, node_label: str, max_depth: int = 5, search_mode: str = "exact", min_degree: int = 0
236
  ) -> KnowledgeGraph:
237
  """
238
  Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
 
245
  Args:
246
  node_label: Label of the starting node
247
  max_depth: Maximum depth of the subgraph
248
+ search_mode (str, optional): Search mode, either "exact" or "inclusive". Defaults to "exact".
249
+ min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0.
250
 
251
  Returns:
252
  KnowledgeGraph object containing nodes and edges
 
264
  graph.copy()
265
  ) # Create a copy to avoid modifying the original graph
266
  else:
267
+ # Find nodes with matching node id based on search_mode
268
  nodes_to_explore = []
269
  for n, attr in graph.nodes(data=True):
270
+ node_str = str(n)
271
+ if search_mode == "exact":
272
+ if node_label == node_str: # Use exact matching
273
+ nodes_to_explore.append(n)
274
+ else: # inclusive mode
275
+ if node_label in node_str: # Use partial matching
276
+ nodes_to_explore.append(n)
277
 
278
  if not nodes_to_explore:
279
  logger.warning(f"No nodes found with label {node_label}")
 
284
  for start_node in nodes_to_explore:
285
  node_subgraph = nx.ego_graph(graph, start_node, radius=max_depth)
286
  combined_subgraph = nx.compose(combined_subgraph, node_subgraph)
287
+
288
+ # Filter nodes based on min_degree
289
+ if min_degree > 0:
290
+ nodes_to_keep = [node for node, degree in combined_subgraph.degree() if degree >= min_degree]
291
+ combined_subgraph = combined_subgraph.subgraph(nodes_to_keep)
292
+
293
  subgraph = combined_subgraph
294
 
295
  # Check if number of nodes exceeds max_graph_nodes
lightrag/lightrag.py CHANGED
@@ -504,10 +504,24 @@ class LightRAG:
504
  return text
505
 
506
  async def get_knowledge_graph(
507
- self, node_label: str, max_depth: int
508
  ) -> KnowledgeGraph:
 
 
 
 
 
 
 
 
 
 
 
509
  return await self.chunk_entity_relation_graph.get_knowledge_graph(
510
- node_label=node_label, max_depth=max_depth
 
 
 
511
  )
512
 
513
  def _get_storage_class(self, storage_name: str) -> Callable[..., Any]:
 
504
  return text
505
 
506
  async def get_knowledge_graph(
507
+ self, node_label: str, max_depth: int, search_mode: str = "exact", min_degree: int = 0
508
  ) -> KnowledgeGraph:
509
+ """Get knowledge graph for a given label
510
+
511
+ Args:
512
+ node_label (str): Label to get knowledge graph for
513
+ max_depth (int): Maximum depth of graph
514
+ search_mode (str, optional): Search mode, either "exact" or "inclusive". Defaults to "exact".
515
+ min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0.
516
+
517
+ Returns:
518
+ KnowledgeGraph: Knowledge graph containing nodes and edges
519
+ """
520
  return await self.chunk_entity_relation_graph.get_knowledge_graph(
521
+ node_label=node_label,
522
+ max_depth=max_depth,
523
+ search_mode=search_mode,
524
+ min_degree=min_degree
525
  )
526
 
527
  def _get_storage_class(self, storage_name: str) -> Callable[..., Any]: