from collections.abc import Mapping from typing import Final, Literal import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. Accepts the same credential set as the `azure` provider: service principal (`tenant_id` / `client_id` / `client_secret`), a pre-fetched `azure_ad_token`, an OIDC federated token, username/password, or `DefaultAzureCredential` / managed identity. """ from litellm.llms.azure.common_utils import get_azure_ad_token params = GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() return get_azure_ad_token(params) def get_azure_ai_auth_headers( api_key: str | None, litellm_params: Mapping[str, object] | None = None, api_key_header: AzureAIApiKeyHeader = "Authorization", api_key_env_var: str = "AZURE_AI_API_KEY", ) -> Mapping[str, str]: """ Build the auth headers for an Azure AI Foundry route. Prefers the API key when one is configured, and otherwise falls back to Entra ID / OAuth, sending the access token as a bearer token. """ if api_key: return {api_key_header: f"Bearer {api_key}" if api_key_header == "Authorization" else api_key} azure_ad_token = get_azure_ai_entra_token(litellm_params=litellm_params) if azure_ad_token: return {"Authorization": f"Bearer {azure_ad_token}"} raise ValueError( f"Missing Azure AI credentials - set an API key (`api_key` or {api_key_env_var}), or Entra ID / OAuth " "credentials (`tenant_id` + `client_id` + `client_secret`, `azure_ad_token`, an OIDC token, or a managed " "identity with `litellm.enable_azure_ad_token_refresh = True`)" ) AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model" class AzureFoundryModelInfo(BaseLLMModelInfo): """Model info for Azure AI / Azure Foundry models.""" def __init__(self, model: str | None = None): self._model = model @staticmethod def get_azure_ai_route(model: str) -> Literal["agents", "model_router", "default"]: """ Get the Azure AI route for the given model. Similar to BedrockModelInfo.get_bedrock_route(). Supported routes: - agents: azure_ai/agents/ - model_router: azure_ai/model_router/ or models with "model-router"/"model_router" in name - default: standard models """ if "agents/" in model: return "agents" # Detect model router by prefix (model_router/) or by name containing "model-router"/"model_router" model_lower: Final = model.lower() if ( "model_router/" in model_lower or "model-router/" in model_lower or "model-router" in model_lower or "model_router" in model_lower ): return "model_router" return "default" @staticmethod def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None: """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``. Reading this beats re-deriving the route from a model string: the stamp is set on the code path that was actually taken, so it holds no matter what the caller named the model. """ if not hidden_params: return None selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY) if isinstance(selected, str) and selected: return selected return None @staticmethod def is_model_router_call( model: str | None = None, hidden_params: Mapping[str, object] | None = None, ) -> bool: """Whether a request went down the Azure Model Router route. Prefers the response stamp, then the deployment's litellm model path, and only then the caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router name heuristic lives in exactly one place. """ if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None: return True deployment_model: Final = ( hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None ) return any( isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router" for candidate in (deployment_model, model) ) @staticmethod def get_api_base(api_base: str | None = None) -> str | None: return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE") @staticmethod def get_api_key(api_key: str | None = None) -> str | None: return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY") @property def api_version(self, api_version: str | None = None) -> str | None: api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") return api_version def get_token_counter(self) -> BaseTokenCounter | None: """ Factory method to create a token counter for Azure AI. Returns: AzureAIAnthropicTokenCounter for Claude models, None otherwise. """ # Only return token counter for Claude models if self._model and "claude" in self._model.lower(): from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import ( AzureAIAnthropicTokenCounter, ) return AzureAIAnthropicTokenCounter() return None def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: """ Returns a list of models supported by Azure AI. Azure AI doesn't have a standard model listing endpoint, so this returns an empty list. """ return [] ######################################################### # Not implemented methods ######################################################### @staticmethod def strip_model_router_prefix(model: str) -> str: """ Strip the model_router prefix from model name. Examples: - "model_router/gpt-4o" -> "gpt-4o" - "model-router/gpt-4o" -> "gpt-4o" - "gpt-4o" -> "gpt-4o" Args: model: Model name potentially with model_router prefix Returns: Model name without the prefix """ if "model_router/" in model: return model.split("model_router/", 1)[1] if "model-router/" in model: return model.split("model-router/", 1)[1] return model @staticmethod def get_base_model(model: str) -> str: """ Get the base model name, stripping any Azure AI routing prefixes. Args: model: Model name potentially with routing prefixes Returns: Base model name """ # Strip model_router prefix if present model = AzureFoundryModelInfo.strip_model_router_prefix(model) return model @staticmethod def get_azure_ai_config_for_model(model: str): """ Get the appropriate Azure AI config class for the given model. Routes to specialized configs based on model type: - Model Router: AzureModelRouterConfig - Claude models: AzureAnthropicConfig - Default: AzureAIStudioConfig Args: model: The model name Returns: The appropriate config instance """ azure_ai_route: Final = AzureFoundryModelInfo.get_azure_ai_route(model) if azure_ai_route == "model_router": from litellm.llms.azure_ai.azure_model_router.transformation import ( AzureModelRouterConfig, ) return AzureModelRouterConfig() elif "claude" in model.lower(): from litellm.llms.azure_ai.anthropic.transformation import ( AzureAnthropicConfig, ) return AzureAnthropicConfig() else: from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig return AzureAIStudioConfig() def validate_environment( self, headers: dict, model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, api_base: str | None = None, ) -> dict: """Azure Foundry sends api key in query params""" raise NotImplementedError("Azure Foundry does not support environment validation")