models.py

"""Handler for REST API call to list available models."""

from typing import Annotated, Any

from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.params import Depends
from llama_stack_client import APIConnectionError

from lightspeed_stack.authentication import get_auth_dependency
from lightspeed_stack.authentication.interface import AuthTuple
from lightspeed_stack.authorization.middleware import authorize
from lightspeed_stack.client import AsyncLlamaStackClientHolder
from lightspeed_stack.configuration import configuration
from lightspeed_stack.log import get_logger
from lightspeed_stack.models.api.requests.catalog import ModelFilter
from lightspeed_stack.models.api.responses.constants import (
    UNAUTHORIZED_OPENAPI_EXAMPLES,
)
from lightspeed_stack.models.api.responses.error import (
    ForbiddenResponse,
    InternalServerErrorResponse,
    ServiceUnavailableResponse,
    UnauthorizedResponse,
)
from lightspeed_stack.models.api.responses.successful import ModelsResponse
from lightspeed_stack.models.config import Action
from lightspeed_stack.utils.endpoints import check_configuration_loaded

logger = get_logger(__name__)
router = APIRouter(tags=["models"])


def parse_llama_stack_model(model: Any) -> dict[str, Any]:
    """
    Parse llama-stack model.

    Converting the new llama-stack model format (0.4.x) with custom_metadata.

    Parameters:
        model: Model object from llama-stack (has id, custom_metadata, object fields)

    Returns:
        dict: Model in legacy format with identifier, provider_id, model_type, etc.
    """
    custom_metadata = getattr(model, "custom_metadata", {}) or {}

    model_type = str(custom_metadata.get("model_type", "unknown"))

    metadata = {
        k: v
        for k, v in custom_metadata.items()
        if k not in ("provider_id", "provider_resource_id", "model_type")
    }

    return {
        "identifier": getattr(model, "id", ""),
        "metadata": metadata,
        "api_model_type": model_type,
        "provider_id": str(custom_metadata.get("provider_id", "")),
        "type": getattr(model, "object", "model"),
        "provider_resource_id": str(custom_metadata.get("provider_resource_id", "")),
        "model_type": model_type,
    }


models_responses: dict[int | str, dict[str, Any]] = {
    200: ModelsResponse.openapi_response(),
    401: UnauthorizedResponse.openapi_response(examples=UNAUTHORIZED_OPENAPI_EXAMPLES),
    403: ForbiddenResponse.openapi_response(examples=["endpoint"]),
    500: InternalServerErrorResponse.openapi_response(examples=["configuration"]),
    503: ServiceUnavailableResponse.openapi_response(
        examples=["llama stack", "kubernetes api"]
    ),
}


@router.get("/models", responses=models_responses)
@authorize(Action.GET_MODELS)
async def models_endpoint_handler(
    request: Request,
    auth: Annotated[AuthTuple, Depends(get_auth_dependency())],
    model_type: Annotated[ModelFilter, Query()],
) -> ModelsResponse:
    """
    Handle requests to the /models endpoint.

    Process GET requests to the /models endpoint, returning a list of available
    models from the Llama Stack service. It is possible to specify "model_type"
    query parameter that is used as a filter. For example, if model type is set
    to "llm", only LLM models will be returned:

        curl http://localhost:8080/v1/models?model_type=llm

    The "model_type" query parameter is optional. When not specified, all models
    will be returned.

    ### Parameters:
    - request: The incoming HTTP request (used by middleware).
    - auth: Authentication tuple from the auth dependency (used by middleware).
    - model_type: Optional filter to return only models matching this type.

    ### Raises:
    - HTTPException: with status 401 for unauthorized access.
    - HTTPException: with status 403 if permission is denied.
    - HTTPException: with status 422 if model_type parameter is
      improper.
    - HTTPException: with status 500 and a detail object containing `response`
      and `cause` when service configuration is wrong or incomplete.
    - HTTPException: with status 503 and a detail object containing `response`
      and `cause` when unable to connect to Llama Stack.

    ### Returns:
    - ModelsResponse: An object containing the list of available models.
    """
    # Used only by the middleware
    _ = auth

    # Nothing interesting in the request
    _ = request

    check_configuration_loaded(configuration)

    llama_stack_configuration = configuration.llama_stack_configuration
    logger.info("Llama stack config: %s", llama_stack_configuration)

    try:
        # try to get Llama Stack client
        client = AsyncLlamaStackClientHolder().get_client()
        # retrieve models
        models = await client.models.list()

        # parse models to legacy format
        parsed_models = [parse_llama_stack_model(model) for model in models]

        # optional filtering by model type
        if model_type.model_type is not None:
            parsed_models = [
                model
                for model in parsed_models
                if model["model_type"] == model_type.model_type
            ]

        return ModelsResponse(models=parsed_models)

    # Connection to Llama Stack server failed
    except APIConnectionError as e:
        logger.error("Unable to connect to Llama Stack: %s", e)
        response = ServiceUnavailableResponse(backend_name="Llama Stack", cause=str(e))
        raise HTTPException(**response.model_dump()) from e