Added minimum degree filter for graph queries
- Introduced min_degree parameter in graph query - Updated UI to include minimum degree setting - Modified API to handle min_degree parameter - Updated graph query logic in LightRAG
This commit is contained in:
@@ -5,7 +5,6 @@ This module contains all graph-related routes for the LightRAG API.
|
||||
from typing import Optional
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from ...utils import logger
|
||||
from ..utils_api import get_api_key_dependency
|
||||
|
||||
router = APIRouter(tags=["graph"])
|
||||
@@ -25,7 +24,9 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
|
||||
return await rag.get_graph_labels()
|
||||
|
||||
@router.get("/graphs", dependencies=[Depends(optional_api_key)])
|
||||
async def get_knowledge_graph(label: str, max_depth: int = 3, inclusive: bool = False, min_degree: int = 0):
|
||||
async def get_knowledge_graph(
|
||||
label: str, max_depth: int = 3, min_degree: int = 0, inclusive: bool = False
|
||||
):
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the label includes the specified label.
|
||||
Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000).
|
||||
@@ -44,7 +45,11 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
|
||||
Returns:
|
||||
Dict[str, List[str]]: Knowledge graph for label
|
||||
"""
|
||||
logger.info(f"Inclusive search : {inclusive}, Min degree: {min_degree}, Label: {label}")
|
||||
return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth, inclusive=inclusive, min_degree=min_degree)
|
||||
return await rag.get_knowledge_graph(
|
||||
node_label=label,
|
||||
max_depth=max_depth,
|
||||
inclusive=inclusive,
|
||||
min_degree=min_degree,
|
||||
)
|
||||
|
||||
return router
|
||||
|
@@ -232,7 +232,11 @@ class NetworkXStorage(BaseGraphStorage):
|
||||
return sorted(list(labels))
|
||||
|
||||
async def get_knowledge_graph(
|
||||
self, node_label: str, max_depth: int = 5, search_mode: str = "exact", min_degree: int = 0
|
||||
self,
|
||||
node_label: str,
|
||||
max_depth: int = 5,
|
||||
min_degree: int = 0,
|
||||
inclusive: bool = False,
|
||||
) -> KnowledgeGraph:
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
|
||||
@@ -268,7 +272,7 @@ class NetworkXStorage(BaseGraphStorage):
|
||||
nodes_to_explore = []
|
||||
for n, attr in graph.nodes(data=True):
|
||||
node_str = str(n)
|
||||
if search_mode == "exact":
|
||||
if not inclusive:
|
||||
if node_label == node_str: # Use exact matching
|
||||
nodes_to_explore.append(n)
|
||||
else: # inclusive mode
|
||||
@@ -284,12 +288,16 @@ 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]
|
||||
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
|
||||
|
@@ -504,7 +504,11 @@ class LightRAG:
|
||||
return text
|
||||
|
||||
async def get_knowledge_graph(
|
||||
self, node_label: str, max_depth: int, inclusive: bool = False, min_degree: int = 0
|
||||
self,
|
||||
node_label: str,
|
||||
max_depth: int,
|
||||
min_degree: int = 0,
|
||||
inclusive: bool = False,
|
||||
) -> KnowledgeGraph:
|
||||
"""Get knowledge graph for a given label
|
||||
|
||||
@@ -520,6 +524,8 @@ class LightRAG:
|
||||
return await self.chunk_entity_relation_graph.get_knowledge_graph(
|
||||
node_label=node_label,
|
||||
max_depth=max_depth,
|
||||
min_degree=min_degree,
|
||||
inclusive=inclusive,
|
||||
)
|
||||
|
||||
def _get_storage_class(self, storage_name: str) -> Callable[..., Any]:
|
||||
|
Reference in New Issue
Block a user