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 +37 -1
- lightrag/kg/networkx_impl.py +17 -4
- lightrag/lightrag.py +16 -2
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 266 |
nodes_to_explore = []
|
| 267 |
for n, attr in graph.nodes(data=True):
|
| 268 |
-
|
| 269 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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,
|
|
|
|
|
|
|
|
|
|
| 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]:
|