main.py

"""Definition of FastAPI based web service."""

import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Final

import sentry_sdk  # pyright: ignore[reportMissingImports]
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from llama_stack_client import APIConnectionError, AsyncLlamaStackClient
from starlette.routing import Mount, Route, WebSocketRoute
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from lightspeed_stack import version
from lightspeed_stack.a2a_storage import A2AStorageFactory
from lightspeed_stack.app import routers
from lightspeed_stack.app.database import create_tables, initialize_database
from lightspeed_stack.app.endpoints.streaming_query import (
    shutdown_background_topic_summary_tasks,
)
from lightspeed_stack.authorization.azure_token_manager import AzureEntraIDManager
from lightspeed_stack.client import AsyncLlamaStackClientHolder
from lightspeed_stack.configuration import configuration
from lightspeed_stack.log import get_logger
from lightspeed_stack.metrics import recording
from lightspeed_stack.metrics.utils import setup_model_metrics
from lightspeed_stack.models.api.responses.error import InternalServerErrorResponse
from lightspeed_stack.sentry import initialize_sentry
from lightspeed_stack.utils.common import register_mcp_servers_async
from lightspeed_stack.utils.degraded_mode import DegradedModeTracker
from lightspeed_stack.utils.llama_stack_version import check_llama_stack_version

logger = get_logger(__name__)

logger.info("Initializing app")


service_name = configuration.configuration.name

# Global OpenAPI tags so every operation tag is declared (Spectral: operation-tag-defined).
_OPENAPI_TAGS: Final[list[dict[str, str]]] = [
    {"name": "a2a", "description": "Agent-to-Agent (A2A) protocol."},
    {"name": "authorized", "description": "Authorization probe."},
    {"name": "config", "description": "Service configuration."},
    {"name": "conversations_v1", "description": "Conversations API v1."},
    {"name": "conversations_v2", "description": "Conversations API v2."},
    {"name": "feedback", "description": "User feedback."},
    {"name": "health", "description": "Health and readiness probes."},
    {"name": "info", "description": "Service information."},
    {"name": "mcp-auth", "description": "MCP client authentication options."},
    {"name": "mcp-servers", "description": "MCP server registration."},
    {"name": "metrics", "description": "Prometheus metrics."},
    {"name": "models", "description": "LLM models."},
    {"name": "prompts", "description": "Prompt management."},
    {"name": "providers", "description": "Inference providers."},
    {"name": "query", "description": "Non-streaming query."},
    {"name": "rags", "description": "RAG configuration."},
    {"name": "responses", "description": "OpenAI-compatible Responses API."},
    {"name": "rlsapi-v1", "description": "RLS API v1 (inference)."},
    {"name": "root", "description": "Service root."},
    {
        "name": "saved-prompts",
        "description": "Saved prompts configuration and management.",
    },
    {"name": "shields", "description": "Safety shields."},
    {"name": "streaming_query", "description": "Streaming query (SSE)."},
    {"name": "streaming_query_interrupt", "description": "Streaming interrupt."},
    {"name": "tools", "description": "Tools."},
    {"name": "vector-stores", "description": "Vector stores and files."},
]


# running on FastAPI startup
@asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
    """
    Initialize app resources.

    FastAPI lifespan context: initializes configuration, Llama client, MCP servers,
    logger, and database before serving requests.
    """
    configuration.load_configuration(os.environ["LIGHTSPEED_STACK_CONFIG_PATH"])

    initialize_sentry()

    llama_stack_config = configuration.configuration.llama_stack
    await AsyncLlamaStackClientHolder().load(llama_stack_config)
    client: AsyncLlamaStackClient = AsyncLlamaStackClientHolder().get_client()
    logger.debug("Llama Stack client initialized, trying to connect to Llama Stack")
    # Check connectivity to Llama Stack and set degraded mode if unavailable
    degraded_tracker = DegradedModeTracker()
    try:
        llama_stack_version = await check_llama_stack_version(
            client, llama_stack_config.max_retries, llama_stack_config.retry_delay
        )
        if llama_stack_version is None:
            logger.error("Cannot retrieve Llama Stack version, check connection")
            if llama_stack_config.allow_degraded_mode:
                degraded_tracker.set_degraded("Llama Stack connection check failed")
        else:
            logger.debug("Llama Stack version: %s", llama_stack_version)
            degraded_tracker.set_healthy()
    except APIConnectionError as e:
        # if degraded mode is allowed, simply ignore the exception
        llama_stack_url = llama_stack_config.url
        logger.error(
            "Failed to connect to Llama Stack at '%s'. "
            "Please verify that the 'llama_stack.url' configuration is correct "
            "and that the Llama Stack service is running and accessible. "
            "Original error: %s",
            llama_stack_url,
            e,
        )
        if llama_stack_config.allow_degraded_mode:
            logger.info("Entering degraded mode: LCORE running w/o Llama Stack")
            degraded_tracker.set_degraded(f"Failed to connect to Llama Stack: {e!s}")
        else:
            raise

    azure_entra_id_config = configuration.configuration.azure_entra_id
    if azure_entra_id_config is not None:
        AzureEntraIDManager().set_config(azure_entra_id_config)
        azure_base_url = await AsyncLlamaStackClientHolder().get_azure_base_url()
        AzureEntraIDManager().set_base_url(azure_base_url)
    logger.info("Registering MCP servers")
    await register_mcp_servers_async(logger, configuration.configuration)

    # Set up model metrics if in healthy mode
    if not degraded_tracker.is_degraded():
        try:
            await setup_model_metrics()
        except APIConnectionError as e:
            logger.warning("Failed to set up model metrics: %s", e, exc_info=True)

    logger.info("App startup complete")

    initialize_database()
    create_tables()

    yield

    # Cleanup resources on shutdown
    try:
        await shutdown_background_topic_summary_tasks()
        await A2AStorageFactory.cleanup()
    finally:
        # Flush pending Sentry events after cleanup so any errors during
        # shutdown are captured before the process exits.
        sentry_sdk.flush(timeout=2)
    logger.info("App shutdown complete")


app = FastAPI(
    root_path=configuration.service_configuration.root_path,
    title=f"{service_name} service - OpenAPI",
    summary=f"{service_name} service API specification.",
    description=f"{service_name} service API specification.",
    version=version.__version__,
    contact={
        "name": "Pavel Tisnovsky",
        "url": "https://github.com/tisnik/",
        "email": "ptisnovs@redhat.com",
    },
    license_info={
        "name": "Apache 2.0",
        "url": "https://www.apache.org/licenses/LICENSE-2.0.html",
    },
    servers=[
        {"url": "http://localhost:8080", "description": "Locally running service"}
    ],
    openapi_tags=_OPENAPI_TAGS,
    lifespan=lifespan,
)

cors = configuration.service_configuration.cors

app.add_middleware(
    CORSMiddleware,
    allow_origins=cors.allow_origins,
    allow_credentials=cors.allow_credentials,
    allow_methods=cors.allow_methods,
    allow_headers=cors.allow_headers,
)


class RestApiMetricsMiddleware:  # pylint: disable=too-few-public-methods
    """Pure ASGI middleware for REST API metrics.

    Record REST API request metrics for application routes and forward the
    request to the next ASGI handler.

    Only requests whose path is listed in the application's routes are
    measured.  For measured requests, this middleware records request duration
    and increments a per-path / per-status counter; it does not increment
    counters for the ``/metrics`` endpoint.

    This is implemented as a pure ASGI middleware (instead of using Starlette's
    ``BaseHTTPMiddleware``) to avoid the ``RuntimeError: No response returned``
    bug that occurs when ``call_next`` is used with long-running handlers such
    as LLM inference.  See https://issues.redhat.com/browse/RSPEED-2413.
    """

    def __init__(self, app: ASGIApp) -> None:  # pylint: disable=redefined-outer-name
        """Initialize the middleware."""
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process an ASGI request."""
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        # When root_path is set (e.g., /api/lightspeed), the proxy forwards
        # requests with the full prefixed path (/api/lightspeed/v1/infer) but
        # app_routes_paths contains only application-level paths (/v1/infer).
        # Strip the prefix so the path check and metric labels match the routes.
        root_path = scope.get("root_path", "")
        path: str = scope["path"]
        if root_path and path.startswith(root_path + "/"):
            path = path[len(root_path) :]
        logger.debug("Received request for path: %s", path)

        # Ignore paths that are not part of the app routes.
        if path not in app_routes_paths:
            await self.app(scope, receive, send)
            return

        logger.debug("Processing API request for path: %s", path)

        status_code = 500

        async def send_wrapper(message: Message) -> None:
            nonlocal status_code
            if message["type"] == "http.response.start":
                status_code = message["status"]
            await send(message)

        # Measure duration and forward the request.  Use try/finally so the
        # call counter is always incremented, even when the inner app raises.
        try:
            with recording.measure_response_duration(path):
                await self.app(scope, receive, send_wrapper)
        finally:
            # Ignore /metrics endpoint that will be called periodically.
            if not path.endswith("/metrics"):
                recording.record_rest_api_call(path, status_code)


class GlobalExceptionMiddleware:  # pylint: disable=too-few-public-methods
    """Pure ASGI middleware to handle uncaught exceptions from all endpoints.

    This is implemented as a pure ASGI middleware (instead of using Starlette's
    ``BaseHTTPMiddleware``) to avoid the ``RuntimeError: No response returned``
    bug that occurs when ``call_next`` is used with long-running handlers such
    as LLM inference.  See https://issues.redhat.com/browse/RSPEED-2413.
    """

    def __init__(self, app: ASGIApp) -> None:  # pylint: disable=redefined-outer-name
        """Initialize the middleware."""
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process an ASGI request."""
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        response_started = False

        async def send_wrapper(message: Message) -> None:
            nonlocal response_started
            if message["type"] == "http.response.start":
                response_started = True
            await send(message)

        try:
            await self.app(scope, receive, send_wrapper)
        except HTTPException:
            raise
        except Exception as exc:  # pylint: disable=broad-exception-caught
            logger.exception("Uncaught exception in endpoint: %s", exc)
            if response_started:
                raise
            error_response = InternalServerErrorResponse.generic()
            model_dump = error_response.detail.model_dump()  # pylint: disable=no-member
            response = JSONResponse(
                status_code=error_response.status_code,
                content={"detail": model_dump},
            )
            await response(scope, receive, send)


logger.info("Including routers")
routers.include_routers(app)

app_routes_paths = [
    route.path
    for route in app.routes
    if isinstance(route, (Mount, Route, WebSocketRoute))
]

# Register pure ASGI middlewares.  Middleware execution order is the reverse of
# registration order: GlobalExceptionMiddleware (registered first) is innermost,
# RestApiMetricsMiddleware (registered last) is outermost.  This ensures metrics
# always observe a status code — including 500s synthesised by the exception
# middleware — rather than seeing a raw exception with no response.
app.add_middleware(GlobalExceptionMiddleware)
app.add_middleware(RestApiMetricsMiddleware)