ax-backend/app/core/exceptions.py

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse


class DocumentNotFoundError(Exception):
    def __init__(self, document_id: str):
        self.document_id = document_id
        super().__init__(f"Document '{document_id}' not found")


class ContentTooLargeError(Exception):
    def __init__(self, size: int, limit: int = 200_000):
        super().__init__(f"Content size {size} bytes exceeds limit of {limit} bytes")


class ExportError(Exception):
    def __init__(self, reason: str):
        super().__init__(f"Export failed: {reason}")


class RecordNotFoundError(Exception):
    def __init__(self, record_id: str):
        self.record_id = record_id
        super().__init__(f"Export record '{record_id}' not found")


def register_exception_handlers(app: FastAPI) -> None:
    @app.exception_handler(DocumentNotFoundError)
    async def document_not_found_handler(
        request: Request, exc: DocumentNotFoundError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=404,
            content={"code": 404, "message": str(exc), "data": None},
        )

    @app.exception_handler(ContentTooLargeError)
    async def content_too_large_handler(
        request: Request, exc: ContentTooLargeError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=413,
            content={"code": 413, "message": str(exc), "data": None},
        )

    @app.exception_handler(ExportError)
    async def export_error_handler(
        request: Request, exc: ExportError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=500,
            content={"code": 500, "message": str(exc), "data": None},
        )

    @app.exception_handler(RecordNotFoundError)
    async def record_not_found_handler(
        request: Request, exc: RecordNotFoundError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=404,
            content={"code": 4041, "message": str(exc), "data": None},
        )