ai-robot-core/ai-service/app/api/admin/kb.py

"""
Knowledge Base management endpoints.
[AC-ASA-01, AC-ASA-02, AC-ASA-08] Document upload, list, and index job status.
"""

import logging
from typing import Annotated, Any, Optional

from fastapi import APIRouter, Depends, Header, Query, UploadFile, File, Form
from fastapi.responses import JSONResponse

from app.core.tenant import get_tenant_id
from app.models import ErrorResponse

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/admin/kb", tags=["KB Management"])


@router.get(
    "/documents",
    operation_id="listDocuments",
    summary="Query document list",
    description="[AC-ASA-08] Get list of documents with pagination and filtering.",
    responses={
        200: {"description": "Document list with pagination"},
        401: {"description": "Unauthorized", "model": ErrorResponse},
        403: {"description": "Forbidden", "model": ErrorResponse},
    },
)
async def list_documents(
    tenant_id: Annotated[str, Depends(get_tenant_id)],
    kb_id: Annotated[Optional[str], Query()] = None,
    status: Annotated[Optional[str], Query()] = None,
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
) -> JSONResponse:
    """
    [AC-ASA-08] List documents with filtering and pagination.
    """
    logger.info(
        f"[AC-ASA-08] Listing documents: tenant={tenant_id}, kb_id={kb_id}, "
        f"status={status}, page={page}, page_size={page_size}"
    )

    mock_documents = [
        {
            "docId": "doc_001",
            "kbId": kb_id or "kb_default",
            "fileName": "product_manual.pdf",
            "status": "completed",
            "createdAt": "2026-02-20T10:00:00Z",
            "updatedAt": "2026-02-20T10:30:00Z",
        },
        {
            "docId": "doc_002",
            "kbId": kb_id or "kb_default",
            "fileName": "faq.docx",
            "status": "processing",
            "createdAt": "2026-02-21T14:00:00Z",
            "updatedAt": "2026-02-21T14:15:00Z",
        },
        {
            "docId": "doc_003",
            "kbId": kb_id or "kb_default",
            "fileName": "invalid_file.txt",
            "status": "failed",
            "createdAt": "2026-02-22T09:00:00Z",
            "updatedAt": "2026-02-22T09:05:00Z",
        },
    ]

    filtered = mock_documents
    if kb_id:
        filtered = [d for d in filtered if d["kbId"] == kb_id]
    if status:
        filtered = [d for d in filtered if d["status"] == status]

    total = len(filtered)
    total_pages = (total + page_size - 1) // page_size

    return JSONResponse(
        content={
            "data": filtered,
            "pagination": {
                "page": page,
                "pageSize": page_size,
                "total": total,
                "totalPages": total_pages,
            },
        }
    )


@router.post(
    "/documents",
    operation_id="uploadDocument",
    summary="Upload/import document",
    description="[AC-ASA-01] Upload document and trigger indexing job.",
    responses={
        202: {"description": "Accepted - async indexing job started"},
        401: {"description": "Unauthorized", "model": ErrorResponse},
        403: {"description": "Forbidden", "model": ErrorResponse},
    },
)
async def upload_document(
    tenant_id: Annotated[str, Depends(get_tenant_id)],
    file: UploadFile = File(...),
    kb_id: str = Form(...),
) -> JSONResponse:
    """
    [AC-ASA-01] Upload document and create indexing job.
    """
    logger.info(
        f"[AC-ASA-01] Uploading document: tenant={tenant_id}, "
        f"kb_id={kb_id}, filename={file.filename}"
    )

    import uuid

    job_id = f"job_{uuid.uuid4().hex[:8]}"

    return JSONResponse(
        status_code=202,
        content={
            "jobId": job_id,
            "status": "pending",
        },
    )


@router.get(
    "/index/jobs/{job_id}",
    operation_id="getIndexJob",
    summary="Query index job status",
    description="[AC-ASA-02] Get indexing job status and progress.",
    responses={
        200: {"description": "Job status details"},
        401: {"description": "Unauthorized", "model": ErrorResponse},
        403: {"description": "Forbidden", "model": ErrorResponse},
    },
)
async def get_index_job(
    tenant_id: Annotated[str, Depends(get_tenant_id)],
    job_id: str,
) -> JSONResponse:
    """
    [AC-ASA-02] Get indexing job status with progress.
    """
    logger.info(
        f"[AC-ASA-02] Getting job status: tenant={tenant_id}, job_id={job_id}"
    )

    mock_job_statuses = {
        "job_pending": {
            "jobId": job_id,
            "status": "pending",
            "progress": 0,
            "errorMsg": None,
        },
        "job_processing": {
            "jobId": job_id,
            "status": "processing",
            "progress": 45,
            "errorMsg": None,
        },
        "job_completed": {
            "jobId": job_id,
            "status": "completed",
            "progress": 100,
            "errorMsg": None,
        },
        "job_failed": {
            "jobId": job_id,
            "status": "failed",
            "progress": 30,
            "errorMsg": "Failed to parse PDF: Invalid format",
        },
    }

    job_status = mock_job_statuses.get(job_id, mock_job_statuses["job_processing"])

    return JSONResponse(content=job_status)