litellm/litellm/proxy/auth/model_checks.py

# What is this?
## Common checks for /v1/models and `/model/info`
import copy
from collections.abc import Sequence
from typing import Any, Final

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.router import Router
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params
from litellm.types.utils import LlmProviders
from litellm.utils import get_valid_models

_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)


_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)


def _check_wildcard_routing(model: str) -> bool:
    """
    Returns True if a model is a provider wildcard.

    eg:
    - anthropic/*
    - openai/*
    - *
    """
    if "*" in model:
        return True
    return False


def get_provider_models(provider: str, litellm_params: LiteLLM_Params | None = None) -> list[str] | None:
    """
    Returns the list of known models by provider
    """
    if provider == "*":
        return get_valid_models(litellm_params=litellm_params)

    if provider in litellm.models_by_provider:
        provider_models: Final = get_valid_models(custom_llm_provider=provider, litellm_params=litellm_params)
        return provider_models
    return None


def _get_models_from_access_groups(
    model_access_groups: dict[str, list[str]],
    all_models: list[str],
    include_model_access_groups: bool | None = False,
    proxy_model_list: Sequence[str] | None = None,
) -> list[str]:
    # a grant naming both a deployed model and an access group means both at runtime
    # (_check_model_access_helper unions them), so listings must keep the literal too
    deployed_model_names: Final = frozenset(proxy_model_list or ())
    kept_models: Final = [
        model
        for model in all_models
        if model not in model_access_groups or include_model_access_groups or model in deployed_model_names
    ]
    member_models: Final = [
        member for model in all_models if model in model_access_groups for member in model_access_groups[model]
    ]
    return kept_models + member_models


async def get_mcp_server_ids(
    user_api_key_dict: UserAPIKeyAuth,
) -> list[str]:
    """
    Returns the list of MCP server ids for a given key by querying the object_permission table
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        return []

    if user_api_key_dict.object_permission_id is None:
        return []

    # Make a direct SQL query to get just the mcp_servers
    try:
        result: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
            where={"object_permission_id": user_api_key_dict.object_permission_id},
        )
        if result and result.mcp_servers:
            return result.mcp_servers
        return []
    except Exception:
        return []


def get_key_models(
    user_api_key_dict: UserAPIKeyAuth,
    proxy_model_list: list[str],
    model_access_groups: dict[str, list[str]],
    include_model_access_groups: bool | None = False,
    only_model_access_groups: bool | None = False,
) -> list[str]:
    """
    Returns:
    - List of model name strings
    - Empty list if no models set
    - If model_access_groups is provided, only return models that are in the access groups
    - If include_model_access_groups is True, it includes the 'keys' of the model_access_groups
      in the response - {"beta-models": ["gpt-4", "claude-v1"]} -> returns 'beta-models'
    """
    all_models: list[str] = []
    if len(user_api_key_dict.models) > 0:
        all_models = list(user_api_key_dict.models)  # copy to avoid mutating cached objects
        if SpecialModelNames.all_team_models.value in all_models:
            all_models = list(user_api_key_dict.team_models)
            if SpecialModelNames.all_team_models.value in all_models:
                all_models = [model for model in all_models if model != SpecialModelNames.all_team_models.value]
                all_models.extend(proxy_model_list)
                if include_model_access_groups:
                    all_models.extend(model_access_groups.keys())
        if SpecialModelNames.all_proxy_models.value in all_models:
            # Union team-specific models with proxy-wide list, matching get_team_models()
            # which uses set.update(). Replacing discards team-specific models.
            merged = set(all_models) | set(proxy_model_list)
            merged.discard(SpecialModelNames.all_proxy_models.value)
            all_models = list(merged)
            if include_model_access_groups:
                all_models.extend(model_access_groups.keys())

    all_models = _get_models_from_access_groups(
        model_access_groups=model_access_groups,
        all_models=all_models,
        include_model_access_groups=include_model_access_groups,
        proxy_model_list=proxy_model_list,
    )

    # deduplicate while preserving order
    all_models = list(dict.fromkeys(all_models))

    verbose_proxy_logger.debug("ALL KEY MODELS - %s", len(all_models))
    return all_models


def get_team_models(
    team_models: list[str],
    proxy_model_list: list[str],
    model_access_groups: dict[str, list[str]],
    include_model_access_groups: bool | None = False,
) -> list[str]:
    """
    Returns:
    - List of model name strings
    - Empty list if no models set
    - If model_access_groups is provided, only return models that are in the access groups
    """
    all_models_set: Final[set[str]] = set()
    if len(team_models) > 0:
        all_models_set.update(team_models)
        if SpecialModelNames.all_team_models.value in all_models_set:
            all_models_set.update(team_models)
            # GH#30619: expand all-team-models sentinel
            # to the actual proxy model list
            all_models_set.discard(SpecialModelNames.all_team_models.value)
            all_models_set.update(proxy_model_list)
            if include_model_access_groups:
                all_models_set.update(model_access_groups.keys())
        if SpecialModelNames.all_proxy_models.value in all_models_set:
            all_models_set.update(proxy_model_list)
            if include_model_access_groups:
                all_models_set.update(model_access_groups.keys())

    all_models = _get_models_from_access_groups(
        model_access_groups=model_access_groups,
        all_models=list(all_models_set),
        include_model_access_groups=include_model_access_groups,
        proxy_model_list=proxy_model_list,
    )

    # deduplicate while preserving order
    all_models = list(dict.fromkeys(all_models))

    verbose_proxy_logger.debug("ALL TEAM MODELS - %s", len(all_models))
    return all_models


def get_complete_model_list(
    key_models: Sequence[str],
    team_models: Sequence[str],
    proxy_model_list: list[str],
    user_model: str | None,
    infer_model_from_keys: bool | None,
    return_wildcard_routes: bool | None = False,
    llm_router: Router | None = None,
    model_access_groups: dict[str, list[str]] = {},
    include_model_access_groups: bool | None = False,
    only_model_access_groups: bool | None = False,
    team_id: str | None = None,
) -> list[str]:
    """Logic for returning complete model list for a given key + team pair"""

    """
    - If key list is empty -> defer to team list
    - If team list is empty -> defer to proxy model list

    If list contains wildcard -> return known provider models
    """

    unique_models: Final = []

    def append_unique(models):
        for model in models:
            if model not in unique_models and model != SpecialModelNames.no_default_models.value:
                unique_models.append(model)

    if key_models:
        append_unique(key_models)
    elif team_models:
        append_unique(team_models)
    else:
        append_unique(proxy_model_list)
        if include_model_access_groups:
            append_unique(list(model_access_groups.keys()))  # TODO: keys order

        if user_model:
            append_unique([user_model])

        if infer_model_from_keys:
            valid_models: Final = get_valid_models()
            append_unique(valid_models)

    if only_model_access_groups:
        model_access_groups_to_return: Final[list[str]] = []
        for model in unique_models:
            if model in model_access_groups:
                model_access_groups_to_return.append(model)
        return model_access_groups_to_return

    all_wildcard_models: Final = _get_wildcard_models(
        unique_models=unique_models,
        return_wildcard_routes=return_wildcard_routes,
        llm_router=llm_router,
        team_id=team_id,
    )

    complete_model_list: Final = unique_models + all_wildcard_models

    return complete_model_list


def _hydrate_litellm_credential_name(
    litellm_params: LiteLLM_Params | None,
) -> LiteLLM_Params | None:
    if litellm_params is None or litellm_params.litellm_credential_name is None:
        return litellm_params

    credential_values: Final = CredentialAccessor.get_credential_values(litellm_params.litellm_credential_name)
    if not credential_values:
        return litellm_params

    litellm_params = litellm_params.model_copy()
    for key, value in credential_values.items():
        if key in _CREDENTIAL_LITELLM_PARAM_FIELDS and getattr(litellm_params, key, None) is None:
            setattr(litellm_params, key, value)
    litellm_params.litellm_credential_name = None
    return litellm_params


def get_known_models_from_wildcard(wildcard_model: str, litellm_params: LiteLLM_Params | None = None) -> list[str]:
    wildcard_model_to_expand: Final = (
        litellm_params.model
        if wildcard_model == "*"
        and litellm_params is not None
        and _check_wildcard_routing(litellm_params.model)
        and "/" in litellm_params.model
        else wildcard_model
    )
    try:
        wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1)
    except ValueError:  # safely fail
        return []

    # Use provider from litellm_params when available, otherwise from wildcard prefix
    # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router)
    if litellm_params is not None:
        try:
            provider = litellm_params.model.split("/", 1)[0]
        except ValueError:
            provider = wildcard_provider_prefix
    else:
        provider = wildcard_provider_prefix

    litellm_params = _hydrate_litellm_credential_name(litellm_params)

    wildcard_models = get_provider_models(provider=provider, litellm_params=litellm_params)

    if wildcard_models is None:
        return []
    if wildcard_suffix != "*":
        ## CHECK IF PARTIAL FILTER e.g. `gemini-*`
        model_prefix: Final = wildcard_suffix.replace("*", "")

        is_partial_filter: Final = any(wc_model.startswith(model_prefix) for wc_model in wildcard_models)
        if is_partial_filter:
            filtered_wildcard_models = [wc_model for wc_model in wildcard_models if wc_model.startswith(model_prefix)]
            wildcard_models = filtered_wildcard_models
        else:
            # add model prefix to wildcard models
            wildcard_models = [f"{model_prefix}{model}" for model in wildcard_models]

    known_providers: Final = {provider.value for provider in LlmProviders}
    suffix_appended_wildcard_models: Final = []
    for model in wildcard_models:
        if not model.startswith(wildcard_provider_prefix):
            # `get_provider_models` returns provider-prefixed ids (e.g. "ollama/gemma3:1b").
            # When the wildcard uses a custom prefix (e.g. "ollama_server1/*" to distinguish
            # multiple instances), replace that existing provider prefix instead of stacking
            # both, which would otherwise yield an uncallable "ollama_server1/ollama/gemma3:1b".
            # Only strip the leading segment when it is a known provider, so ids whose first
            # segment is an org rather than a provider (e.g. "meta-llama/Llama-3-8B") keep it.
            leading, sep, model_suffix = model.partition("/")
            if sep and leading in known_providers:
                model = f"{wildcard_provider_prefix}/{model_suffix}"
            else:
                model = f"{wildcard_provider_prefix}/{model}"
        suffix_appended_wildcard_models.append(model)
    return suffix_appended_wildcard_models or []


def expand_wildcard_deployments_for_model_info(
    deployments: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Expand wildcard deployments into one row per known provider model.

    PR #30025 changed /model/info to read from llm_router.model_list (correct,
    so team-scoped rows are included). This function restores wildcard expansion
    on top of that: a wildcard deployment like model_name="*" / litellm_params.model="openai/*"
    becomes one entry per known openai model, matching /v1/models behaviour.
    """
    expanded: Final[list[dict[str, Any]]] = []
    for deployment in deployments:
        model_name = str(deployment.get("model_name") or "")
        raw_params = deployment.get("litellm_params")
        litellm_params_dict: dict[str, Any] = raw_params if isinstance(raw_params, dict) else {}
        litellm_model = str(litellm_params_dict.get("model") or "")

        # Determine the wildcard pattern to expand.
        # Branch order matters: only fall to litellm_model when model_name is
        # also a wildcard, so a concrete model_name is never overwritten.
        if _check_wildcard_routing(model_name) and "/" in model_name:
            wildcard_pattern = model_name
        elif _check_wildcard_routing(model_name) and _check_wildcard_routing(litellm_model):
            wildcard_pattern = litellm_model
        elif _check_wildcard_routing(model_name):
            wildcard_pattern = model_name
        else:
            expanded.append(deployment)
            continue

        try:
            litellm_params = LiteLLM_Params.model_validate(litellm_params_dict) if litellm_params_dict else None
        except Exception:
            expanded.append(deployment)
            continue
        expanded_names = get_known_models_from_wildcard(
            wildcard_model=wildcard_pattern,
            litellm_params=litellm_params,
        )
        if not expanded_names:
            expanded.append(deployment)
            continue

        for name in expanded_names:
            row = copy.deepcopy(deployment)
            row["model_name"] = name
            params = row.get("litellm_params")
            if isinstance(params, dict):
                params["model"] = name
            expanded.append(row)

    return expanded


def _get_wildcard_models(
    unique_models: list[str],
    return_wildcard_routes: bool | None = False,
    llm_router: Router | None = None,
    team_id: str | None = None,
) -> list[str]:
    models_to_remove: Final = set()
    all_wildcard_models: Final = []
    for model in unique_models:
        if _check_wildcard_routing(model=model):
            if return_wildcard_routes:
                all_wildcard_models.append(model)

            models_to_remove.add(model)

            model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None
            if model_list:
                for router_model in model_list:
                    all_wildcard_models.extend(
                        get_known_models_from_wildcard(
                            wildcard_model=model,
                            litellm_params=LiteLLM_Params(**router_model["litellm_params"]),
                        )
                    )
            else:
                all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None))

    for model in models_to_remove:
        unique_models.remove(model)

    return all_wildcard_models


def get_all_fallbacks(
    model: str,
    llm_router: Router | None = None,
    fallback_type: str = "general",
) -> list[str]:
    """
    Get all fallbacks for a given model from the router's fallback configuration.

    Args:
        model: The model name to get fallbacks for
        llm_router: The LiteLLM router instance
        fallback_type: Type of fallback ("general", "context_window", "content_policy")

    Returns:
        List of fallback model names. Empty list if no fallbacks found.
    """
    if llm_router is None:
        return []

    # Get the appropriate fallback list based on type
    fallbacks_config: list = []
    if fallback_type == "general":
        fallbacks_config = getattr(llm_router, "fallbacks", [])
    elif fallback_type == "context_window":
        fallbacks_config = getattr(llm_router, "context_window_fallbacks", [])
    elif fallback_type == "content_policy":
        fallbacks_config = getattr(llm_router, "content_policy_fallbacks", [])
    else:
        verbose_proxy_logger.warning("Unknown fallback_type: %s", fallback_type)
        return []

    if not fallbacks_config:
        return []

    try:
        # Use existing function to get fallback model group
        fallback_model_group, _ = get_fallback_model_group(fallbacks=fallbacks_config, model_group=model)

        if fallback_model_group is None:
            return []

        return fallback_model_group
    except Exception as e:
        verbose_proxy_logger.error("Error getting fallbacks for model %s: %s", model, e)
        return []