Add dynamic parameter handling for storage

This commit is contained in:
yangdx
2025-03-06 14:29:48 +08:00
parent 252d1f57e0
commit 5b29e760f0
2 changed files with 17 additions and 8 deletions

View File

@@ -247,7 +247,6 @@ class NetworkXStorage(BaseGraphStorage):
3. Followed by nodes directly connected to the matching nodes 3. Followed by nodes directly connected to the matching nodes
4. Finally, the degree of the nodes 4. Finally, the degree of the nodes
Args: Args:
node_label: Label of the starting node node_label: Label of the starting node
max_depth: Maximum depth of the subgraph max_depth: Maximum depth of the subgraph

View File

@@ -515,18 +515,28 @@ class LightRAG:
Args: Args:
node_label (str): Label to get knowledge graph for node_label (str): Label to get knowledge graph for
max_depth (int): Maximum depth of graph 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. min_degree (int, optional): Minimum degree of nodes to include. Defaults to 0.
inclusive (bool, optional): Whether to use inclusive search mode. Defaults to False.
Returns: Returns:
KnowledgeGraph: Knowledge graph containing nodes and edges KnowledgeGraph: Knowledge graph containing nodes and edges
""" """
return await self.chunk_entity_relation_graph.get_knowledge_graph( # get params supported by get_knowledge_graph of specified storage
node_label=node_label, import inspect
max_depth=max_depth, storage_params = inspect.signature(self.chunk_entity_relation_graph.get_knowledge_graph).parameters
min_degree=min_degree,
inclusive=inclusive, kwargs = {
) 'node_label': node_label,
'max_depth': max_depth
}
if 'min_degree' in storage_params and min_degree > 0:
kwargs['min_degree'] = min_degree
if 'inclusive' in storage_params:
kwargs['inclusive'] = inclusive
return await self.chunk_entity_relation_graph.get_knowledge_graph(**kwargs)
def _get_storage_class(self, storage_name: str) -> Callable[..., Any]: def _get_storage_class(self, storage_name: str) -> Callable[..., Any]:
import_path = STORAGES[storage_name] import_path = STORAGES[storage_name]