mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-07-28 03:13:04 +00:00
89 lines
3.4 KiB
Python
89 lines
3.4 KiB
Python
"""
|
|
This module contains all graph-related routes for the LightRAG API.
|
|
"""
|
|
|
|
from typing import Optional, List, Dict, Any
|
|
from fastapi import APIRouter, Depends, Query
|
|
from pydantic import BaseModel, Field
|
|
|
|
from ..utils_api import get_combined_auth_dependency
|
|
|
|
router = APIRouter(tags=["graph"])
|
|
|
|
# Pydantic models for graph routes
|
|
class GraphLabelsResponse(BaseModel):
|
|
"""Response model: List of graph labels"""
|
|
labels: List[str] = Field(description="List of graph labels")
|
|
|
|
class KnowledgeGraphNode(BaseModel):
|
|
"""Model for a node in the knowledge graph"""
|
|
id: str = Field(description="Unique identifier of the node")
|
|
label: str = Field(description="Label of the node")
|
|
properties: Dict[str, Any] = Field(default_factory=dict, description="Properties of the node")
|
|
|
|
class KnowledgeGraphEdge(BaseModel):
|
|
"""Model for an edge in the knowledge graph"""
|
|
source: str = Field(description="Source node ID")
|
|
target: str = Field(description="Target node ID")
|
|
type: str = Field(description="Type of the relationship")
|
|
properties: Dict[str, Any] = Field(default_factory=dict, description="Properties of the edge")
|
|
|
|
class KnowledgeGraphResponse(BaseModel):
|
|
"""Response model: Knowledge graph data"""
|
|
nodes: List[KnowledgeGraphNode] = Field(description="List of nodes in the graph")
|
|
edges: List[KnowledgeGraphEdge] = Field(description="List of edges in the graph")
|
|
|
|
|
|
def create_graph_routes(rag, api_key: Optional[str] = None):
|
|
combined_auth = get_combined_auth_dependency(api_key)
|
|
|
|
@router.get("/graph/label/list",
|
|
dependencies=[Depends(combined_auth)],
|
|
response_model=GraphLabelsResponse)
|
|
async def get_graph_labels():
|
|
"""
|
|
Get all graph labels
|
|
|
|
Returns:
|
|
GraphLabelsResponse: List of graph labels
|
|
"""
|
|
labels = await rag.get_graph_labels()
|
|
return GraphLabelsResponse(labels=labels)
|
|
|
|
@router.get("/graphs",
|
|
dependencies=[Depends(combined_auth)],
|
|
response_model=KnowledgeGraphResponse)
|
|
async def get_knowledge_graph(
|
|
label: str = Query(..., description="Label to get knowledge graph for"),
|
|
max_depth: int = Query(3, description="Maximum depth of graph", ge=1),
|
|
max_nodes: int = Query(1000, description="Maxiumu nodes to return", ge=1),
|
|
):
|
|
"""
|
|
Retrieve a connected subgraph of nodes where the label includes the specified label.
|
|
When reducing the number of nodes, the prioritization criteria are as follows:
|
|
1. Hops(path) to the staring node take precedence
|
|
2. Followed by the degree of the nodes
|
|
|
|
Args:
|
|
label (str): Label of the starting node
|
|
max_depth (int, optional): Maximum depth of the subgraph,Defaults to 3
|
|
max_nodes: Maxiumu nodes to return
|
|
|
|
Returns:
|
|
KnowledgeGraphResponse: Knowledge graph containing nodes and edges
|
|
"""
|
|
graph_data = await rag.get_knowledge_graph(
|
|
node_label=label,
|
|
max_depth=max_depth,
|
|
max_nodes=max_nodes,
|
|
)
|
|
|
|
# Convert the returned dictionary to our response model format
|
|
# Assuming the returned dictionary has 'nodes' and 'edges' keys
|
|
return KnowledgeGraphResponse(
|
|
nodes=graph_data.get("nodes", []),
|
|
edges=graph_data.get("edges", [])
|
|
)
|
|
|
|
return router
|