yangdx commited on
Commit
f18b719
·
1 Parent(s): 3564bbb

Add dynamic parameter handling for storage

Browse files
lightrag/kg/networkx_impl.py CHANGED
@@ -247,7 +247,6 @@ class NetworkXStorage(BaseGraphStorage):
247
  3. Followed by nodes directly connected to the matching nodes
248
  4. Finally, the degree of the nodes
249
 
250
-
251
  Args:
252
  node_label: Label of the starting node
253
  max_depth: Maximum depth of the subgraph
 
247
  3. Followed by nodes directly connected to the matching nodes
248
  4. Finally, the degree of the nodes
249
 
 
250
  Args:
251
  node_label: Label of the starting node
252
  max_depth: Maximum depth of the subgraph
lightrag/lightrag.py CHANGED
@@ -515,18 +515,28 @@ class LightRAG:
515
  Args:
516
  node_label (str): Label to get knowledge graph for
517
  max_depth (int): Maximum depth of graph
518
- search_mode (str, optional): Search mode, either "exact" or "inclusive". Defaults to "exact".
519
  min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0.
 
520
 
521
  Returns:
522
  KnowledgeGraph: Knowledge graph containing nodes and edges
523
  """
524
- return await self.chunk_entity_relation_graph.get_knowledge_graph(
525
- node_label=node_label,
526
- max_depth=max_depth,
527
- min_degree=min_degree,
528
- inclusive=inclusive,
529
- )
 
 
 
 
 
 
 
 
 
 
530
 
531
  def _get_storage_class(self, storage_name: str) -> Callable[..., Any]:
532
  import_path = STORAGES[storage_name]
 
515
  Args:
516
  node_label (str): Label to get knowledge graph for
517
  max_depth (int): Maximum depth of graph
 
518
  min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0.
519
+ inclusive (bool, optional): Whether to use inclusive search mode. Defaults to False.
520
 
521
  Returns:
522
  KnowledgeGraph: Knowledge graph containing nodes and edges
523
  """
524
+ # get params supported by get_knowledge_graph of specified storage
525
+ import inspect
526
+ storage_params = inspect.signature(self.chunk_entity_relation_graph.get_knowledge_graph).parameters
527
+
528
+ kwargs = {
529
+ 'node_label': node_label,
530
+ 'max_depth': max_depth
531
+ }
532
+
533
+ if 'min_degree' in storage_params and min_degree > 0:
534
+ kwargs['min_degree'] = min_degree
535
+
536
+ if 'inclusive' in storage_params:
537
+ kwargs['inclusive'] = inclusive
538
+
539
+ return await self.chunk_entity_relation_graph.get_knowledge_graph(**kwargs)
540
 
541
  def _get_storage_class(self, storage_name: str) -> Callable[..., Any]:
542
  import_path = STORAGES[storage_name]