AgentEvalTool/backend/agenteval/web/routers/files.py

"""API routes for file management."""

from typing import NoReturn

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from sqlmodel import Session

from agenteval.config import get_settings
from agenteval.services.files import (
    CategoryNotFoundError,
    FileManagementError,
    FileManagementService,
    FileRecordNotFoundError,
    FileValidationError,
    PhysicalFileNotFoundError,
    UploadTooLargeError,
)
from agenteval.storage.db import FILES_DIR
from agenteval.storage.file_storage import FileStorage
from agenteval.web.deps import get_db
from agenteval.web.file_schemas import (
    CategoryCreate,
    CategoryUpdate,
    FileCategoryResponse,
    FileRecordResponse,
    FileUploadConfigResponse,
    OkResponse,
)

router = APIRouter()


def _allowed_extensions() -> set[str]:
    raw = get_settings().allowed_extensions
    return {extension.strip().lower() for extension in raw.split(",") if extension.strip()}


def get_file_service(session: Session = Depends(get_db)) -> FileManagementService:
    settings = get_settings()
    return FileManagementService(
        session=session,
        storage=FileStorage(FILES_DIR),
        allowed_extensions=_allowed_extensions(),
        max_upload_size_mb=settings.max_upload_size_mb,
    )


def _raise_http_error(error: FileManagementError) -> NoReturn:
    if isinstance(error, UploadTooLargeError):
        raise HTTPException(
            status_code=413,
            detail=f"文件大小超过限制 ({get_settings().max_upload_size_mb}MB)",
        ) from error
    if isinstance(error, FileValidationError):
        raise HTTPException(status_code=400, detail=str(error)) from error
    if isinstance(error, CategoryNotFoundError):
        raise HTTPException(status_code=404, detail="分类不存在") from error
    if isinstance(error, FileRecordNotFoundError):
        raise HTTPException(status_code=404, detail="文件不存在") from error
    if isinstance(error, PhysicalFileNotFoundError):
        raise HTTPException(status_code=404, detail="物理文件不存在") from error
    raise HTTPException(status_code=500, detail="文件操作失败") from error


@router.get("/config", response_model=FileUploadConfigResponse)
def get_upload_config() -> FileUploadConfigResponse:
    return FileUploadConfigResponse(
        allowed_extensions=sorted(_allowed_extensions()),
        max_upload_size_mb=get_settings().max_upload_size_mb,
    )


@router.get("/categories", response_model=list[FileCategoryResponse])
def list_categories(service: FileManagementService = Depends(get_file_service)):
    return service.list_categories()


@router.post("/categories", response_model=FileCategoryResponse)
def create_category(body: CategoryCreate, service: FileManagementService = Depends(get_file_service)):
    try:
        return FileCategoryResponse.from_record(service.create_category(body.name, body.parent_id))
    except FileManagementError as error:
        _raise_http_error(error)


@router.put("/categories/{category_id}", response_model=FileCategoryResponse)
def update_category(
    category_id: str,
    body: CategoryUpdate,
    service: FileManagementService = Depends(get_file_service),
):
    try:
        return FileCategoryResponse.from_record(service.update_category(category_id, body.name))
    except FileManagementError as error:
        _raise_http_error(error)


@router.delete("/categories/{category_id}", response_model=OkResponse)
def delete_category(category_id: str, service: FileManagementService = Depends(get_file_service)):
    try:
        service.delete_category(category_id)
        return OkResponse()
    except FileManagementError as error:
        _raise_http_error(error)


@router.get("", response_model=list[FileRecordResponse])
def list_files(category_id: str | None = None, service: FileManagementService = Depends(get_file_service)):
    try:
        return [FileRecordResponse.from_record(record) for record in service.list_files(category_id)]
    except FileManagementError as error:
        _raise_http_error(error)


@router.post("/upload", response_model=FileRecordResponse)
async def upload_file(
    file: UploadFile = File(...),
    category_id: str | None = Form(default=None),
    service: FileManagementService = Depends(get_file_service),
):
    try:
        record = await service.upload(file, category_id)
        return FileRecordResponse.from_record(record)
    except FileManagementError as error:
        _raise_http_error(error)


@router.get("/{file_id}/download")
def download_file(file_id: str, service: FileManagementService = Depends(get_file_service)):
    try:
        record, file_path = service.get_download(file_id)
        return FileResponse(
            path=str(file_path),
            filename=record.original_name,
            media_type=record.mime_type or "application/octet-stream",
        )
    except FileManagementError as error:
        _raise_http_error(error)


@router.delete("/{file_id}", response_model=OkResponse)
def delete_file(file_id: str, service: FileManagementService = Depends(get_file_service)):
    try:
        service.delete_file(file_id)
        return OkResponse()
    except FileManagementError as error:
        _raise_http_error(error)