590 lines
20 KiB
Python
590 lines
20 KiB
Python
"""Tenant management routes for multi-tenant LightRAG API.
|
|
|
|
Provides CRUD endpoints for managing tenants and knowledge bases.
|
|
"""
|
|
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, status, Request, Header
|
|
from pydantic import BaseModel
|
|
|
|
from lightrag.models.tenant import TenantContext, Permission
|
|
from lightrag.services.tenant_service import TenantService
|
|
from lightrag.api.dependencies import (
|
|
get_tenant_context,
|
|
check_permission,
|
|
get_admin_context,
|
|
get_tenant_context_no_kb,
|
|
resolve_default_tenant,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Request/Response Models
|
|
class TenantCreateRequest(BaseModel):
|
|
name: str
|
|
description: Optional[str] = ""
|
|
metadata: Optional[dict] = None
|
|
|
|
|
|
class TenantUpdateRequest(BaseModel):
|
|
name: Optional[str] = None
|
|
description: Optional[str] = None
|
|
metadata: Optional[dict] = None
|
|
|
|
|
|
class TenantResponse(BaseModel):
|
|
tenant_id: str
|
|
name: str
|
|
description: str
|
|
created_at: str
|
|
updated_at: str
|
|
num_knowledge_bases: int
|
|
num_documents: int
|
|
storage_used_gb: float
|
|
|
|
|
|
class KBCreateRequest(BaseModel):
|
|
name: str
|
|
description: Optional[str] = ""
|
|
metadata: Optional[dict] = None
|
|
|
|
|
|
class KBUpdateRequest(BaseModel):
|
|
name: Optional[str] = None
|
|
description: Optional[str] = None
|
|
metadata: Optional[dict] = None
|
|
|
|
|
|
class KBResponse(BaseModel):
|
|
kb_id: str
|
|
tenant_id: str
|
|
name: str
|
|
description: str
|
|
created_at: str
|
|
updated_at: str
|
|
num_documents: int
|
|
num_entities: int
|
|
num_relations: int
|
|
|
|
|
|
# Pagination Models
|
|
class PaginatedKBResponse(BaseModel):
|
|
"""Paginated response for knowledge bases."""
|
|
|
|
items: List[KBResponse]
|
|
total: int
|
|
page: int
|
|
page_size: int
|
|
total_pages: int
|
|
has_next: bool
|
|
has_prev: bool
|
|
|
|
|
|
class PaginatedTenantResponse(BaseModel):
|
|
"""Paginated response for tenants."""
|
|
|
|
items: List[TenantResponse]
|
|
total: int
|
|
page: int
|
|
page_size: int
|
|
total_pages: int
|
|
has_next: bool
|
|
has_prev: bool
|
|
|
|
|
|
def create_tenant_routes(tenant_service: TenantService) -> APIRouter:
|
|
"""Create tenant management routes.
|
|
|
|
Args:
|
|
tenant_service: Service instance for tenant operations
|
|
|
|
Returns:
|
|
APIRouter with tenant routes
|
|
"""
|
|
router = APIRouter(prefix="/api/v1", tags=["tenants"])
|
|
|
|
# Tenant management endpoints
|
|
|
|
@router.get("/tenants", response_model=PaginatedTenantResponse)
|
|
async def list_tenants(
|
|
page: int = 1,
|
|
page_size: int = 10,
|
|
search: Optional[str] = None,
|
|
authorization: Optional[str] = Header(None),
|
|
):
|
|
"""List all available tenants with pagination.
|
|
|
|
Useful for tenant selection. This endpoint is public to allow
|
|
unauthenticated access for tenant selection on the login page.
|
|
"""
|
|
# Note: This endpoint is intentionally public to support tenant selection
|
|
# on the login page before authentication. In production, you may want to
|
|
# restrict this to specific IPs or use rate limiting.
|
|
|
|
try:
|
|
# Validate pagination parameters
|
|
page = max(1, page)
|
|
page_size = min(max(1, page_size), 100) # Max 100 per page
|
|
|
|
# Get tenants from service
|
|
tenants_data = await tenant_service.list_tenants(
|
|
skip=(page - 1) * page_size,
|
|
limit=page_size,
|
|
search=search,
|
|
tenant_id_filter=None,
|
|
)
|
|
|
|
total_count = tenants_data.get("total", 0)
|
|
tenants_list = tenants_data.get("items", [])
|
|
|
|
# Convert to response models
|
|
items = [
|
|
TenantResponse(
|
|
tenant_id=t.tenant_id,
|
|
name=t.tenant_name,
|
|
description=t.description,
|
|
created_at=t.created_at.isoformat(),
|
|
updated_at=t.updated_at.isoformat(),
|
|
num_knowledge_bases=t.kb_count,
|
|
num_documents=t.total_documents,
|
|
storage_used_gb=t.total_storage_mb / 1024.0,
|
|
)
|
|
for t in tenants_list
|
|
]
|
|
|
|
# Calculate pagination metadata
|
|
total_pages = (total_count + page_size - 1) // page_size
|
|
|
|
return PaginatedTenantResponse(
|
|
items=items,
|
|
total=total_count,
|
|
page=page,
|
|
page_size=page_size,
|
|
total_pages=total_pages,
|
|
has_next=page < total_pages,
|
|
has_prev=page > 1,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error listing tenants: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to list tenants",
|
|
)
|
|
|
|
@router.get("/tenants/me", response_model=TenantResponse)
|
|
async def get_current_tenant(
|
|
context: TenantContext = Depends(get_tenant_context_no_kb),
|
|
):
|
|
"""Get current tenant details based on context.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
"""
|
|
try:
|
|
tenant = await tenant_service.get_tenant(context.tenant_id)
|
|
if not tenant:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND, detail="Tenant not found"
|
|
)
|
|
|
|
return TenantResponse(
|
|
tenant_id=tenant.tenant_id,
|
|
name=tenant.tenant_name,
|
|
description=tenant.description,
|
|
created_at=tenant.created_at.isoformat(),
|
|
updated_at=tenant.updated_at.isoformat(),
|
|
num_knowledge_bases=tenant.kb_count,
|
|
num_documents=tenant.total_documents,
|
|
storage_used_gb=tenant.total_storage_mb / 1024.0,
|
|
)
|
|
except HTTPException:
|
|
# Let explicit HTTPExceptions (404 etc.) bubble up unchanged
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error getting current tenant {context.tenant_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to get tenant details",
|
|
)
|
|
|
|
@router.post(
|
|
"/tenants", response_model=TenantResponse, status_code=status.HTTP_201_CREATED
|
|
)
|
|
async def create_tenant(
|
|
request: TenantCreateRequest, admin_context: dict = Depends(get_admin_context)
|
|
):
|
|
"""Create a new tenant.
|
|
|
|
Requires admin authentication (no tenant_id required).
|
|
This allows creating the initial tenant(s) for new organizations.
|
|
"""
|
|
try:
|
|
username = admin_context.get("username")
|
|
|
|
tenant = await tenant_service.create_tenant(
|
|
tenant_name=request.name,
|
|
description=request.description or "",
|
|
created_by=username,
|
|
)
|
|
|
|
# Add creator as owner
|
|
if username:
|
|
try:
|
|
await tenant_service.add_user_to_tenant(
|
|
user_id=username,
|
|
tenant_id=tenant.tenant_id,
|
|
role="owner",
|
|
created_by=username,
|
|
)
|
|
logger.info(
|
|
f"Added user {username} as owner of new tenant {tenant.tenant_id}"
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to add creator as owner: {e}")
|
|
# Continue anyway, as tenant was created
|
|
|
|
return TenantResponse(
|
|
tenant_id=tenant.tenant_id,
|
|
name=tenant.tenant_name,
|
|
description=tenant.description,
|
|
created_at=tenant.created_at.isoformat(),
|
|
updated_at=tenant.updated_at.isoformat(),
|
|
num_knowledge_bases=tenant.kb_count,
|
|
num_documents=tenant.total_documents,
|
|
storage_used_gb=tenant.total_storage_mb / 1024.0,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error creating tenant: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to create tenant",
|
|
)
|
|
|
|
# Knowledge base management endpoints
|
|
|
|
@router.post(
|
|
"/knowledge-bases",
|
|
response_model=KBResponse,
|
|
status_code=status.HTTP_201_CREATED,
|
|
)
|
|
async def create_knowledge_base_context(
|
|
request: KBCreateRequest,
|
|
context: TenantContext = Depends(get_tenant_context_no_kb),
|
|
):
|
|
"""Create a new knowledge base within the current tenant.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
"""
|
|
try:
|
|
kb = await tenant_service.create_knowledge_base(
|
|
tenant_id=context.tenant_id,
|
|
kb_name=request.name,
|
|
description=request.description or "",
|
|
)
|
|
|
|
return KBResponse(
|
|
kb_id=kb.kb_id,
|
|
tenant_id=kb.tenant_id,
|
|
name=kb.kb_name,
|
|
description=kb.description,
|
|
created_at=kb.created_at.isoformat(),
|
|
updated_at=kb.updated_at.isoformat(),
|
|
num_documents=kb.document_count,
|
|
num_entities=kb.entity_count,
|
|
num_relations=kb.relationship_count,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error creating KB for tenant {context.tenant_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to create knowledge base",
|
|
)
|
|
|
|
@router.get("/knowledge-bases", response_model=PaginatedKBResponse)
|
|
async def list_knowledge_bases_context(
|
|
page: int = 1,
|
|
page_size: int = 10,
|
|
search: Optional[str] = None,
|
|
context: TenantContext = Depends(get_tenant_context_no_kb),
|
|
):
|
|
"""List all knowledge bases in the current tenant with pagination.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
|
|
Query Parameters:
|
|
page: Page number (1-indexed), defaults to 1
|
|
page_size: Number of items per page, defaults to 10 (max 100)
|
|
search: Optional search string to filter KBs by name or description
|
|
"""
|
|
try:
|
|
# Validate pagination parameters
|
|
page = max(1, page)
|
|
page_size = min(max(1, page_size), 100) # Max 100 per page
|
|
|
|
# Get KBs from service
|
|
kbs_data = await tenant_service.list_knowledge_bases(
|
|
tenant_id=context.tenant_id,
|
|
skip=(page - 1) * page_size,
|
|
limit=page_size,
|
|
search=search,
|
|
)
|
|
|
|
total_count = kbs_data.get("total", 0)
|
|
kbs_list = kbs_data.get("items", [])
|
|
|
|
# Convert to response models
|
|
items = [
|
|
KBResponse(
|
|
kb_id=kb.kb_id,
|
|
tenant_id=kb.tenant_id,
|
|
name=kb.kb_name,
|
|
description=kb.description,
|
|
created_at=kb.created_at.isoformat(),
|
|
updated_at=kb.updated_at.isoformat(),
|
|
num_documents=kb.document_count,
|
|
num_entities=kb.entity_count,
|
|
num_relations=kb.relationship_count,
|
|
)
|
|
for kb in kbs_list
|
|
]
|
|
|
|
# Calculate pagination metadata
|
|
total_pages = (total_count + page_size - 1) // page_size
|
|
|
|
return PaginatedKBResponse(
|
|
items=items,
|
|
total=total_count,
|
|
page=page,
|
|
page_size=page_size,
|
|
total_pages=total_pages,
|
|
has_next=page < total_pages,
|
|
has_prev=page > 1,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error listing KBs for tenant {context.tenant_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to list knowledge bases",
|
|
)
|
|
|
|
@router.get("/knowledge-bases/{kb_id}", response_model=KBResponse)
|
|
async def get_knowledge_base_context(
|
|
kb_id: str, context: TenantContext = Depends(get_tenant_context)
|
|
):
|
|
"""Get knowledge base details.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
"""
|
|
# Note: get_tenant_context already validates that context.kb_id matches kb_id if kb_id is in context
|
|
# But here kb_id is a path param, so we should double check if context has a specific kb_id
|
|
if context.kb_id and context.kb_id != kb_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="Cannot access other knowledge bases",
|
|
)
|
|
|
|
try:
|
|
kb = await tenant_service.get_knowledge_base(context.tenant_id, kb_id)
|
|
if not kb or kb.tenant_id != context.tenant_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Knowledge base not found",
|
|
)
|
|
|
|
return KBResponse(
|
|
kb_id=kb.kb_id,
|
|
tenant_id=kb.tenant_id,
|
|
name=kb.kb_name,
|
|
description=kb.description,
|
|
created_at=kb.created_at.isoformat(),
|
|
updated_at=kb.updated_at.isoformat(),
|
|
num_documents=kb.document_count,
|
|
num_entities=kb.entity_count,
|
|
num_relations=kb.relationship_count,
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error getting KB {kb_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to get knowledge base",
|
|
)
|
|
|
|
@router.put("/knowledge-bases/{kb_id}", response_model=KBResponse)
|
|
async def update_knowledge_base_context(
|
|
kb_id: str,
|
|
request: KBUpdateRequest,
|
|
context: TenantContext = Depends(check_permission(Permission.MANAGE_KB.value)),
|
|
):
|
|
"""Update knowledge base settings.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
"""
|
|
if context.kb_id and context.kb_id != kb_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="Cannot update other knowledge bases",
|
|
)
|
|
|
|
try:
|
|
kb = await tenant_service.update_knowledge_base(
|
|
tenant_id=context.tenant_id,
|
|
kb_id=kb_id,
|
|
name=request.name,
|
|
description=request.description,
|
|
metadata=request.metadata,
|
|
)
|
|
|
|
if not kb or kb.tenant_id != context.tenant_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Knowledge base not found",
|
|
)
|
|
|
|
return KBResponse(
|
|
kb_id=kb.kb_id,
|
|
tenant_id=kb.tenant_id,
|
|
name=kb.kb_name,
|
|
description=kb.description,
|
|
created_at=kb.created_at.isoformat(),
|
|
updated_at=kb.updated_at.isoformat(),
|
|
num_documents=kb.document_count,
|
|
num_entities=kb.entity_count,
|
|
num_relations=kb.relationship_count,
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error updating KB {kb_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to update knowledge base",
|
|
)
|
|
|
|
@router.delete("/knowledge-bases/{kb_id}", status_code=status.HTTP_204_NO_CONTENT)
|
|
async def delete_knowledge_base_context(
|
|
kb_id: str,
|
|
context: TenantContext = Depends(check_permission(Permission.DELETE_KB.value)),
|
|
):
|
|
"""Delete a knowledge base.
|
|
|
|
The tenant is identified by the X-Tenant-ID header or authentication token.
|
|
"""
|
|
if context.kb_id and context.kb_id != kb_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="Cannot delete other knowledge bases",
|
|
)
|
|
|
|
try:
|
|
success = await tenant_service.delete_knowledge_base(
|
|
context.tenant_id, kb_id
|
|
)
|
|
if not success:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Knowledge base not found",
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error deleting KB {kb_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to delete knowledge base",
|
|
)
|
|
|
|
@router.get(
|
|
"/tenants/{tenant_id}/knowledge-bases", response_model=PaginatedKBResponse
|
|
)
|
|
async def list_knowledge_bases_path(
|
|
tenant_id: str,
|
|
request: Request,
|
|
page: int = 1,
|
|
page_size: int = 10,
|
|
search: Optional[str] = None,
|
|
authorization: Optional[str] = Header(None),
|
|
):
|
|
"""List all knowledge bases for a specific tenant (path param).
|
|
|
|
Legacy endpoint support.
|
|
"""
|
|
if not authorization:
|
|
raise HTTPException(status_code=401, detail="Missing authorization header")
|
|
|
|
try:
|
|
# Extract username from token
|
|
from lightrag.api.auth import auth_handler
|
|
|
|
try:
|
|
print(f"DEBUG: Validating token: {authorization[:30]}...")
|
|
logger.info(f"Validating token: {authorization[:30]}...")
|
|
scheme, token = authorization.split()
|
|
token_data = auth_handler.validate_token(token)
|
|
username = token_data.get("username")
|
|
print(f"DEBUG: Token valid for user: {username}")
|
|
logger.info(f"Token valid for user: {username}")
|
|
except Exception as e:
|
|
print(f"DEBUG: Token validation failed: {e}")
|
|
logger.error(f"Token validation failed: {e}")
|
|
username = None
|
|
|
|
# Resolve default tenant
|
|
resolved_tenant_id = await resolve_default_tenant(
|
|
request, tenant_id, user_id=username
|
|
)
|
|
|
|
# Validate pagination parameters
|
|
page = max(1, page)
|
|
page_size = min(max(1, page_size), 100) # Max 100 per page
|
|
|
|
# Get KBs from service
|
|
kbs_data = await tenant_service.list_knowledge_bases(
|
|
tenant_id=resolved_tenant_id,
|
|
skip=(page - 1) * page_size,
|
|
limit=page_size,
|
|
search=search,
|
|
)
|
|
|
|
total_count = kbs_data.get("total", 0)
|
|
kbs_list = kbs_data.get("items", [])
|
|
|
|
# Convert to response models
|
|
items = [
|
|
KBResponse(
|
|
kb_id=kb.kb_id,
|
|
tenant_id=kb.tenant_id,
|
|
name=kb.kb_name,
|
|
description=kb.description,
|
|
created_at=kb.created_at.isoformat(),
|
|
updated_at=kb.updated_at.isoformat(),
|
|
num_documents=kb.document_count,
|
|
num_entities=kb.entity_count,
|
|
num_relations=kb.relationship_count,
|
|
)
|
|
for kb in kbs_list
|
|
]
|
|
|
|
# Calculate pagination metadata
|
|
total_pages = (total_count + page_size - 1) // page_size
|
|
|
|
return PaginatedKBResponse(
|
|
items=items,
|
|
total=total_count,
|
|
page=page,
|
|
page_size=page_size,
|
|
total_pages=total_pages,
|
|
has_next=page < total_pages,
|
|
has_prev=page > 1,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error listing KBs for tenant {tenant_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="Failed to list knowledge bases",
|
|
)
|
|
|
|
return router
|