app.py

"""AI Gateway FastAPI application."""

from __future__ import annotations

import logging
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass
from uuid import uuid4

import gradio as gr
from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastmcp.utilities.lifespan import combine_lifespans
from gradio.context import LocalContext

from config import Settings, get_settings
from core.errors import GatewayError
from core.loader import ModelLoader
from core.manager import AIService, TaskManager
from core.runtime import initialize_zerogpu
from core.workflow import WorkflowService
from gateway_mcp.auth import AuthenticationError, authenticate_headers
from gateway_mcp.server import MCPServerBundle, create_mcp_server
from gradio_ui import create_gradio_ui
from routes import audio, health, image, video, workflow
from routes.schemas import ErrorResponse
from utils.files import OutputManager
from utils.logger import configure_logging, request_id_context


logger = logging.getLogger(__name__)


@dataclass(slots=True)
class GatewayServices:
    """Process-wide services made available to route dependencies."""

    settings: Settings
    outputs: OutputManager
    loader: ModelLoader
    tasks: TaskManager
    workflows: WorkflowService
    ai: AIService
    mcp: MCPServerBundle


def create_app(settings: Settings | None = None) -> FastAPI:
    """Create an application instance, allowing isolated test settings."""
    resolved_settings = settings or get_settings()
    configure_logging(resolved_settings.log_level)
    outputs = OutputManager(resolved_settings)
    outputs.initialize()
    loader = ModelLoader(resolved_settings)
    tasks = TaskManager(resolved_settings, loader)
    workflows = WorkflowService(resolved_settings, tasks, outputs)
    ai = AIService(resolved_settings, tasks, loader, outputs, workflows)
    mcp = create_mcp_server(ai)
    gradio_ui = create_gradio_ui(ai)
    services = GatewayServices(
        resolved_settings, outputs, loader, tasks, workflows, ai, mcp
    )

    @asynccontextmanager
    async def lifespan(application: FastAPI):
        application.state.gateway = services
        removed = outputs.cleanup_stale_tmp()
        logger.info("gateway starting", extra={"stale_tmp_removed": removed})
        initialize_zerogpu()
        await tasks.start()
        try:
            yield
        finally:
            await tasks.stop()
            loader.close()
            logger.info("gateway stopped")

    application = FastAPI(
        title=resolved_settings.app_name,
        version=resolved_settings.app_version,
        description="A serialized, multi-model media inference gateway.",
        lifespan=combine_lifespans(lifespan, mcp.app.lifespan),
    )
    application.state.gateway = services
    application.state.gateway_mcp = mcp

    @application.middleware("http")
    async def request_context(request: Request, call_next):
        request_id = uuid4().hex
        request.state.request_id = request_id
        token = request_id_context.set(request_id)
        gradio_token = LocalContext.request.set(gr.Request(request))
        started = time.perf_counter()
        try:
            public_paths = {"/", "/health", "/docs", "/openapi.json", "/redoc"}
            gradio_path = request.url.path == "/ui" or request.url.path.startswith(
                "/ui/"
            )
            if request.url.path not in public_paths and not gradio_path:
                try:
                    identity = authenticate_headers(
                        request.headers,
                        resolved_settings,
                        mcp_request=(
                            request.url.path == "/mcp"
                            or request.url.path.startswith("/mcp/")
                        ),
                    )
                except AuthenticationError:
                    body = ErrorResponse(
                        error="Invalid or missing API credentials",
                        code="unauthorized",
                        request_id=request_id,
                    )
                    response = JSONResponse(status_code=401, content=body.model_dump())
                else:
                    request.state.auth_user = identity.user
                    response = await call_next(request)
            else:
                request.state.auth_user = "anonymous"
                response = await call_next(request)
        finally:
            elapsed = round(time.perf_counter() - started, 3)
            logger.info(
                "request completed",
                extra={
                    "method": request.method,
                    "path": request.url.path,
                    "status_code": getattr(locals().get("response"), "status_code", 500),
                    "execution_time": elapsed,
                },
            )
            request_id_context.reset(token)
            LocalContext.request.reset(gradio_token)
        response.headers["X-Request-ID"] = request_id
        return response

    @application.exception_handler(GatewayError)
    async def gateway_error_handler(request: Request, exc: GatewayError) -> JSONResponse:
        logger.warning(
            "gateway request failed",
            extra={"error_code": exc.code, "status_code": exc.status_code},
        )
        body = ErrorResponse(
            error=exc.message,
            code=exc.code,
            request_id=getattr(request.state, "request_id", "-"),
        )
        return JSONResponse(status_code=exc.status_code, content=body.model_dump())

    @application.exception_handler(RequestValidationError)
    async def validation_error_handler(
        request: Request, exc: RequestValidationError
    ) -> JSONResponse:
        errors = "; ".join(
            f"{'.'.join(str(item) for item in error['loc'])}: {error['msg']}"
            for error in exc.errors()
        )
        logger.warning(
            "request validation failed",
            extra={"error_code": "validation_error", "status_code": 422},
        )
        body = ErrorResponse(
            error=errors,
            code="validation_error",
            request_id=getattr(request.state, "request_id", "-"),
        )
        return JSONResponse(status_code=422, content=body.model_dump())

    @application.exception_handler(HTTPException)
    async def http_error_handler(request: Request, exc: HTTPException) -> JSONResponse:
        body = ErrorResponse(
            error=str(exc.detail),
            code="http_error",
            request_id=getattr(request.state, "request_id", "-"),
        )
        return JSONResponse(status_code=exc.status_code, content=body.model_dump())

    @application.exception_handler(Exception)
    async def unexpected_error_handler(request: Request, exc: Exception) -> JSONResponse:
        logger.exception("unhandled request error")
        body = ErrorResponse(
            error="Internal server error",
            code="internal_error",
            request_id=getattr(request.state, "request_id", "-"),
        )
        return JSONResponse(status_code=500, content=body.model_dump())

    application.include_router(health.router)
    application.include_router(image.router)
    application.include_router(video.router)
    application.include_router(audio.router)
    application.include_router(workflow.router)
    application.mount("/mcp", mcp.app, name="mcp")

    @application.get("/", include_in_schema=False)
    @application.get("/ui", include_in_schema=False)
    async def ui_entry(request: Request) -> RedirectResponse:
        query = f"?{request.url.query}" if request.url.query else ""
        return RedirectResponse(url=f"/ui/{query}")

    application.mount(
        "/output",
        StaticFiles(directory=resolved_settings.output_folder, check_dir=False),
        name="output",
    )
    return gr.mount_gradio_app(
        application,
        gradio_ui,
        path="/ui",
        server_name=resolved_settings.host,
        server_port=resolved_settings.port,
        allowed_paths=[str(outputs.root)],
        show_error=False,
        max_file_size=f"{resolved_settings.max_upload_mb}mb",
        ssr_mode=False,
        mcp_server=True,
    )


app = create_app()


if __name__ == "__main__":
    import uvicorn

    settings = get_settings()
    uvicorn.run(app, host=settings.host, port=settings.port, proxy_headers=True)