LightRAG/lightrag/api/routers/tenant_routes.py

"""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 uuid import uuid4

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