From 784754ca4837b1cd4c27afb44aa221c7029dc549 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:44:55 -0800 Subject: [PATCH 01/19] fix(azure)!: answer 401 when Entra rejects the deployment's credential (#49863) * fix(azure): answer 401 when Entra rejects the deployment's credential A ClientAuthenticationError from the Entra token provider (wrong, expired or rotated client secret) used to surface as a retryable 500 carrying the app's client id and Entra trace ids. Map it to a 401 AuthenticationError that keeps the AADSTS code and first sentence with GUIDs redacted, on both the direct token fetch and the provider handed to the OpenAI SDK client. Network failures stay 500. * fix(azure): keep network failures to Entra a retryable 500 azure-identity wraps a transport failure in ClientAuthenticationError, so the 401 mapping now skips errors whose cause chain holds ServiceRequestError or ServiceResponseError * fix(azure): cover async token providers and Entra 5xx, and scope caller-sent Entra credentials * fix(azure): map only AADSTS-coded Entra rejections to 401 and keep caller credential failures out of cooldown An Entra ClientAuthenticationError now answers 401 only when its message carries an AADSTS code, its HTTP status is not 429 or 5xx, and the code is not the AADSTS50196 throttle, so a chained credential's outage or throttle stays a retryable error. Caller-sent Entra fields no longer change the deployment id; a 401 from a value the caller sent skips the shared deployment's cooldown instead. The Azure AI Foundry agent token fetch maps the same rejection, and the A2A endpoint answers it with HTTP 401. * fix(router): count router default_litellm_params as configured Entra credentials when deciding cooldown * fix(azure): keep transient Entra codes retryable and scope 401 handling to Entra rejections ClientSecretCredential drops Entra's HTTP status, so throttling and transient codes (AADSTS50033, 50196, 90006, 90024, 90033, 90055) now stay retryable instead of mapping to 401. The A2A endpoint answers 401 only for an AuthenticationError caused by an Entra token rejection, and the router skips cooldown for a caller-sent Entra field only on Azure-family deployments. The azure/ Decisions route and the Bing grounding search tool now map an environment-credential rejection to the same 401. * fix(azure): map only known Entra credential rejections to 401 and cover the OIDC token exchange Unlisted AADSTS codes keep the retryable 500, the router skips the 401 cooldown only when an Entra rejection caused it, and a federated token Entra rejects answers the same 401 on Azure chat, Responses and A2A agents * fix(azure): answer 401 for Entra rejections on the /azure passthrough and token counting The fallback hop skips cooldown when Entra rejected a credential the caller sent, Prometheus counts a deployment's own Entra rejection as a deployment failure again, and an OIDC exchange error keeps Entra's headers so Retry-After reaches the router. * fix(azure): keep a no-op caller Entra field from skipping cooldown and redact the passthrough 401 message A caller-sent Entra field now skips the deployment's cooldown only when it differs from the deployment, its named credential, the router defaults, the AZURE_* environment variables and the default scope. An empty value counts as unset, so a deployment whose own secret is wrong still cools down when the caller's field changes nothing about the token request. The /azure passthrough's Entra 401 message goes through the same internal hostname redaction as other client-facing errors. * fix(router): compare a caller's Entra field with the value the deployment uses A caller-sent Entra field skipped the cooldown only when it matched no configured source, so a value equal to an environment default the deployment overrides (a stale AZURE_CLIENT_SECRET, say) counted as the deployment's own and let a caller's rejected credential cool down a working deployment. The check now resolves the value the deployment would use without the caller, in order from the router defaults, the deployment, its named credential, the AZURE_* environment variables and the default scope, and skips the cooldown only when the caller's value differs from it. * fix(azure): follow the settings merge for empty Entra fields and keep non-rejection OIDC exchange errors unchanged * fix(azure): count only the failed attempt's Entra rejection and ignore a non-callable token provider * fix(azure): keep the streaming Entra rejection inside the failed attempt --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/prometheus.py | 3 +- litellm/llms/azure/common_utils.py | 211 +++++++- .../llms/azure/decisions/transformation.py | 3 +- litellm/llms/azure/search/transformation.py | 5 +- .../anthropic/count_tokens/handler.py | 3 + litellm/llms/azure_ai/common_utils.py | 67 ++- .../proxy/agent_endpoints/a2a_endpoints.py | 3 + .../llm_passthrough_endpoints.py | 13 +- litellm/router_api.py | 3 + litellm/router_python_core.py | 16 +- litellm/router_utils/cooldown_handlers.py | 39 ++ .../router_utils/fallback_event_handlers.py | 9 + .../test_prometheus_invalid_key_filtering.py | 28 + .../azure/decisions/test_transformation.py | 23 + ...st_bing_grounding_search_transformation.py | 36 ++ .../llms/azure/test_azure_common_utils.py | 493 +++++++++++++++++- .../count_tokens/test_token_counter.py | 73 +++ .../llms/azure_ai/test_azure_ai_entra_auth.py | 51 ++ .../agent_endpoints/test_a2a_endpoints.py | 147 ++++++ .../test_llm_pass_through_endpoints.py | 172 ++++++ .../test_fallback_event_handlers.py | 57 ++ tests/unit/test_router/test_router.py | 435 ++++++++++++++++ 22 files changed, 1858 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1c9a443b546..fc1a601f9da 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -54,6 +54,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_cust from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.models.end_user import LiteLLM_EndUserTable @@ -448,7 +449,7 @@ def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str, return not ( isinstance(exception, CallerCredentialAuthenticationError) or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)) - or logged_call.get("api_call_start_time") is None + or (logged_call.get("api_call_start_time") is None and not caused_by_entra_rejection(exception)) or _used_a_caller_credential(logged_call) ) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6bb6d584385..16c8bf6f9f1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,15 +1,16 @@ import asyncio import hashlib +import itertools import json import os import re -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import closing from types import MappingProxyType -from typing import Final, Literal, NamedTuple, cast +from typing import Final, Literal, NamedTuple, NoReturn, cast import httpx -from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from openai import APIError, AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict @@ -20,7 +21,7 @@ from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler, MaskedHTTPStatusError from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -52,6 +53,13 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object): _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams) +_AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitable[str]] | None]( + Callable[[], str | Awaitable[str]] | None +) + + +def _callable_azure_ad_token_provider(provider: object) -> Callable[[], str | Awaitable[str]] | None: + return _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(provider) if callable(provider) else None class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): @@ -81,6 +89,160 @@ class AzureOpenAIError(BaseLLMException): ) +class EntraRejectionError(AzureOpenAIError): + pass + + +_GUID: Final = re.compile(r"[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}") +_AADSTS_REASON: Final = re.compile(r"\bAADSTS(\d+)\b.*?(?:\.(?=[\s'\"\\]|$)|$)", re.MULTILINE) +_ENTRA_CREDENTIAL_REJECTION_CODES: Final = frozenset( + { + 50012, + 50013, + 50027, + 50034, + 50053, + 50055, + 50057, + 50059, + 50076, + 50079, + 50126, + 50128, + 53003, + 65001, + 70001, + 70002, + 70011, + 70021, + 70025, + 70052, + 90002, + 90094, + 500011, + 500014, + 700016, + 700023, + 700024, + 700025, + 700027, + 700211, + 700213, + 700223, + 700238, + 7000112, + 7000215, + 7000218, + 7000222, + } +) + + +def _entra_rejection_reason(message: str) -> str | None: + reason: Final = _AADSTS_REASON.search(message) + if reason is None or int(reason.group(1)) not in _ENTRA_CREDENTIAL_REJECTION_CODES: + return None + return _GUID.sub("", reason.group(0)) + + +def _entra_rejection(reason: str) -> EntraRejectionError: + return EntraRejectionError(status_code=401, message=f"Failed to get Azure AD token: {reason}") + + +def entra_rejection_as_auth_error(error: Exception) -> EntraRejectionError | None: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return None + if not isinstance(error, ClientAuthenticationError): + return None + status_code: Final = error.status_code + if status_code is not None and (status_code == 429 or status_code >= 500): + return None + reason: Final = _entra_rejection_reason(str(error)) + return None if reason is None else _entra_rejection(reason) + + +def _raise_entra_rejection_as_auth_error(error: Exception) -> NoReturn: + auth_error: Final = entra_rejection_as_auth_error(error) + if auth_error is None: + raise error + raise auth_error from error + + +def entra_token(provider: Callable[[], str]) -> str: + try: + return provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def entra_rejection_types() -> tuple[type[Exception], ...]: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return () + return (ClientAuthenticationError,) + + +_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16 + + +def _exception_chain(error: BaseException) -> Iterator[BaseException]: + link: BaseException | None = error # rebind-ok: advances one link of the cause chain per iteration + for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH): + if link is None: + return + yield link + link = link.__cause__ or link.__context__ + + +_REQUEST_FAILURE_TYPES: Final = (APIError, httpx.HTTPError) + + +def _ends_the_failed_attempt(link: BaseException, failed_attempt: BaseException) -> bool: + return not isinstance(link, EntraRejectionError) and ( + isinstance(link, _REQUEST_FAILURE_TYPES) or type(link) is type(failed_attempt) + ) + + +def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]: + yield error + yield from itertools.takewhile( + lambda link: not _ends_the_failed_attempt(link, error), itertools.islice(_exception_chain(error), 1, None) + ) + + +def entra_rejection_in(error: BaseException) -> EntraRejectionError | None: + return next( + (link for link in _links_within_the_failed_attempt(error) if isinstance(link, EntraRejectionError)), None + ) + + +def caused_by_entra_rejection(error: BaseException) -> bool: + return entra_rejection_in(error) is not None + + +async def _awaited_token(token: Awaitable[str]) -> str: + try: + return await token + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def _entra_rejections_as_auth_errors( + provider: Callable[[], str | Awaitable[str]], +) -> Callable[[], str | Awaitable[str]]: + def call_provider() -> str | Awaitable[str]: + try: + token: Final = provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + return token if isinstance(token, str) else _awaited_token(token) + + return call_provider + + def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "x-ratelimit-limit-requests" in headers: @@ -91,7 +253,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"] if "x-ratelimit-remaining-tokens" in headers: openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"] - llm_response_headers: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()} + llm_response_headers: Final = {f"llm_provider-{k}": v for k, v in headers.items()} return {**llm_response_headers, **openai_headers} @@ -257,22 +419,34 @@ def _cached_oidc_token(cache_key: str) -> str | None: return cached if isinstance(cached, str) else None +def _oidc_exchange_rejection_reason(error: MaskedHTTPStatusError) -> str | None: + if error.response.status_code not in (400, 401): + return None + return _entra_rejection_reason(error.text or error.response.text) + + def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str: cached_access_token: Final = _cached_oidc_token(cache_key) if cached_access_token is not None: return cached_access_token with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: - req_token: Final = client.post( - token_url, - data={ - "client_id": client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, - ) + try: + req_token: Final = client.post( + token_url, + data={ + "client_id": client_id, + "grant_type": "client_credentials", + "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": oidc_token, + }, + ) + except MaskedHTTPStatusError as error: + reason: Final = _oidc_exchange_rejection_reason(error) + if reason is None: + raise + raise _entra_rejection(reason) from error if req_token.status_code != 200: raise AzureOpenAIError( @@ -483,6 +657,9 @@ def get_azure_ad_token( raise except Exception as e: verbose_logger.error("Error calling Azure AD token provider: %s", e) + auth_error: Final = entra_rejection_as_auth_error(e) + if auth_error is not None: + raise auth_error from e raise RuntimeError(f"Failed to get Azure AD token: {e}") from e return azure_ad_token @@ -696,7 +873,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_version: str | None, is_async: bool, ) -> dict: - azure_ad_token_provider = litellm_params.get("azure_ad_token_provider") + azure_ad_token_provider = _callable_azure_ad_token_provider(litellm_params.get("azure_ad_token_provider")) # If we have api_key, then we have higher priority azure_ad_token = litellm_params.get("azure_ad_token") @@ -775,7 +952,7 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params["timeout"] = timeout if azure_ad_token_provider is not None: - azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider + azure_client_params["azure_ad_token_provider"] = _entra_rejections_as_auth_errors(azure_ad_token_provider) # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client diff --git a/litellm/llms/azure/decisions/transformation.py b/litellm/llms/azure/decisions/transformation.py index 3fb40e39702..bdbd784185f 100644 --- a/litellm/llms/azure/decisions/transformation.py +++ b/litellm/llms/azure/decisions/transformation.py @@ -8,6 +8,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, cached_azure_ad_token_refresh_provider, + entra_token, ) from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig from litellm.secret_managers.main import get_secret_str @@ -61,7 +62,7 @@ class AzureDecisionsConfig(OpenAIDecisionsConfig): scope: Final = ( params.azure_scope or get_secret_str("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default" ) - azure_headers["Authorization"] = f"Bearer {cached_azure_ad_token_refresh_provider(scope)()}" + azure_headers["Authorization"] = f"Bearer {entra_token(cached_azure_ad_token_refresh_provider(scope))}" return azure_headers async def avalidate_environment( diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index e255a667606..6d0241d4d3b 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -36,7 +36,7 @@ import httpx from pydantic import ConfigDict, ValidationError from litellm.litellm_core_utils.url_utils import ensure_api_route -from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider +from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider, entra_rejection_as_auth_error from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.responses import ( ResponsesSearchAnnotation, @@ -270,6 +270,9 @@ class BingGroundingSearchConfig(BaseSearchConfig): try: return minter() except Exception as e: + rejection: Final = entra_rejection_as_auth_error(e) + if rejection is not None: + raise rejection from e raise ValueError( f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} " f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / " diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 97ee44bd3bb..3609871fdf8 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -16,6 +16,7 @@ from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.azure.common_utils import EntraRejectionError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) @@ -135,6 +136,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): except AnthropicError: # Re-raise Anthropic exceptions as-is raise + except EntraRejectionError as rejection: + raise AnthropicError(status_code=rejection.status_code, message=rejection.message) from rejection except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code verbose_logger.error("HTTP error in CountTokens handler: %s", e) diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 147bf69e91e..1635b9b5679 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,7 @@ import asyncio -from collections.abc import Mapping +import os +from collections.abc import Callable, Mapping +from functools import partial from types import MappingProxyType from typing import Final, Literal from urllib.parse import urlparse @@ -52,6 +54,28 @@ AZURE_ENTRA_CREDENTIAL_PARAM_KEYS: Final = frozenset({"azure_ad_token", "client_ AZURE_ENTRA_LITELLM_PARAM_KEYS: Final = AZURE_ENTRA_CREDENTIAL_PARAM_KEYS | frozenset( {"tenant_id", "client_id", "azure_username", "azure_scope"} ) +AZURE_DEFAULT_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_AZURE_ENTRA_ENVIRONMENT_VARIABLES: Final = MappingProxyType( + { + "tenant_id": "AZURE_TENANT_ID", + "client_id": "AZURE_CLIENT_ID", + "client_secret": "AZURE_CLIENT_SECRET", + "azure_username": "AZURE_USERNAME", + "azure_password": "AZURE_PASSWORD", + "azure_ad_token": "AZURE_AD_TOKEN", + } +) + + +def azure_entra_environment_defaults() -> Mapping[str, object]: + return MappingProxyType( + { + **{key: os.getenv(variable) for key, variable in _AZURE_ENTRA_ENVIRONMENT_VARIABLES.items()}, + "azure_scope": os.getenv("AZURE_SCOPE", AZURE_DEFAULT_SCOPE), + } + ) + + AZURE_ENTRA_CREDENTIAL_HELP: Final = ( "Set `tenant_id` + `client_id` + `client_secret`, `azure_ad_token` (an `oidc/` token also needs " "`tenant_id` + `client_id`), or `client_id` + `azure_username` + `azure_password` in the agent's `litellm_params`" @@ -70,6 +94,17 @@ def _resolve_config_secret(value: object) -> str | None: return get_secret_str(value) if value.startswith("os.environ/") else value +def _agent_entra_token(provider: Callable[[], str], model: str) -> str: + from litellm.llms.azure.common_utils import EntraRejectionError, entra_token + + try: + return entra_token(provider) + except EntraRejectionError as rejection: + raise litellm.AuthenticationError( + message=rejection.message, llm_provider="azure_ai", model=model + ) from rejection + + def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: """Mints the Entra bearer from the agent's own litellm_params, never from process-wide AZURE_* env vars.""" from litellm.llms.azure.common_utils import ( @@ -88,18 +123,32 @@ def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: azure_username: Final = resolved["azure_username"] azure_password: Final = resolved["azure_password"] azure_ad_token: Final = resolved["azure_ad_token"] + model: Final = str(litellm_params.get("model") or "") if tenant_id and client_id and client_secret: - return get_azure_ad_token_from_entra_id( - tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_entra_id( + tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope + ), + model, + ) if client_id and azure_username and azure_password: - return get_azure_ad_token_from_username_password( - client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_username_password( + client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope + ), + model, + ) federated: Final = azure_ad_token is not None and azure_ad_token.startswith("oidc/") if azure_ad_token and federated and tenant_id and client_id: - return get_azure_ad_token_from_oidc( - azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope + return _agent_entra_token( + partial( + get_azure_ad_token_from_oidc, + azure_ad_token=azure_ad_token, + azure_client_id=client_id, + azure_tenant_id=tenant_id, + scope=scope, + ), + model, ) if azure_ad_token and not federated: return azure_ad_token diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 3dfaa7f213e..f0a4d47f1e4 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy @@ -1202,6 +1203,8 @@ async def invoke_agent_a2a( pass if isinstance(e, litellm.BadRequestError): return _jsonrpc_error(body.get("id"), -32602, e.message, 400) + if isinstance(e, litellm.AuthenticationError) and caused_by_entra_rejection(e): + return _jsonrpc_error(body.get("id"), -32000, e.message, 401) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 64cde072dc5..627164abb93 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,7 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider -from litellm._logging import verbose_proxy_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, @@ -58,6 +58,7 @@ from litellm.llms.anthropic.common_utils import ( anthropic_pass_through_api_base, merge_anthropic_beta_headers, ) +from litellm.llms.azure.common_utils import entra_rejection_in from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, @@ -2381,6 +2382,16 @@ async def _relay_router_model( ) except httpx.HTTPStatusError as upstream_error: return await _relay_upstream_response(upstream_error.response) + except Exception as error: + rejection: Final = entra_rejection_in(error) + if rejection is None: + raise + raise ProxyException( + message=redact_internal_details_from_client_message(rejection.message), + type=ProxyErrorTypes.auth_error, + param=None, + code=rejection.status_code, + ) from error if not is_streaming_request: return await _relay_upstream_response(cast(httpx.Response, result)) diff --git a/litellm/router_api.py b/litellm/router_api.py index 3007b5345c7..f22e9b1f351 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1402,6 +1402,9 @@ class RoutingContext: def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: return self._router._params_sent_to_deployment(request, deployment) + def default_litellm_params(self) -> Mapping[str, object]: + return self._router._active_default_litellm_params() + def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e251248d675..c0c5d74a88a 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -306,6 +306,7 @@ from litellm.router_utils.cooldown_handlers import ( get_cooldown_deployment_info, get_cooldown_deployments, is_background_response_poll_not_found, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, @@ -10309,8 +10310,9 @@ class PythonCore(DeprecatedRouterMembers): return False exception_status: Final = getattr(exception, "status_code", "") + failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs) - if is_caller_timeout_408(kwargs, exception_status): + if is_caller_timeout_408(failure_call_details, exception_status): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " "A timeout the caller set caused this 408, not the deployment's health." @@ -10334,6 +10336,18 @@ class PythonCore(DeprecatedRouterMembers): "Caller-scoped OAuth authentication failed, not the deployment." ) return False + if is_caller_entra_credential_failure( + caller_failure_deployment, + self._active_default_litellm_params(), + failure_call_details, + exception_status, + exception, + ): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "Entra ID rejected a credential the caller sent, not the deployment's own." + ) + return False exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e50d5bb47a6..7a8a645ad13 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -847,6 +847,45 @@ def is_stream_stall(exception: BaseException) -> bool: return isinstance(exception, litellm.Timeout) and isinstance(exception.__cause__, MidStreamFallbackError) +def is_caller_entra_credential_failure( + deployment: Deployment | None, + router_defaults: Mapping[str, object], + model_call_details: Mapping[str, object], + exception_status: str | int, + exception: BaseException, +) -> bool: + if cast_exception_status_to_int(exception_status) != 401 or deployment is None: + return False + + from litellm.llms.azure.common_utils import caused_by_entra_rejection + + if not caused_by_entra_rejection(exception): + return False + from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, azure_entra_environment_defaults + + call_params: Final = _CREDENTIAL_VALUES_ADAPTER.validate_python(model_call_details.get("litellm_params") or {}) + + credential_name: Final = deployment.litellm_params.litellm_credential_name + configured_by_precedence: Final[tuple[Mapping[str, object], ...]] = ( + router_defaults, + _CREDENTIAL_VALUES_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)), + _CREDENTIAL_VALUES_ADAPTER.validate_python( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ), + ) + environment: Final = azure_entra_environment_defaults() + return any( + call_params.get(key) is not None + and (call_params.get(key) or environment.get(key)) + != (_first_configured(configured_by_precedence, key) or environment.get(key)) + for key in AZURE_ENTRA_LITELLM_PARAM_KEYS + ) + + +def _first_configured(sources: tuple[Mapping[str, object], ...], key: str) -> object: + return next((value for source in sources if (value := source.get(key)) is not None), None) + + def is_caller_scoped_auth_failure(deployment: Deployment | None, exception_status: str | int) -> bool: if cast_exception_status_to_int(exception_status) not in (401, 403) or deployment is None: return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index c929cca44be..e8f832b4b2d 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -28,6 +28,7 @@ from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, is_advisor_orchestration_failure, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_stream_stall, @@ -132,6 +133,14 @@ def _trigger_cooldown_for_failed_deployment( deployment_id, ) return + if is_caller_entra_credential_failure( + deployment, ctx.default_litellm_params(), model_call_details, exception_status, exception + ): + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment %s: Entra ID rejected a credential the caller sent.", + deployment_id, + ) + return # Priority: deployment config > response header > router default, matching # Router.deployment_callback_on_failure's precedence for the primary path. diff --git a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py index 453fd28636d..3f3e2e0d465 100644 --- a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py @@ -512,6 +512,20 @@ def _provider_auth_failure_kwargs() -> dict[str, object]: } +def _deployment_entra_rejection_kwargs() -> dict[str, object]: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + return { + key: value for key, value in _provider_auth_failure_kwargs().items() if key not in CALLED_THE_PROVIDER + } | {"exception": raised} + + def _metric_values(name: str, model_id: str) -> list[float]: return [ sample.value @@ -541,6 +555,20 @@ class TestProviderAuthFailures: assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0] + @pytest.mark.asyncio + async def test_log_failure_event_records_entra_rejecting_the_deployment_credential_before_the_call( + self, prometheus_logger + ): + await prometheus_logger.async_log_failure_event( + kwargs=_deployment_entra_rejection_kwargs(), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] + def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger): prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs()) diff --git a/tests/unit/llms/azure/decisions/test_transformation.py b/tests/unit/llms/azure/decisions/test_transformation.py index b6de6084b9c..3f5e97ee000 100644 --- a/tests/unit/llms/azure/decisions/test_transformation.py +++ b/tests/unit/llms/azure/decisions/test_transformation.py @@ -125,6 +125,29 @@ async def test_default_credential_refreshes_without_opt_in( ) +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("systemone", (False, True)) +async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + monkeypatch: pytest.MonkeyPatch, is_async: bool, systemone: bool +) -> None: + from azure.core.exceptions import ClientAuthenticationError + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + await call_decisions(is_async, systemone, {}) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + + @pytest.mark.parametrize("kind", ("service-principal", "oidc", "provider")) async def test_azure_decisions_reuses_existing_azure_credential_resolution( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, kind: str diff --git a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py index dfd1474211f..53f426d8505 100644 --- a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py @@ -466,6 +466,12 @@ def _mock_response(): response.content = json.dumps(MOCK_BING_GROUNDING_RESPONSE).encode() return response +async def _search_what_is_litellm(is_async: bool) -> object: + if is_async: + return await litellm.asearch(query="what is litellm", search_provider="bing_grounding") + return litellm.search(query="what is litellm", search_provider="bing_grounding") + + @pytest.mark.usefixtures("_vcr_outcome_gate") class TestBingGroundingSearchTransformation: """ @@ -546,6 +552,36 @@ class TestBingGroundingSearchTransformation: } ] + @pytest.mark.parametrize("is_async", (False, True)) + async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + self, monkeypatch: pytest.MonkeyPatch, is_async: bool + ): + from azure.core.exceptions import ClientAuthenticationError + + from litellm.llms.azure import common_utils + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.delenv("BING_GROUNDING_TOKEN") + for name in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + try: + with pytest.raises(litellm.AuthenticationError) as excinfo: + await _search_what_is_litellm(is_async) + finally: + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + @pytest.mark.asyncio async def test_bing_grounding_asearch(self): with patch( # test-quality-ok: litellm.asearch has no client injection seam diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 0a41861e088..d66ba3f324e 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1,16 +1,23 @@ +import asyncio +import io import json import os +import sys import threading import traceback import uuid +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from queue import SimpleQueue -from typing import Callable, Final, Optional +from typing import Awaitable, Callable, Final, Optional from urllib.parse import parse_qs from unittest.mock import ANY, MagicMock, patch import httpx +import openai import pytest +import requests +from azure.core.pipeline.transport import RequestsTransport import litellm from litellm.llms.azure.common_utils import ( @@ -2500,3 +2507,487 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m closer.reap() assert wrapper.is_closed() is True + + +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" +_ENTRA_REJECTIONS: Final = { + "wrong_secret": ( + "AADSTS7000215", + "Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent in the " + "request is the client secret value, not the client secret ID, for a secret added to app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID} Correlation ID: ecfd2c73-ec0a-474c-bba7-205be903159d " + "Timestamp: 2026-10-29 07:19:40Z", + ), + "expired_secret": ( + "AADSTS7000222", + f"Authentication failed: AADSTS7000222: The provided client secret keys for app '{_ENTRA_APP_CLIENT_ID}' " + f"are expired. Visit the Azure portal to create new keys for your app. Trace ID: {_ENTRA_TRACE_ID} " + "Correlation ID: 354f9b57-5fed-47e4-a67d-89a7c56f6ddc Timestamp: 2026-10-29 07:19:50Z", + ), +} + + +def _raising_token_provider(error: Exception) -> Callable[[], str]: + def token_provider() -> str: + raise error + + return token_provider + + +def _azure_chat(token_provider: Callable[[], str]) -> object: + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + max_retries=0, + ) + + +def _azure_async_chat(token_provider: Callable[[], str]) -> object: + return asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + ) + ) + + +def _azure_responses(token_provider: Callable[[], str]) -> object: + return litellm.responses( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + input="hi", + azure_ad_token_provider=token_provider, + ) + + +def _azure_ai_decisions(token_provider: Callable[[], str]) -> object: + return litellm.systemone( + model="azure_ai/decision-1-qa", + api_base="https://litellm-qa.services.ai.azure.com", + state={"ticket": "The login page returns a 500 error after deploy"}, + questions={"is_bug": {"type": "noul", "instructions": "Is this a software bug report?"}}, + azure_ad_token_provider=token_provider, + ) + + +_AZURE_CALLS: Final = { + "chat": _azure_chat, + "async_chat": _azure_async_chat, + "responses": _azure_responses, + "decisions": _azure_ai_decisions, +} + + +@pytest.fixture +def no_azure_api_key(monkeypatch): + for env_var in ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "azure_key", None) + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize(("aadsts_code", "entra_message"), _ENTRA_REJECTIONS.values(), ids=_ENTRA_REJECTIONS.keys()) +def test_entra_rejecting_the_credential_answers_401_with_only_the_aadsts_reason( + no_azure_api_key, call, aadsts_code, entra_message +): + from azure.core.exceptions import ClientAuthenticationError + + with pytest.raises(litellm.AuthenticationError) as excinfo: + call(_raising_token_provider(ClientAuthenticationError(message=entra_message))) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +@pytest.mark.parametrize("api_version", ["2025-04-01-preview", "v1"]) +def test_entra_rejecting_an_async_token_provider_answers_401(no_azure_api_key, api_version): + from azure.core.exceptions import ClientAuthenticationError + + aadsts_code, entra_message = _ENTRA_REJECTIONS["wrong_secret"] + + async def token_provider() -> str: + raise ClientAuthenticationError(message=entra_message) + + async_token_provider: Final[Callable[[], Awaitable[str]]] = token_provider + + with pytest.raises(litellm.AuthenticationError) as excinfo: + asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version=api_version, + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=async_token_provider, + ) + ) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def _azure_identity_wrapped_token_provider(inner_error: Exception) -> Callable[[], str]: + from azure.core.exceptions import ClientAuthenticationError + + def token_provider() -> str: + try: + raise inner_error + except Exception as error: + raise ClientAuthenticationError(message=f"Authentication failed: {error}") from error + + return token_provider + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("failed_leg", ["request", "response"]) +def test_unreachable_entra_endpoint_stays_a_retryable_server_error(no_azure_api_key, call, failed_leg): + from azure.core.exceptions import ServiceRequestError, ServiceResponseError + + transport_error: Final = {"request": ServiceRequestError, "response": ServiceResponseError}[failed_leg]( + "('Unable to connect to proxy', NewConnectionError(\"Failed to establish a new connection: " + "[Errno 111] Connection refused\"))" + ) + + with pytest.raises(openai.APIError) as excinfo: + call(_azure_identity_wrapped_token_provider(transport_error)) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +_ENTRA_TENANT_ID: Final = "5b3e1c7a-2f4d-4e8b-9a61-0c2d7f8e9b14" +_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_ENTRA_OPENID_CONFIGURATION: Final = { + "token_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token", + "authorization_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/authorize", + "issuer": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/v2.0", +} + + +class _FakeEntraAdapter(requests.adapters.BaseAdapter): + def __init__(self, status_code: int | None, body: Mapping[str, object] | str) -> None: + super().__init__() + self.status_code: Final = status_code + self.body: Final = body + + def send(self, request: requests.PreparedRequest, *args: object, **kwargs: object) -> requests.Response: + if "openid-configuration" in str(request.url): + return self._answer(request, 200, _ENTRA_OPENID_CONFIGURATION) + if self.status_code is None: + raise requests.exceptions.ConnectionError("Failed to establish a new connection: Connection refused") + return self._answer(request, self.status_code, self.body) + + def close(self) -> None: + pass + + @staticmethod + def _answer( + request: requests.PreparedRequest, status_code: int, body: Mapping[str, object] | str + ) -> requests.Response: + response: Final = requests.Response() + response.status_code = status_code + response.raw = io.BytesIO((body if isinstance(body, str) else json.dumps(body)).encode()) + response.headers = requests.structures.CaseInsensitiveDict( + {"Content-Type": "application/json", "Retry-After": "0"} + ) + response.url = str(request.url) + response.request = request + return response + + +def _entra_answering(status_code: int | None, body: Mapping[str, object] | str = "") -> RequestsTransport: + session: Final = requests.Session() + session.mount("https://", _FakeEntraAdapter(status_code, body)) + return RequestsTransport(session=session, session_owner=False) + + +def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]: + return { + "error": "invalid_client", + "error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}", + "error_codes": [int(aadsts_code.removeprefix("AADSTS"))], + } + + +def _real_credential_token_provider(credential_kind: str, entra: RequestsTransport) -> Callable[[], str]: + from azure.identity import ( + ChainedTokenCredential, + ClientAssertionCredential, + ClientSecretCredential, + get_bearer_token_provider, + ) + + client_secret: Final = ClientSecretCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, "wrong-secret", transport=entra, disable_instance_discovery=True + ) + credentials: Final = { + "client_secret": client_secret, + "default_azure_credential_chain": ChainedTokenCredential(client_secret), + "workload_identity": ClientAssertionCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, lambda: "federated-token", transport=entra, retry_total=0 + ), + } + return get_bearer_token_provider(credentials[credential_kind], _ENTRA_SCOPE) + + +_REAL_CREDENTIAL_REJECTIONS: Final = { + "wrong_secret": ("client_secret", "AADSTS7000215", "Invalid client secret provided."), + "wrong_secret_through_a_credential_chain": ( + "default_azure_credential_chain", + "AADSTS7000215", + "Invalid client secret provided.", + ), + "unknown_app_through_workload_identity": ( + "workload_identity", + "AADSTS700016", + f"Application with identifier '{_ENTRA_APP_CLIENT_ID}' was not found in the directory.", + ), + "unknown_federated_subject_through_workload_identity": ( + "workload_identity", + "AADSTS700213", + "No matching federated identity record found for presented assertion subject 'system:serviceaccount:qa:proxy'.", + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "aadsts_code", "description"), + _REAL_CREDENTIAL_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_REJECTIONS.keys(), +) +def test_entra_rejection_through_a_real_azure_identity_credential_answers_401( + no_azure_api_key, credential_kind, aadsts_code, description +): + entra: Final = _entra_answering(401, _entra_error(aadsts_code, description)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, entra)) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +_REAL_CREDENTIAL_NON_REJECTIONS: Final = { + "unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""), + "html_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + "Service Unavailable", + ), + "json_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."), + ), + "tenant_throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS90055", "There are too many incoming requests."), + ), + "request_budget_exceeded_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90024", "A transient error has occurred. Try again."), + ), + "retryable_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50033", "A transient error has occurred."), + ), + "external_server_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90006", "The service is temporarily unavailable."), + ), + "graph_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90090", "The service is temporarily unavailable."), + ), + "token_issuance_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."), + ), + "throttled_without_a_code_through_a_credential_chain": ( + "default_azure_credential_chain", + 429, + {"error": "temporarily_unavailable", "error_description": "Too many requests."}, + ), + "json_503_through_workload_identity": ( + "workload_identity", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_workload_identity_naming_a_rejection_code": ( + "workload_identity", + 429, + _entra_error("AADSTS7000215", "Invalid client secret provided."), + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "status_code", "body"), + _REAL_CREDENTIAL_NON_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_NON_REJECTIONS.keys(), +) +def test_entra_outage_or_throttle_through_a_real_azure_identity_credential_stays_retryable( + no_azure_api_key, credential_kind, status_code, body +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, _entra_answering(status_code, body))) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + entra_status_code: int, + entra_body: Mapping[str, object], + entra_headers: Mapping[str, str] | None = None, +): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post(f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token").mock( + return_value=httpx.Response(entra_status_code, json=entra_body, headers=entra_headers) + ) + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token="oidc/env/LITELLM_QA_FEDERATED_TOKEN", + tenant_id=_ENTRA_TENANT_ID, + client_id=_ENTRA_APP_CLIENT_ID, + max_retries=0, + ) + + +def test_entra_rejecting_an_oidc_token_answers_401_with_only_the_aadsts_reason(no_azure_api_key, respx_mock, monkeypatch): + rejection: Final = _entra_error( + "AADSTS700213", f"No matching federated identity record found for app '{_ENTRA_APP_CLIENT_ID}'." + ) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat_with_oidc_token(respx_mock, monkeypatch, 400, rejection) + + assert excinfo.value.status_code == 401 + assert "AADSTS700213" in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def test_entra_answering_an_oidc_exchange_with_a_transient_code_keeps_its_status( + no_azure_api_key, respx_mock, monkeypatch +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, monkeypatch, 400, _entra_error("AADSTS90024", "A transient error has occurred. Try again.") + ) + + assert excinfo.value.status_code == 400 + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + +def test_entra_throttling_an_oidc_exchange_keeps_its_retry_after(no_azure_api_key, respx_mock, monkeypatch): + with pytest.raises(litellm.RateLimitError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + 429, + _entra_error("AADSTS90055", "TenantThrottlingError: There are too many pending requests."), + {"Retry-After": "7"}, + ) + + assert excinfo.value.response.headers.get("retry-after") == "7" + + +class _EntraServiceUnavailableAnswer: + status_code: Final = 503 + reason: Final = "Service Unavailable" + + def text(self, encoding: str | None = None) -> str: + return "Service Unavailable" + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("credential_kind", ["client_secret", "managed_identity"]) +def test_entra_answering_a_server_error_stays_a_retryable_server_error(no_azure_api_key, call, credential_kind): + from azure.core.exceptions import ClientAuthenticationError + from msal.exceptions import MsalServiceError + + token_providers: Final = { + "client_secret": _azure_identity_wrapped_token_provider( + MsalServiceError("HTTP Error: 503", error=None, error_description=None) + ), + "managed_identity": _raising_token_provider( + ClientAuthenticationError(message="No token received.", response=_EntraServiceUnavailableAnswer()) + ), + } + + with pytest.raises(openai.APIError) as excinfo: + call(token_providers[credential_kind]) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def test_token_provider_error_without_azure_core_installed_keeps_the_generic_token_error(monkeypatch): + monkeypatch.setitem(sys.modules, "azure.core.exceptions", None) + + with pytest.raises(RuntimeError, match="Failed to get Azure AD token: token service down"): + get_azure_ad_token( + GenericLiteLLMParams(azure_ad_token_provider=_raising_token_provider(ValueError("token service down"))) + ) + + +def test_image_generation_with_an_api_key_ignores_a_non_callable_azure_ad_token_provider(): + import respx + + with respx.mock: + image_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]}) + ) + response: Final = litellm.image_generation( + model="azure/gpt-image-1", + prompt="a red dot", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + api_key="configured-key", + azure_ad_token_provider="not-a-callable", + ) + + assert image_route.calls.last.request.headers["api-key"] == "configured-key" + assert response.data[0].b64_json == "aGk=" diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index faa33e559ca..7ce5d8b6a02 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -171,6 +171,79 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ ) +@pytest.mark.asyncio +async def test_entra_rejecting_the_deployment_oidc_token_becomes_a_401_error_response_without_counting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 401, + json={ + "error": "invalid_client", + "error_description": "AADSTS700211: No matching federated identity record found for presented " + "assertion issuer 'https://oidc.litellm-qa.invalid'. Trace ID: 0f5c2b1e", + "error_codes": [700211], + }, + ) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "qa-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 401) + assert "AADSTS700211" in (result.error_message or "") + assert not count_route.called + + +@pytest.mark.asyncio +async def test_entra_throttling_the_deployment_oidc_token_exchange_keeps_its_429( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response(429, json={"error": "temporarily_unavailable", "error_codes": [90055]}) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 429) + assert not count_route.called + + @pytest.mark.asyncio async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py index d0c1c4a6fb2..a8958b301db 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py @@ -254,6 +254,57 @@ def test_agent_entra_values_resolve_os_environ_references(monkeypatch): assert get_azure_ai_agent_entra_token({"azure_ad_token": "os.environ/FOUNDRY_AGENT_AD_TOKEN"}) == "token-from-env" +_AGENT_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_AGENT_CREDENTIALS: Final = { + "service_principal": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", + {"tenant_id": "tenant", "client_id": _AGENT_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + ), + "username_password": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password", + {"client_id": _AGENT_APP_CLIENT_ID, "azure_username": "agent@example.com", "azure_password": "wrong"}, + ), +} + + +def _raising(error: Exception): + def token() -> str: + raise error + + return token + + +@pytest.mark.parametrize(("token_fetch", "agent_params"), _AGENT_CREDENTIALS.values(), ids=_AGENT_CREDENTIALS.keys()) +def test_agent_entra_rejection_raises_401_with_only_the_aadsts_reason(token_fetch, agent_params): + from azure.core.exceptions import ClientAuthenticationError + + rejection: Final = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent " + f"in the request is the client secret value, for a secret added to app '{_AGENT_APP_CLIENT_ID}'." + ) + + with patch(token_fetch, return_value=_raising(rejection)): # test-quality-ok: stands in for the Entra token endpoint answering a wrong credential, the shape azure-identity raises + with pytest.raises(litellm.AuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert _AGENT_APP_CLIENT_ID not in str(excinfo.value) + + +def test_agent_entra_endpoint_unreachable_is_not_an_auth_error(): + from azure.core.exceptions import ClientAuthenticationError + + unreachable: Final = ClientAuthenticationError(message="Authentication failed: Connection refused") + token_fetch, agent_params = _AGENT_CREDENTIALS["service_principal"] + + with patch(token_fetch, return_value=_raising(unreachable)): # test-quality-ok: stands in for an unreachable Entra token endpoint, the shape azure-identity raises + with pytest.raises(ClientAuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + def test_agent_entra_token_failure_names_the_credential_fields(): with pytest.raises(ValueError, match="client_secret"): get_azure_ai_agent_entra_token({"azure_scope": "https://ai.azure.com/.default"}) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 88da07c612b..0aacc1f8bda 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1268,6 +1268,153 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): assert failure_data.get("agent_id") == "test-agent" +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" + + +class _JsonRpcError(BaseModel): + model_config = ConfigDict(frozen=True) + + code: int + message: str + + +class _JsonRpcErrorBody(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _JsonRpcError + + +async def _send_message_to_entra_agent( + litellm_params: Mapping[str, str], *patches: AbstractContextManager[object] +) -> tuple[int, _JsonRpcError]: + agent = _make_agent_mock() + agent.litellm_params = dict(litellm_params) + mock_request = _make_request_mock( + "message/send", + {"message": {"role": "user", "parts": [{"kind": "text", "text": "Hello"}], "messageId": "msg-1"}}, + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in (*_base_patches(agent), *patches): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + return response.status_code, _JsonRpcErrorBody.model_validate_json(response.body).error + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_credential_returns_401_with_only_the_aadsts_reason(): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{_ENTRA_APP_CLIENT_ID}'." + ) + + status_code, error = await _send_message_to_entra_agent( + {"tenant_id": "tenant", "client_id": _ENTRA_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + patch( # test-quality-ok: stands in for the Entra token endpoint answering a wrong secret, the shape azure-identity raises + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=rejected_token + ), + ) + + assert status_code == 401 + assert "AADSTS7000215" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_oidc_token_returns_401_with_only_the_aadsts_reason(monkeypatch, respx_mock): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 400, + json={ + "error": "invalid_client", + "error_description": "AADSTS700213: No matching federated identity record found for app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID}", + }, + ) + ) + + status_code, error = await _send_message_to_entra_agent( + { + "tenant_id": "qa-tenant", + "client_id": _ENTRA_APP_CLIENT_ID, + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + ) + + assert status_code == 401 + assert error.code == -32000 + assert "AADSTS700213" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + assert _ENTRA_TRACE_ID not in error.message + + +@pytest.mark.asyncio +async def test_upstream_authentication_error_without_an_entra_rejection_stays_an_internal_error(): + agent = _make_agent_mock() + mock_handler = MagicMock() + mock_handler.post = AsyncMock( + side_effect=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider="openai", model="gpt-4.1" + ) + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: stands in for the upstream agent rejecting the provider key it was sent + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=mock_handler + ) + ) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=_make_request_mock("tasks/get", {"id": "task-1"}), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + assert response.status_code == 500 + assert json.loads(response.body.decode())["error"]["code"] == -32603 + + @pytest.mark.asyncio async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400(): from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4a76aa77f14..11b07472eaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7241,6 +7241,178 @@ class TestRouterModelRelayUpstreamContract: assert json.loads(result.body) == upstream_body assert result.headers["x-ms-request-id"] == "req-1" + @pytest.mark.asyncio + async def test_azure_relay_answers_401_when_entra_rejects_the_deployment_credential(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + try: + raise EntraRejectionError(status_code=401, message=reason) + except EntraRejectionError as rejection: + raise BaseLLMException(status_code=401, message=str(rejection)) from rejection + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert (excinfo.value.code, excinfo.value.message) == ("401", reason) + + @pytest.mark.asyncio + async def test_azure_relay_entra_rejection_hides_an_internal_issuer_host(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = ( + "Failed to get Azure AD token: AADSTS700211: No matching federated identity record found for " + "presented assertion issuer 'https://oidc.build.internal:8443/realms/ci'." + ) + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + raise EntraRejectionError(status_code=401, message=reason) + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert excinfo.value.code == "401" + assert "AADSTS700211" in excinfo.value.message + assert "oidc.build.internal" not in excinfo.value.message + + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_failure_entra_did_not_cause(self, monkeypatch): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + upstream_failure: Final = BaseLLMException(status_code=500, message="Connection reset by Azure") + + class FailingRouter: + async def allm_passthrough_route(self, **kwargs): + raise upstream_failure + + self._install_router(monkeypatch, FailingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(BaseLLMException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert excinfo.value is upstream_failure + + @pytest.mark.asyncio + async def test_azure_relay_keeps_entra_throttling_the_oidc_token_exchange_as_a_429(self, monkeypatch): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + }, + } + ], + num_retries=0, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 429, + headers={"Retry-After": "7"}, + json={"error": "temporarily_unavailable", "error_codes": [90055]}, + ) + ) + deployment_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert (response.status_code, response.headers.get("retry-after")) == (429, "7") + assert not deployment_route.called + + + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_later_deployment_failure_after_an_entra_rejection( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from azure.core.exceptions import ClientAuthenticationError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + "order": 1, + }, + }, + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa-unreachable.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "configured-key", + "order": 2, + }, + }, + ], + num_retries=1, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + unreachable_route: Final = respx.post(url__startswith="https://litellm-qa-unreachable.openai.azure.com").mock( + side_effect=httpx.ConnectError("Connection refused") + ) + with pytest.raises(BaseLLMException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert unreachable_route.called + assert excinfo.value.status_code == 500 + NIM_INFER_BODY = { "input": [ diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 26dbb96c60a..e9dd91d4362 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -302,6 +302,63 @@ def _two_deployment_router() -> Router: ) +def _entra_fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +def _raised_from_an_entra_rejection(failed_deployment_id: str) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + raised.failed_deployment_id = failed_deployment_id + return raised + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_params", "cooled_down"), + [ + pytest.param({"client_secret": "not-the-credential"}, [], id="caller-secret"), + pytest.param({"client_secret": "configured-secret"}, ["fallback-dep-1"], id="deployment-secret"), + ], +) +async def test_fallback_hop_cools_down_only_a_deployment_entra_rejected_its_own_credential( + call_params: Mapping[str, object], cooled_down: list[str] +) -> None: + router: Final = _entra_fallback_router() + + _trigger_cooldown_for_failed_deployment( + ctx=RoutingContext(router), + kwargs={}, + exception=_raised_from_an_entra_rejection("fallback-dep-1"), + model_call_details={"litellm_params": call_params}, + ) + + assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == cooled_down + + @pytest.mark.asyncio @pytest.mark.parametrize("failure_already_logged", [False, True]) async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index f815b107648..c16c1c62202 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -25836,6 +25836,441 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} +def _caused_by_an_entra_rejection(error: litellm.AuthenticationError) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + try: + raise error from EntraRejectionError(status_code=401, message=error.message) + except litellm.AuthenticationError as raised: + return raised + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize("caller_credential", ["client_secret", "azure_password", "azure_ad_token"]) +async def test_router_caller_entra_credential_rejection_leaves_the_shared_deployment_routable(caller_credential): + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **{caller_credential: "not-the-credential"}, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_field", + [ + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + pytest.param({"client_id": "environment-client"}, id="environment_client_id"), + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"client_secret": ""}, id="blank_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_the_deployment_credential_leaves_it_routable( + caller_field, monkeypatch +): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + monkeypatch.setenv("AZURE_CLIENT_ID", "environment-client") + monkeypatch.delenv("AZURE_SCOPE", raising=False) + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + "azure_scope": "api://configured-app/.default", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_field, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("deployment_secret", "router_default_secret"), + [ + pytest.param({"client_secret": "configured-secret"}, None, id="deployment"), + pytest.param({}, {"client_secret": "configured-secret"}, id="router_default_litellm_params"), + ], +) +async def test_router_configured_entra_credential_rejection_cools_down_the_deployment( + deployment_secret, router_default_secret +): + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + **deployment_secret, + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params=router_default_secret, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000222: The provided client secret keys are expired.", + llm_provider="azure", + model="gpt-4.1", + ), + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_fields", + [ + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"azure_username": ""}, id="empty_username"), + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_nothing_still_cools_down_a_rejected_deployment( + caller_fields, monkeypatch +): + monkeypatch.delenv("AZURE_SCOPE", raising=False) + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_fields, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_caller_repeating_the_secret_an_empty_router_default_selects_still_cools_down(monkeypatch): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params={"client_secret": ""}, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + client_secret="environment-secret", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_caller_entra_field_keeps_the_shared_deployment(): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "configured-key"}, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + response: Final = await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + client_id="caller-client", + ) + + assert response._hidden_params["model_id"] == "shared-deployment" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("litellm_params", "provider"), + [ + pytest.param({"model": "openai/gpt-4.1-mini", "api_key": "revoked-key"}, "openai", id="openai"), + pytest.param( + { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "azure", + id="azure_api_key", + ), + ], +) +@pytest.mark.parametrize("caller_entra_field", ["client_id", "azure_scope"]) +async def test_router_auth_failure_entra_did_not_cause_cools_down_despite_a_caller_entra_field( + litellm_params, provider, caller_entra_field +): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": litellm_params, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider=provider, model=litellm_params["model"] + ), + **{caller_entra_field: "caller-value"}, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_fallback_auth_failure_after_an_entra_rejection_still_cools_down_the_fallback(monkeypatch): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + }, + "model_info": {"id": "entra-deployment"}, + }, + { + "model_name": "key-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa-fallback.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "model_info": {"id": "key-deployment"}, + }, + ], + fallbacks=[{"entra-model": ["key-model"]}], + ) + + with respx.mock: + fallback_route: Final = respx.post(url__startswith="https://litellm-qa-fallback.openai.azure.com").mock( + return_value=httpx.Response( + 401, json={"error": {"code": "401", "message": "Access denied due to invalid subscription key."}} + ) + ) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + client_secret="not-the-credential", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="key-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert fallback_route.call_count == 1 + assert excinfo.value.all_deployments_in_cooldown is True + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "stream_request", + [ + pytest.param( + lambda router, **caller_fields: router.acompletion( + model="entra-model", messages=[{"role": "user", "content": "hi"}], stream=True, **caller_fields + ), + id="chat_completions", + ), + pytest.param( + lambda router, **caller_fields: router.aanthropic_messages( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + **caller_fields, + ), + id="messages", + ), + ], +) +async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deployment_routable( + monkeypatch: pytest.MonkeyPatch, + stream_request: Callable[..., Awaitable[object]], +) -> None: + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + }, + "model_info": {"id": "entra-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router, client_secret="not-the-credential") + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router( From de85f59539fbe0729aa4de480b259853388a588c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 21:00:39 +0000 Subject: [PATCH 02/19] feat(cost-map): add baseten Kimi-K3.1-Pro (#50093) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ model_prices_and_context_window.json | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 09308ea2a60..767b455e7fb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -84541,5 +84541,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "baseten/moonshotai/Kimi-K3.1-Pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1000000, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 09308ea2a60..767b455e7fb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -84541,5 +84541,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "baseten/moonshotai/Kimi-K3.1-Pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1000000, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } From 81e8d652b08dc275ef50d8a8647a7cb615336d8e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 21:07:21 +0000 Subject: [PATCH 03/19] test(observability): reserve sink port throughout outage (#50100) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_arize_otel_v2_openinference_chaos.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py index 79b0d9084aa..0095080d944 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -1,12 +1,13 @@ from __future__ import annotations +import socket import threading import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from pathlib import Path +from queue import SimpleQueue from typing import Final -from urllib.parse import urlsplit import httpx import pytest @@ -44,7 +45,7 @@ from _openinference_support import ( ) from integration._support.client import Gateway, eventually from integration._support.process import graceful_stop_seconds -from integration._support.wire import Reply, Request, wire_server +from integration._support.wire import Reply, Request, Wire, wire_server def _call( @@ -249,11 +250,12 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P _assert_response(response, marker, surface, stream, model_name) with ExitStack() as servers: - initial_stack: Final = servers.enter_context(ExitStack()) - stopped_destination: Final = initial_stack.enter_context(wire_server(_owned_sink_handler(sink))) - sink_port: Final = urlsplit(stopped_destination.url).port - assert sink_port is not None, stopped_destination.url - initial_stack.close() + outage_socket: Final = servers.enter_context(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) + outage_socket.bind(("127.0.0.1", 0)) # test-quality-ok: socket stays bound throughout the outage + sink_port: Final = outage_socket.getsockname()[1] + stopped_destination: Final = Wire( + f"http://127.0.0.1:{sink_port}", SimpleQueue(), SimpleQueue(), SimpleQueue(), SimpleQueue() + ) with _rig( gateway, tmp_path, @@ -272,6 +274,7 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P with httpx.Client(trust_env=False) as client, pytest.raises(httpx.ConnectError): client.get(stopped_destination.url + "/health", timeout=2) recovered_stack: Final = servers.enter_context(ExitStack()) + outage_socket.close() recovered_destination: Final = recovered_stack.enter_context( wire_server(_owned_sink_handler(sink), port=sink_port) ) From e39caf7bf89c7c755a895d95f21f22d0c5da8c3b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:22:21 -0800 Subject: [PATCH 04/19] chore(databricks): add retirement successors and kimi-k2-7 retirement date from the Databricks retired models page (#50094) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 13 ++++++++++++- model_prices_and_context_window.json | 13 ++++++++++++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 767b455e7fb..dfd33601e8c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22047,6 +22047,7 @@ "output_cost_per_token": 2.8e-07, "output_dbu_cost_per_token": 4e-06, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -22060,8 +22061,12 @@ "supports_vision": false }, "databricks/system.ai.deepseek-v4-flash-0731": { + "deprecation_date": "2026-11-05", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "litellm_provider": "databricks", "mode": "chat", + "source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { @@ -22081,6 +22086,7 @@ "output_cost_per_token": 3.95997e-06, "output_dbu_cost_per_token": 5.6571e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -23158,6 +23164,7 @@ "output_cost_per_token": 4.04999e-06, "output_dbu_cost_per_token": 5.7857e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-glm-5-3", "supported_modalities": [ "text", "image" @@ -81729,7 +81736,8 @@ "input_cost_per_token_priority": 5.25e-07, "output_cost_per_token_priority": 2.1e-06, "cache_read_input_token_cost_priority": 5.25e-08, - "cache_creation_input_token_cost_priority": 5.25e-07 + "cache_creation_input_token_cost_priority": 5.25e-07, + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017" }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -82181,6 +82189,8 @@ "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, + "deprecation_date": "2026-10-30", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "input_cost_per_token": 9.4997e-07, "input_dbu_cost_per_token": 1.3571e-05, "litellm_provider": "databricks", @@ -82193,6 +82203,7 @@ "output_cost_per_token": 4.00001e-06, "output_dbu_cost_per_token": 5.7143e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "successor_model": "databricks/databricks-kimi-k3", "supported_modalities": [ "text", "image" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 767b455e7fb..dfd33601e8c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22047,6 +22047,7 @@ "output_cost_per_token": 2.8e-07, "output_dbu_cost_per_token": 4e-06, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -22060,8 +22061,12 @@ "supports_vision": false }, "databricks/system.ai.deepseek-v4-flash-0731": { + "deprecation_date": "2026-11-05", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "litellm_provider": "databricks", "mode": "chat", + "source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { @@ -22081,6 +22086,7 @@ "output_cost_per_token": 3.95997e-06, "output_dbu_cost_per_token": 5.6571e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -23158,6 +23164,7 @@ "output_cost_per_token": 4.04999e-06, "output_dbu_cost_per_token": 5.7857e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-glm-5-3", "supported_modalities": [ "text", "image" @@ -81729,7 +81736,8 @@ "input_cost_per_token_priority": 5.25e-07, "output_cost_per_token_priority": 2.1e-06, "cache_read_input_token_cost_priority": 5.25e-08, - "cache_creation_input_token_cost_priority": 5.25e-07 + "cache_creation_input_token_cost_priority": 5.25e-07, + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017" }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -82181,6 +82189,8 @@ "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, + "deprecation_date": "2026-10-30", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "input_cost_per_token": 9.4997e-07, "input_dbu_cost_per_token": 1.3571e-05, "litellm_provider": "databricks", @@ -82193,6 +82203,7 @@ "output_cost_per_token": 4.00001e-06, "output_dbu_cost_per_token": 5.7143e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "successor_model": "databricks/databricks-kimi-k3", "supported_modalities": [ "text", "image" From 1967b41368833ea4c9a150a74e4c71df7246060d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:23:26 -0800 Subject: [PATCH 05/19] chore(cost-map): openai context fixes, azure gpt-6 flex pricing, bedrock govcloud and gmi rows (#50082) * chore(cost-map): set openai gpt-5.6-cyber and Daybreak Red max input to 272K and source the sora-2-pro-high-res deprecation * Absorb GovCloud Bedrock, Azure GPT-6 Flex and GMI registry rows Roll the verified registry data and Azure Flex cost coverage into the rolling price sync Co-authored-by: kerry-berri <326648561+kerry-berri@users.noreply.github.com> Co-authored-by: dguthruf <102602010+danielguthruf@users.noreply.github.com> Co-authored-by: hisamotodev Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: hisamotodev --- ...odel_prices_and_context_window_backup.json | 1320 ++++++++++++++++- model_prices_and_context_window.json | 1320 ++++++++++++++++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 31 + 3 files changed, 2545 insertions(+), 126 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dfd33601e8c..2ebf3dd0ae1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9003,11 +9003,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9039,7 +9035,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -9055,11 +9055,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9091,7 +9087,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -9107,12 +9107,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9139,7 +9135,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, @@ -9155,12 +9159,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9187,7 +9187,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -9239,7 +9247,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6-sol-2026-09-22": { "cache_creation_input_token_cost": 2.5e-06, @@ -9291,7 +9307,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, @@ -9354,11 +9378,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9386,7 +9406,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, @@ -9402,11 +9430,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9434,7 +9458,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/us/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, @@ -15566,7 +15598,7 @@ "output_vector_size": 1536 }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 3e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "mode": "embedding", @@ -38133,7 +38165,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38172,7 +38204,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38211,7 +38243,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token": 7.5e-05, @@ -55968,6 +56000,7 @@ "openai/sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://platform.openai.com/docs/api-reference/videos", @@ -59830,6 +59863,7 @@ "sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://developers.openai.com/api/docs/pricing", @@ -69671,14 +69705,14 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-06, - "input_cost_per_token_above_272k_tokens": 5.4e-06, - "cache_creation_input_token_cost": 3.375e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-06, - "cache_read_input_token_cost": 2.7e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, - "output_cost_per_token": 1.62e-05, - "output_cost_per_token_above_272k_tokens": 2.43e-05, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", "supports_web_search": true @@ -69708,14 +69742,14 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-07, - "input_cost_per_token_above_272k_tokens": 5.4e-07, - "cache_creation_input_token_cost": 3.375e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-07, - "cache_read_input_token_cost": 2.7e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, - "output_cost_per_token": 1.62e-06, - "output_cost_per_token_above_272k_tokens": 2.43e-06, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", "supports_web_search": true @@ -69744,9 +69778,9 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 3.375e-06, - "cache_read_input_token_cost": 3.375e-07, - "output_cost_per_token": 2.025e-05, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", "supports_web_search": true @@ -84529,6 +84563,1166 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gmi/MiniMaxAI/MiniMax-M2": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.5": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.7": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3": { + "cache_read_input_token_cost": 4.8e-08, + "input_cost_per_token": 2.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9.6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.125e-08, + "input_cost_per_token": 1.125e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-FP8": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.09e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Instruct-2507-FP8": { + "cache_read_input_token_cost": 1.75e-08, + "input_cost_per_token": 8.75e-08, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Thinking-2507-FP8": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262141, + "mode": "chat", + "output_cost_per_token": 4.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-30B-A3B": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 40946, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-32B-FP8": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 9e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262128, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-122B-A10B": { + "input_cost_per_token": 4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-27B": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-35B-A3B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-35B-A3B": { + "input_cost_per_token": 2.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.485e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus-2026-04-02": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.7-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-27B": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Flash": { + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 1.6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max-0902": { + "cache_read_input_token_cost": 2.06e-07, + "input_cost_per_token": 1.65e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.951e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5": { + "cache_read_input_token_cost": 2.38e-09, + "input_cost_per_token": 1.19e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.38e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5-Pro": { + "cache_read_input_token_cost": 2.52e-09, + "input_cost_per_token": 3.045e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.09e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-0528": { + "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1-Terminus": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Exp": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4.1e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Speciale": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 1.82e-08, + "input_cost_per_token": 9.1e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.82e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 9.1e-09, + "input_cost_per_token": 2.86e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.58e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Pro-0813": { + "cache_read_input_token_cost": 3.52e-08, + "input_cost_per_token": 1.056e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.168e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-pro-preview": { + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1.3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2-Instruct-0905": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.6": { + "cache_read_input_token_cost": 1.44e-07, + "input_cost_per_token": 8.55e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code-highspeed": { + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3.1-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/nemotron-3-ultra-550b-a55b": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-pro": { + "input_cost_per_token": 3e-05, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 0.00018, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-luna": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-sol": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-terra": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-astra": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-sol": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6.1-sol": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-oss-120b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-3.7-Flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.15e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-5-Preview": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.7e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/HunyuanOCR": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/Hy3": { + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 5.8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy3-preview": { + "cache_read_input_token_cost": 6e-08, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 8.34e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/x-ai/grok-4.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/xai/grok4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-Air-FP8": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131076, + "mode": "chat", + "output_cost_per_token": 1.1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-FP8": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.6": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 204800, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.7-Flash": { + "input_cost_per_token": 7e-08, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5-FP8": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 1.92e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.1-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.2-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "fireworks_ai/accounts/fireworks/models/kimi-k3p1-pro": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_priority": 3e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dfd33601e8c..2ebf3dd0ae1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9003,11 +9003,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9039,7 +9035,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -9055,11 +9055,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9091,7 +9087,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -9107,12 +9107,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9139,7 +9135,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, @@ -9155,12 +9159,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9187,7 +9187,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -9239,7 +9247,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6-sol-2026-09-22": { "cache_creation_input_token_cost": 2.5e-06, @@ -9291,7 +9307,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, @@ -9354,11 +9378,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9386,7 +9406,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, @@ -9402,11 +9430,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9434,7 +9458,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/us/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, @@ -15566,7 +15598,7 @@ "output_vector_size": 1536 }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 3e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "mode": "embedding", @@ -38133,7 +38165,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38172,7 +38204,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38211,7 +38243,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token": 7.5e-05, @@ -55968,6 +56000,7 @@ "openai/sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://platform.openai.com/docs/api-reference/videos", @@ -59830,6 +59863,7 @@ "sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://developers.openai.com/api/docs/pricing", @@ -69671,14 +69705,14 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-06, - "input_cost_per_token_above_272k_tokens": 5.4e-06, - "cache_creation_input_token_cost": 3.375e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-06, - "cache_read_input_token_cost": 2.7e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, - "output_cost_per_token": 1.62e-05, - "output_cost_per_token_above_272k_tokens": 2.43e-05, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", "supports_web_search": true @@ -69708,14 +69742,14 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-07, - "input_cost_per_token_above_272k_tokens": 5.4e-07, - "cache_creation_input_token_cost": 3.375e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-07, - "cache_read_input_token_cost": 2.7e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, - "output_cost_per_token": 1.62e-06, - "output_cost_per_token_above_272k_tokens": 2.43e-06, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", "supports_web_search": true @@ -69744,9 +69778,9 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 3.375e-06, - "cache_read_input_token_cost": 3.375e-07, - "output_cost_per_token": 2.025e-05, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", "supports_web_search": true @@ -84529,6 +84563,1166 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gmi/MiniMaxAI/MiniMax-M2": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.5": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.7": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3": { + "cache_read_input_token_cost": 4.8e-08, + "input_cost_per_token": 2.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9.6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.125e-08, + "input_cost_per_token": 1.125e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-FP8": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.09e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Instruct-2507-FP8": { + "cache_read_input_token_cost": 1.75e-08, + "input_cost_per_token": 8.75e-08, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Thinking-2507-FP8": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262141, + "mode": "chat", + "output_cost_per_token": 4.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-30B-A3B": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 40946, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-32B-FP8": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 9e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262128, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-122B-A10B": { + "input_cost_per_token": 4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-27B": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-35B-A3B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-35B-A3B": { + "input_cost_per_token": 2.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.485e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus-2026-04-02": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.7-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-27B": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Flash": { + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 1.6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max-0902": { + "cache_read_input_token_cost": 2.06e-07, + "input_cost_per_token": 1.65e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.951e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5": { + "cache_read_input_token_cost": 2.38e-09, + "input_cost_per_token": 1.19e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.38e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5-Pro": { + "cache_read_input_token_cost": 2.52e-09, + "input_cost_per_token": 3.045e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.09e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-0528": { + "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1-Terminus": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Exp": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4.1e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Speciale": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 1.82e-08, + "input_cost_per_token": 9.1e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.82e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 9.1e-09, + "input_cost_per_token": 2.86e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.58e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Pro-0813": { + "cache_read_input_token_cost": 3.52e-08, + "input_cost_per_token": 1.056e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.168e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-pro-preview": { + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1.3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2-Instruct-0905": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.6": { + "cache_read_input_token_cost": 1.44e-07, + "input_cost_per_token": 8.55e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code-highspeed": { + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3.1-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/nemotron-3-ultra-550b-a55b": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-pro": { + "input_cost_per_token": 3e-05, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 0.00018, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-luna": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-sol": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-terra": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-astra": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-sol": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6.1-sol": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-oss-120b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-3.7-Flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.15e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-5-Preview": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.7e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/HunyuanOCR": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/Hy3": { + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 5.8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy3-preview": { + "cache_read_input_token_cost": 6e-08, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 8.34e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/x-ai/grok-4.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/xai/grok4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-Air-FP8": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131076, + "mode": "chat", + "output_cost_per_token": 1.1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-FP8": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.6": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 204800, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.7-Flash": { + "input_cost_per_token": 7e-08, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5-FP8": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 1.92e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.1-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.2-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "fireworks_ai/accounts/fireworks/models/kimi-k3p1-pro": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_priority": 3e-07, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 609b28270d7..3964f472bd9 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4310,6 +4310,37 @@ def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base): assert azure_eu_info[field] == pytest.approx(1.2 * base) +@pytest.mark.parametrize("model", ["azure/gpt-6-luna", "azure/gpt-6-luna-2026-09-22"]) +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "cache_read_rate", "output_rate"), + [ + (100_000, 5e-08, 5e-09, 2.5e-07), + (300_000, 1e-07, 1e-08, 3.75e-07), + ], +) +def test_generic_cost_per_token_azure_gpt_6_luna_flex_tiers( + _local_model_cost_map, model, prompt_tokens, input_rate, cache_read_rate, output_rate +): + """azure/gpt-6-luna bills Azure's Flex Global rates, half of Standard, with the long context step past 272K.""" + cached_tokens = 20_000 + completion_tokens = 1_000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="azure", + service_tier="flex", + ) + expected_prompt = (prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate + assert prompt_cost == pytest.approx(expected_prompt) + assert completion_cost == pytest.approx(completion_tokens * output_rate) + + @pytest.mark.parametrize("region_prefix", ["azure/", "azure/us/", "azure/eu/"]) def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_prefix): """The bare gpt-5.6 alias routes to GPT-5.6 Sol, so every Azure region must bill the From 9703e01cd4704b3b6a850efaee3f15d314403ce8 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Wed, 4 Nov 2026 13:55:50 -0800 Subject: [PATCH 06/19] fix(responses): keep the system prompt in input for json_object on the chat to responses bridge (#49957) The bridge moves leading system messages into the Responses instructions field. With response_format json_object, Azure requires the word json in input and ignores instructions, so a request that follows the documented rule 400s after translation. Send the system text as a system input item when json_object is requested, the same shape already used for system-only requests Resolves LIT-10426 --- .../transformation.py | 18 ++++---- ...responses_transformation_transformation.py | 45 +++++++++++++++++++ 2 files changed, 55 insertions(+), 8 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 8dd4d87ae13..8e6341805db 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1180,11 +1180,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): include_namespace=True, replay_openai_history=provider == "openai", ) - # OpenAI's Responses API rejects an empty input. For a system-only - # request, carry the system message as a system-role input item instead - # of instructions, mirroring how non-string system content is already - # handled in convert_chat_completion_messages_to_responses_api. - is_system_only_request: Final = not converted_input_items and converted_instructions is not None + # Keep the system text in input instead of instructions when the Responses API needs it there. + # It rejects an empty input, and json_object requires the word "json" in input, ignoring instructions. + typed_optional_params: Final = cast("Mapping[str, object]", optional_params) # cast-ok: untyped dict + system_goes_in_input: Final = converted_instructions is not None and ( + not converted_input_items or _payload_type(typed_optional_params.get("response_format")) == "json_object" + ) target_input_items: Final = ( _without_audio_input_parts(converted_input_items) if _drops_audio_input(model, litellm_params) @@ -1196,12 +1197,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": "system", "content": [{"type": "input_text", "text": converted_instructions}], - } + }, + *target_input_items, ] - if is_system_only_request + if system_goes_in_input else target_input_items ) - instructions: Final = None if is_system_only_request else converted_instructions + instructions: Final = None if system_goes_in_input else converted_instructions # Build responses API request using the reverse transformation logic responses_api_request: Final = ResponsesAPIOptionalRequestParams() diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 04ba823d08f..19aefb430c7 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1856,6 +1856,51 @@ def test_transform_request_system_only_message_maps_to_system_input_item(): assert not result.get("instructions") +_JSON_SYSTEM_PROMPT: Final = "You output JSON objects." +_PLAIN_USER_ITEM: Final = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "List three colors."}], +} +_COLORS_SCHEMA: Final = { + "type": "json_schema", + "json_schema": {"name": "colors", "schema": {"type": "object", "properties": {}}}, +} + + +@pytest.mark.parametrize( + ("response_format", "expected_input", "expected_instructions"), + [ + ( + {"type": "json_object"}, + [ + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": _JSON_SYSTEM_PROMPT}]}, + _PLAIN_USER_ITEM, + ], + None, + ), + (_COLORS_SCHEMA, [_PLAIN_USER_ITEM], _JSON_SYSTEM_PROMPT), + ], + ids=["json_object_keeps_system_in_input", "json_schema_keeps_system_in_instructions"], +) +def test_transform_request_json_object_keeps_system_prompt_where_the_json_word_check_reads_it( + response_format, expected_input, expected_instructions +): + result = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.4-mini", + messages=[ + {"role": "system", "content": _JSON_SYSTEM_PROMPT}, + {"role": "user", "content": "List three colors."}, + ], + optional_params={"response_format": response_format}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert (result["input"], result.get("instructions")) == (expected_input, expected_instructions), result + + def test_transform_request_single_char_keys_not_matched(): """Test that single-character keys are not incorrectly matched to 'metadata' or 'previous_response_id' From 732a50318b8bddc6cf79f2c21b31bcd63474b6d6 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 4 Nov 2026 14:08:40 -0800 Subject: [PATCH 07/19] feat(proxy): authenticate billing metrics with a cognito bearer token (#49852) * feat(billing): support Cognito bearer auth for metrics Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refetch billing token on 401 and require https Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): import Unpack from typing_extensions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): make billing metrics bearer only and drop mtls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format billing metrics module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm-helm/Chart.yaml | 2 +- helm/litellm-helm/README.md | 11 +- helm/litellm-helm/templates/_helpers.tpl | 50 +- helm/litellm-helm/templates/deployment.yaml | 6 - .../tests/billing_metrics_tests.yaml | 283 ++----- helm/litellm-helm/tests/collector_tests.yaml | 6 +- helm/litellm-helm/values.yaml | 16 +- helm/litellm/Chart.yaml | 2 +- helm/litellm/README.md | 17 + helm/litellm/templates/_helpers.tpl | 53 +- .../litellm/templates/backend/deployment.yaml | 10 +- .../litellm/templates/gateway/deployment.yaml | 10 +- helm/litellm/tests/billing_metrics_tests.yaml | 224 ++---- helm/litellm/tests/migration_job_tests.yaml | 2 +- helm/litellm/values.yaml | 21 +- litellm/litellm_core_utils/url_utils.py | 8 + .../claude_code_marketplace.py | 12 +- .../enterprise_billing/billing_metrics.py | 282 ++++--- terraform/litellm/aws/README.md | 28 +- terraform/litellm/aws/ecs.tf | 60 +- terraform/litellm/aws/iam.tf | 4 +- terraform/litellm/aws/secrets.tf | 54 +- terraform/litellm/aws/variables.tf | 54 +- terraform/litellm/gcp/README.md | 28 +- terraform/litellm/gcp/cloudrun.tf | 37 +- terraform/litellm/gcp/iam.tf | 28 +- terraform/litellm/gcp/secrets.tf | 52 +- terraform/litellm/gcp/variables.tf | 55 +- .../unit/litellm_core_utils/test_url_utils.py | 14 + .../test_billing_metrics.py | 707 +++++++++--------- 30 files changed, 773 insertions(+), 1363 deletions(-) diff --git a/helm/litellm-helm/Chart.yaml b/helm/litellm-helm/Chart.yaml index a3cb388ffc6..059889cb667 100644 --- a/helm/litellm-helm/Chart.yaml +++ b/helm/litellm-helm/Chart.yaml @@ -18,7 +18,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.3 +version: 1.1.4 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index 694cca05b52..bc6e29632f1 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented): | `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` | | `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` | -| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` | -| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` | +| `billingMetrics.tokenUrl` | Cognito OAuth2 client-credentials token endpoint. | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` | +| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys. | `litellm-billing-metrics` | | `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` | #### Example `proxy_config` ConfigMap from values (default): @@ -109,10 +109,10 @@ type: Opaque #### Enterprise billable-request metering -Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on: +Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values: ``` -kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key +kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=... ``` ``` @@ -120,7 +120,8 @@ billingMetrics: enabled: true ``` -Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time. +The endpoint and token URL must use HTTPS. The chart requires both URLs and +the Secret name when billing metrics are enabled ### Database Settings diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 929fc415163..230d2acaeef 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} {{- end }} -{{/* -Enterprise billable-request metering. The client certificate identifies the -deployment to LiteLLM's collector, so it is mounted read-only from an existing -Secret rather than passed through the environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - {{- define "litellm.billingMetricsEnv" -}} - name: LITELLM_BILLING_METRICS_ENDPOINT value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} +- name: LITELLM_BILLING_METRICS_TOKEN_URL + value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-id +- name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-secret {{- with .Values.billingMetrics.exportIntervalMs }} - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: {{ . | quote }} {{- end }} {{- end -}} -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} - {{/* Create the name of the service account to use */}} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 37e7eaec9e9..a37193f0efa 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -148,9 +148,6 @@ spec: - name: npm mountPath: /.npm {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- if include "litellm.collector.socketDir" . }} - name: collector-socket mountPath: {{ include "litellm.collector.socketDir" . }} @@ -234,9 +231,6 @@ spec: items: - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} path: "config.yaml" - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- if include "litellm.collector.socketDir" . }} - name: collector-socket emptyDir: diff --git a/helm/litellm-helm/tests/billing_metrics_tests.yaml b/helm/litellm-helm/tests/billing_metrics_tests.yaml index 71803df378c..52185c8dd9d 100644 --- a/helm/litellm-helm/tests/billing_metrics_tests.yaml +++ b/helm/litellm-helm/tests/billing_metrics_tests.yaml @@ -1,31 +1,21 @@ -suite: test billingMetrics wiring on the proxy deployment +suite: test bearer-only billingMetrics wiring templates: - deployment.yaml - configmap-litellm.yaml - - migrations-job.yaml tests: - - it: is off by default, adding no env, volume, or mount + - it: is disabled by default template: deployment.yaml asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - notContains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - notContains: path: spec.template.spec.containers[0].env content: name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics - - it: renders the endpoint and the mounted cert paths when enabled + - it: renders required URLs and Cognito secret references when enabled template: deployment.yaml set: billingMetrics: @@ -39,119 +29,34 @@ tests: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt + name: LITELLM_BILLING_METRICS_TOKEN_URL + value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - # The conventional Secret name is the default, so enabling the block is enough. - - it: mounts the default cert secret read-only alongside the config volume - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: honours a secretName override - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: my-billing-mtls - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: my-billing-mtls - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - - it: honours an endpoint override - template: deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: https://collector.internal:4318 - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-id - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://collector.internal:4318 - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca + name: billing-metrics - notContains: path: spec.template.spec.containers[0].volumeMounts content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt + name: billing-metrics - - it: mounts the CA secret when caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - caSecretName: billing-ca - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - - it: passes the export interval through only when set + - it: renders the export interval only when set template: deployment.yaml set: billingMetrics: @@ -164,129 +69,7 @@ tests: name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: "5000" - - it: omits the export interval when unset - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS - value: "60000" - - # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing - # entries must render after .Values.envVars or a user could silently redirect - # the metering export. The three billing entries are the last ones emitted here - # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case). - - it: renders the billing endpoint after envVars so it cannot be shadowed - template: deployment.yaml - set: - migrationJob: - enabled: false - billingMetrics: - enabled: true - envVars: - LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://shadowed.example - - equal: - path: spec.template.spec.containers[0].env[-3] - value: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - equal: - path: spec.template.spec.containers[0].env[-2].name - value: LITELLM_BILLING_METRICS_CLIENT_CERT - - equal: - path: spec.template.spec.containers[0].env[-1].name - value: LITELLM_BILLING_METRICS_CLIENT_KEY - - - it: keeps user-supplied volumes and mounts alongside the billing secret - template: deployment.yaml - set: - billingMetrics: - enabled: true - volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: custom-callbacks - mountPath: /app/callbacks - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: still mounts the proxy config when enabled - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - - # Only the proxy serves billable traffic. The migrations Job must never mount - # the client certificate, and it renders its own env and volumes, so nothing - # stops a future edit from wiring the billing include into it by mistake. - - it: does not touch the migrations job when enabled - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - notExists: - path: spec.template.spec.containers[0].volumeMounts - - notExists: - path: spec.template.spec.volumes - - - it: fails loudly when enabled with an emptied secretName - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: "" - asserts: - - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint + - it: requires the endpoint when enabled template: deployment.yaml set: billingMetrics: @@ -295,3 +78,23 @@ tests: asserts: - failedTemplate: errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true + + - it: requires the token URL when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + tokenUrl: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true + + - it: requires the secret name when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + secretName: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true diff --git a/helm/litellm-helm/tests/collector_tests.yaml b/helm/litellm-helm/tests/collector_tests.yaml index 0340b1161b7..1365290e16c 100644 --- a/helm/litellm-helm/tests/collector_tests.yaml +++ b/helm/litellm-helm/tests/collector_tests.yaml @@ -191,7 +191,7 @@ tests: metricsServer.port: 9090 billingMetrics.enabled: true billingMetrics.endpoint: https://metering.example.com - billingMetrics.secretName: billing-mtls + billingMetrics.secretName: billing-metrics asserts: - contains: path: spec.template.spec.containers[0].env @@ -214,9 +214,9 @@ tests: name: LITELLM_BILLING_METRICS_ENDPOINT any: true - notContains: - path: spec.template.spec.containers[1].volumeMounts + path: spec.template.spec.containers[1].env content: - name: billing-metrics-mtls + name: LITELLM_BILLING_METRICS_CLIENT_SECRET any: true - it: should give the sidecar the same scratch mounts as the proxy on a read-only root diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index f51f8bdd173..9afce050261 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -154,19 +154,13 @@ masterkeySecretName: "" # if set, use this secret key for the master key; otherwise, use the default key masterkeySecretKey: "" -# Optional: enterprise billable-request metering. When enabled, the proxy counts -# successful requests to inference, MCP, and A2A endpoints and pushes them to -# LiteLLM's collector over mutual TLS. Requires an enterprise license. -# The client certificate identifies the deployment, so it is mounted read-only -# from an existing Secret and never passed through the environment. +# Optional: enterprise billable-request metering. Requires an enterprise license. billingMetrics: enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + endpoint: https://telemetry.litellm.ai + tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token + secretName: litellm-billing-metrics + exportIntervalMs: "" proxyConfigMap: # when true, creates a new configmap diff --git a/helm/litellm/Chart.yaml b/helm/litellm/Chart.yaml index e67f5790c7e..341172da925 100644 --- a/helm/litellm/Chart.yaml +++ b/helm/litellm/Chart.yaml @@ -2,7 +2,7 @@ apiVersion: v2 name: litellm description: LiteLLM componentized — gateway, UI backend, and UI as separate services type: application -version: 0.1.0 +version: 0.1.1 appVersion: "0.1.0" annotations: org.opencontainers.image.source: "https://github.com/BerriAI/litellm" diff --git a/helm/litellm/README.md b/helm/litellm/README.md index 97b94754402..d66bb3284bf 100644 --- a/helm/litellm/README.md +++ b/helm/litellm/README.md @@ -22,6 +22,23 @@ backend: gatewayUrl: https://gateway.example.com ``` +## Enterprise billing metrics + +Billing metrics use Cognito bearer authentication. Create the Secret with +`client-id` and `client-secret` keys, then enable billing metrics: + +```bash +kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=... +``` + +```yaml +billingMetrics: + enabled: true +``` + +The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name +when billing metrics are enabled + ## Routes to another Service For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 18e938564e9..b704782f5d1 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }} helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }} {{- end -}} -{{/* -Enterprise billable-request metering. Wired into gateway and backend, not the -migrations job. The gateway serves nearly all billable traffic, but the backend -keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a -SpendLogs row, so metering only the gateway would silently drop that traffic. -The client certificate identifies the deployment to LiteLLM's collector, so it is -mounted read-only from an existing Secret rather than passed through the -environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - {{- define "litellm.billingMetricsEnv" -}} - name: LITELLM_BILLING_METRICS_ENDPOINT value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} +- name: LITELLM_BILLING_METRICS_TOKEN_URL + value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-id +- name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-secret {{- with .Values.billingMetrics.exportIntervalMs }} - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: {{ . | quote }} {{- end }} {{- end -}} -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} {{/* Per-component selector labels — used in both Service selectors and Deployment matchLabels. diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 7360261dab1..7db14ad87a5 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -79,16 +79,13 @@ spec: {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} {{- include "litellm.envFrom" .Values.backend | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- with .Values.backend.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -114,16 +111,13 @@ spec: {{- with .Values.backend.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumes }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- with .Values.backend.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index b0e7f0dcebb..987109e361e 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -86,7 +86,7 @@ spec: {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config @@ -101,9 +101,6 @@ spec: - name: collector-socket mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- with .Values.gateway.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -206,7 +203,7 @@ spec: {{- with .Values.gateway.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config @@ -222,9 +219,6 @@ spec: emptyDir: sizeLimit: 1Mi {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- with .Values.gateway.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/tests/billing_metrics_tests.yaml b/helm/litellm/tests/billing_metrics_tests.yaml index ceba0bd1430..d17b8192fce 100644 --- a/helm/litellm/tests/billing_metrics_tests.yaml +++ b/helm/litellm/tests/billing_metrics_tests.yaml @@ -1,39 +1,30 @@ -suite: test billingMetrics wiring on gateway and backend +suite: test bearer-only billingMetrics wiring templates: - gateway/deployment.yaml - gateway/configmap.yaml - backend/deployment.yaml + - backend/configmap.yaml - migrations-job.yaml values: - ./values/required.yaml tests: - - it: is off by default, adding no env, volume, or mount + - it: is disabled by default template: gateway/deployment.yaml asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - equal: - path: spec.template.spec.containers[0].volumeMounts - value: - - name: gateway-config - mountPath: /app/config/config.yaml - subPath: config.yaml - notContains: path: spec.template.spec.containers[0].env content: name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics - - it: renders the endpoint and the mounted cert paths when enabled + - it: wires gateway URLs and bearer credentials from Secret keys template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls asserts: - contains: path: spec.template.spec.containers[0].env @@ -43,87 +34,68 @@ tests: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt + name: LITELLM_BILLING_METRICS_TOKEN_URL + value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - - it: mounts the cert secret read-only alongside the config volume - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-id - contains: - path: spec.template.spec.volumes + path: spec.template.spec.containers[0].env content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca + name: billing-metrics - notContains: - path: spec.template.spec.containers[0].env + path: spec.template.spec.containers[0].volumeMounts content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt + name: billing-metrics - - it: mounts the CA secret when caSecretName is set - template: gateway/deployment.yaml + - it: wires billing metrics on the backend without volumes + template: backend/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - caSecretName: billing-ca asserts: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret + - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true + name: billing-metrics - - it: passes the export interval through only when set + - it: does not add billing metrics to migrations + template: migrations-job.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + + - it: renders an optional export interval template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls exportIntervalMs: 5000 asserts: - contains: @@ -132,102 +104,27 @@ tests: name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: "5000" - - it: keeps user-supplied gateway volumes alongside the billing secret + - it: requires the endpoint when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - gateway.volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - gateway.volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks + endpoint: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls + - failedTemplate: + errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true - # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp), - # which writes a SpendLogs row, so it must meter too or that traffic is lost. - - it: meters the backend as well, since it serves the MCP transport - template: backend/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - - it: leaves the backend alone when metering is off - template: backend/deployment.yaml - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - # The migrations job runs prisma and serves no traffic; it must never receive - # the client key. - - it: never mounts the billing cert on the migrations job - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - isNull: - path: spec.template.spec.volumes - - # The conventional Secret name is the default, so enabling metering needs no - # secretName at all; the guard below only fires on an explicitly blanked one. - - it: uses the conventional secret name by default + - it: requires the token URL when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true + tokenUrl: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls + - failedTemplate: + errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true - - it: fails loudly when the secretName is explicitly blanked + - it: requires the secret name when enabled template: gateway/deployment.yaml set: billingMetrics: @@ -235,15 +132,4 @@ tests: secretName: "" asserts: - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: "" - secretName: billing-mtls - asserts: - - failedTemplate: - errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml index 2ebb1b44926..b9e52ba7123 100644 --- a/helm/litellm/tests/migration_job_tests.yaml +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -91,7 +91,7 @@ tests: app.kubernetes.io/name: litellm app.kubernetes.io/instance: RELEASE-NAME app.kubernetes.io/managed-by: Helm - helm.sh/chart: litellm-0.1.0 + helm.sh/chart: litellm-0.1.1 app.kubernetes.io/component: migrations - it: renders pod-level and container-level securityContext in their own scopes diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 92a48153c5a..cb37656cb77 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -180,23 +180,14 @@ masterKey: secretKey: master-key # Optional: enterprise billable-request metering. When enabled, the gateway and -# backend count successful requests to inference, MCP, and A2A endpoints and push -# them to LiteLLM's collector over mutual TLS. Both components serve billable -# routes: the backend keeps the named-server MCP transport. Requires an -# enterprise license. The client certificate identifies the deployment, so it is -# mounted read-only from an existing Secret and never passed through the env. +# backend count successful requests to inference, MCP, and A2A endpoints. Requires +# an enterprise license. billingMetrics: enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - # An existing Secret holding the client certificate under tls.crt and its key - # under tls.key, usually created from the onboarding artifact. The default is - # the conventional name, so the common path is to create that Secret and set - # enabled: true. Override only if yours is named differently. - secretName: litellm-billing-metrics-mtls - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + endpoint: https://telemetry.litellm.ai + tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token + secretName: litellm-billing-metrics + exportIntervalMs: "" # External Postgres connection. database: diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0db6a41b503..eae8ac04e8e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -102,6 +102,14 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +def is_https_url_with_host(url: str) -> bool: + try: + parsed_url: Final = urlsplit(url) + return parsed_url.scheme == "https" and bool(parsed_url.hostname) + except ValueError: + return False + + def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 5df73230b6c..b8c90fe7e93 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -26,7 +26,6 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args -from urllib.parse import urlsplit from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from fastapi.responses import JSONResponse, Response @@ -34,6 +33,7 @@ from pydantic import Json, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import is_https_url_with_host from litellm.models.plugin_upload import PluginUpload from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( @@ -377,14 +377,6 @@ _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a- _VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$") -def _is_https_url_with_host(url: str) -> bool: - try: - parts: Final = urlsplit(url) - except ValueError: - return False - return parts.scheme == "https" and bool(parts.hostname) - - def _validate_plugin_source(source: Mapping[str, str]) -> None: """Validate plugin source format, raising HTTPException on invalid input.""" source_type: Final = source.get("source") @@ -421,7 +413,7 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None: }, ) elif source_type == "archive": - if not _is_https_url_with_host(source.get("url", "")): + if not is_https_url_with_host(source.get("url", "")): raise HTTPException( status_code=400, detail={ diff --git a/litellm/proxy/enterprise_billing/billing_metrics.py b/litellm/proxy/enterprise_billing/billing_metrics.py index 166ece28b83..b6f69e8233c 100644 --- a/litellm/proxy/enterprise_billing/billing_metrics.py +++ b/litellm/proxy/enterprise_billing/billing_metrics.py @@ -2,29 +2,33 @@ Push-based OTLP metering for enterprise litellm deployments. Owns a dedicated OpenTelemetry meter provider and an OTLP/HTTP exporter -authenticated to our global collector with a TLS client certificate. The -collector front end terminates mutual TLS: the client certificate presented -here is validated against our CA at the edge, and the verified subject is -what identifies the deployment. It is intentionally isolated from the global -meter provider so the customer's own OTEL metrics are untouched and ours -never leak into their backend. +authenticated to our global collector with a Cognito bearer token. It is +intentionally isolated from the global meter provider so the customer's own +OTEL metrics are untouched and ours never leak into their backend. -The deployment's identity rides on the TLS client certificate, not on the -payload; the secret license key is never sent as an attribute or header. +The secret license key is never sent as an attribute or header. """ import os -import tempfile -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final, Optional +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypedDict +import requests from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.metrics import Counter from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader from opentelemetry.sdk.resources import Resource +from pydantic import BaseModel, ConfigDict +from requests.auth import AuthBase +from typing_extensions import ReadOnly, Unpack from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import is_https_url_with_host from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableCategory, ) @@ -33,23 +37,15 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData ENDPOINT_ENV: Final = "LITELLM_BILLING_METRICS_ENDPOINT" -CLIENT_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_CERT" -CLIENT_KEY_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_KEY" -CA_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CA_CERT" +CLIENT_ID_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_ID" +CLIENT_SECRET_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_SECRET" +TOKEN_URL_ENV: Final = "LITELLM_BILLING_METRICS_TOKEN_URL" EXPORT_INTERVAL_ENV: Final = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS" DEFAULT_EXPORT_INTERVAL_MS: Final = 60_000 SHUTDOWN_FLUSH_TIMEOUT_MS: Final = 5_000 +_TOKEN_REFRESH_SKEW_SECONDS: Final = 60 _METRICS_PATH: Final = "/v1/metrics" -# The cert env vars take a path or the PEM itself. Secret stores that inject -# values as env content cannot mount them as files, so inline PEM is written out. -_PEM_PREFIX: Final = "-----BEGIN" -_PEM_DIR_PREFIX: Final = "litellm-billing-mtls-" -_PEM_FILE_MODE: Final = 0o600 -_CLIENT_CERT_FILENAME: Final = "client.crt" -_CLIENT_KEY_FILENAME: Final = "client.key" -_CA_CERT_FILENAME: Final = "ca.crt" - METRIC_NAME: Final = "litellm.enterprise.billable_requests" METER_NAME: Final = "litellm.enterprise.billing" @@ -59,21 +55,28 @@ AttributeValue = str | int @dataclass(frozen=True, slots=True) class BillingMetricsConfig: endpoint: str - client_cert_path: str - client_key_path: str - ca_cert_path: str | None + token_url: str + client_id: str + client_secret: str = field(repr=False) export_interval_ms: int litellm_version: str license_id: str | None +class _BillingMetricsOAuthToken(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + + access_token: str + expires_in: int + + def _metrics_endpoint(endpoint: str) -> str: """The OTLP/HTTP metric exporter wants the full URL including the signal path.""" trimmed: Final = endpoint.rstrip("/") return trimmed if trimmed.endswith(_METRICS_PATH) else f"{trimmed}{_METRICS_PATH}" -def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]: +def resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]: base: Final[dict[str, AttributeValue]] = { "service.name": "litellm-proxy", "litellm.version": config.litellm_version, @@ -82,7 +85,7 @@ def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeVal return {**base, **license_attr} -def _billable_attributes( +def billable_attributes( category: BillableCategory, route: str, status_code: int, model_id: str | None ) -> dict[str, AttributeValue]: base: Final[dict[str, AttributeValue]] = { @@ -94,23 +97,95 @@ def _billable_attributes( return {**base, **model_attr} -def build_mtls_meter_provider(config: BillingMetricsConfig) -> MeterProvider: - """OTLP/HTTP exporter presenting a TLS client certificate. +class BillingMetricsBearerAuth(AuthBase): + def __init__( + self, + *, + client_id: str, + client_secret: str, + token_url: str, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._client_id: Final = client_id + self._client_secret: Final = client_secret + self._token_url: Final = token_url + self._clock: Final = clock + self._lock: Final = threading.Lock() + self._cached_token: str | None = None + self._refresh_at: float = 0.0 - The collector's load balancer terminates mutual TLS and validates the client - certificate against our CA. Server verification uses the system trust store - (the collector presents a public web-PKI certificate); ca_cert_path overrides - it only for private/test collectors. - """ - exporter: Final = OTLPMetricExporter( - endpoint=_metrics_endpoint(config.endpoint), - # None -> exporter falls back to the system trust store. - certificate_file=config.ca_cert_path, - client_certificate_file=config.client_cert_path, - client_key_file=config.client_key_path, + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: + token: Final = self._get_token() + authenticated_request: Final = request.copy() + authenticated_request.headers["Authorization"] = f"Bearer {token}" + return authenticated_request + + def invalidate(self) -> None: + with self._lock: + self._cached_token = None + self._refresh_at = 0.0 + + def _get_token(self) -> str: + with self._lock: + now: Final = self._clock() + if self._cached_token is not None and now < self._refresh_at: + return self._cached_token + token, cache_seconds = self._fetch_token() + self._cached_token = token + self._refresh_at = now + cache_seconds + return token + + def _fetch_token(self) -> tuple[str, float]: + try: + # The OTel periodic exporter's background thread makes blocking token I/O safe here. + response: Final = requests.post( + self._token_url, + auth=(self._client_id, self._client_secret), + data={"grant_type": "client_credentials", "scope": "telemetry/write"}, + timeout=10, + ) + response.raise_for_status() + token_response: Final = _BillingMetricsOAuthToken.model_validate(response.json()) + return token_response.access_token, max(0.0, float(token_response.expires_in) - _TOKEN_REFRESH_SKEW_SECONDS) + except (requests.RequestException, ValueError): + raise requests.RequestException("Billing metrics token request failed") from None + + +class _SessionHookKwargs(TypedDict, total=False): + stream: ReadOnly[object] + timeout: ReadOnly[object] + verify: ReadOnly[object] + cert: ReadOnly[object] + proxies: ReadOnly[object] + + +class _BillingMetricsBearerSession(requests.Session): + def __init__(self, auth: BillingMetricsBearerAuth) -> None: + super().__init__() + self._bearer_auth: Final = auth + self.auth = auth + self.hooks = {"response": [self._invalidate_token_on_unauthorized]} + + def _invalidate_token_on_unauthorized( + self, + response: requests.Response, + **kwargs: Unpack[_SessionHookKwargs], # kwargs-ok: requests hooks forward send options + ) -> requests.Response: + if response.status_code == 401: + self._bearer_auth.invalidate() + return response + + +def build_bearer_meter_provider(config: BillingMetricsConfig) -> MeterProvider: + auth: Final = BillingMetricsBearerAuth( + client_id=config.client_id, + client_secret=config.client_secret, + token_url=config.token_url, ) + session: Final = _BillingMetricsBearerSession(auth) + exporter: Final = OTLPMetricExporter(endpoint=_metrics_endpoint(config.endpoint), session=session) reader: Final = PeriodicExportingMetricReader(exporter, export_interval_millis=config.export_interval_ms) - return MeterProvider(metric_readers=[reader], resource=Resource.create(_resource_attributes(config))) + return MeterProvider(metric_readers=[reader], resource=Resource.create(resource_attributes(config))) class BillingMetricsRecorder: @@ -125,7 +200,7 @@ class BillingMetricsRecorder: ) def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: str | None) -> None: - self._counter.add(1, _billable_attributes(category, route, status_code, model_id)) + self._counter.add(1, billable_attributes(category, route, status_code, model_id)) def shutdown(self) -> None: """Final flush + exporter-thread stop. Without this, up to one export @@ -146,115 +221,34 @@ def _export_interval_ms() -> int: return DEFAULT_EXPORT_INTERVAL_MS -@dataclass(frozen=True, slots=True) -class _CredentialPaths: - client_cert_path: str - client_key_path: str - ca_cert_path: str | None - - -def _is_pem_content(value: str) -> bool: - return value.lstrip().startswith(_PEM_PREFIX) - - -def _write_pem(directory: str, filename: str, pem: str) -> str: - path: Final = os.path.join(directory, filename) - with open(path, "w", encoding="utf-8") as handle: - handle.write(pem if pem.endswith("\n") else f"{pem}\n") - os.chmod(path, _PEM_FILE_MODE) - return path - - -def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: str | None) -> _CredentialPaths: - """ - Accept either a filesystem path or inline PEM content for each credential. - - Secret stores that inject values as environment content rather than mounted - files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret - Manager) can only deliver the certificate as a string. The OTLP exporter - takes paths, so inline PEM is written to a private directory once, when the - recorder is built. Raises OSError if that write fails; the caller disables - metering rather than propagating. - """ - inline: Final = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value)) - if not inline: - return _CredentialPaths(client_cert, client_key, ca_cert) - - # mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users. - directory: Final = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX) - return _CredentialPaths( - client_cert_path=( - _write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert - ), - client_key_path=( - _write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key - ), - ca_cert_path=( - _write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert - ), - ) - - def load_billing_metrics_config( - *, license_data: Optional["EnterpriseLicenseData"], litellm_version: str + *, license_data: "EnterpriseLicenseData | None", litellm_version: str ) -> BillingMetricsConfig | None: - endpoint: Final = os.getenv(ENDPOINT_ENV) - client_cert: Final = os.getenv(CLIENT_CERT_ENV) - client_key: Final = os.getenv(CLIENT_KEY_ENV) - # Optional: only for private/test collectors whose server cert is not on the - # public web PKI. The production collector needs no CA override. - ca_cert: Final = os.getenv(CA_CERT_ENV) - - missing: Final = [ - name - for name, value in ( - (ENDPOINT_ENV, endpoint), - (CLIENT_CERT_ENV, client_cert), - (CLIENT_KEY_ENV, client_key), - ) - if not value - ] - if not endpoint or not client_cert or not client_key: + env: Final = MappingProxyType( + {name: os.getenv(name) or "" for name in (ENDPOINT_ENV, CLIENT_ID_ENV, CLIENT_SECRET_ENV, TOKEN_URL_ENV)} + ) + missing: Final = tuple(name for name, value in env.items() if not value) + if missing: verbose_proxy_logger.warning( "Enterprise billing metrics disabled: licensed deployment missing config (%s)", ", ".join(missing), ) return None - try: - paths: Final = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert) - except OSError as exc: - verbose_proxy_logger.warning( - "Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc - ) - return None - - # Report the variable names, never their values. A value that is neither a - # readable path nor recognizable PEM is still secret material, and this - # warning would otherwise copy a client key straight into the proxy logs. - unreadable: Final = [ + insecure_urls: Final = tuple( env_name - for env_name, path in ( - (CLIENT_CERT_ENV, paths.client_cert_path), - (CLIENT_KEY_ENV, paths.client_key_path), - (CA_CERT_ENV, paths.ca_cert_path), - ) - if path and not os.path.isfile(path) - ] - if unreadable: - verbose_proxy_logger.warning( - "Enterprise billing metrics disabled: %s did not resolve to a readable certificate file. " - "Set each to a file path, or to inline PEM content beginning with '%s'.", - ", ".join(unreadable), - _PEM_PREFIX, - ) + for env_name, url in ((ENDPOINT_ENV, env[ENDPOINT_ENV]), (TOKEN_URL_ENV, env[TOKEN_URL_ENV])) + if not is_https_url_with_host(url) + ) + if insecure_urls: + verbose_proxy_logger.warning("Enterprise billing metrics disabled: %s must use HTTPS", insecure_urls[0]) return None return BillingMetricsConfig( - endpoint=endpoint, - client_cert_path=paths.client_cert_path, - client_key_path=paths.client_key_path, - ca_cert_path=paths.ca_cert_path, + endpoint=env[ENDPOINT_ENV], + client_id=env[CLIENT_ID_ENV], + client_secret=env[CLIENT_SECRET_ENV], + token_url=env[TOKEN_URL_ENV], export_interval_ms=_export_interval_ms(), litellm_version=litellm_version, license_id=(license_data or {}).get("user_id"), @@ -282,7 +276,7 @@ _ACTIVE_RECORDER: Final = _ActiveRecorderRegistry() def build_billing_metrics_recorder( - *, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str + *, premium: bool, license_data: "EnterpriseLicenseData | None", litellm_version: str ) -> BillingMetricsRecorder | None: """Build the recorder, or None when the deployment is not licensed or metering is unconfigured.""" if not premium: @@ -296,14 +290,12 @@ def build_billing_metrics_recorder( return None try: - recorder: Final = BillingMetricsRecorder(build_mtls_meter_provider(config)) + provider: Final = build_bearer_meter_provider(config) + recorder: Final = BillingMetricsRecorder(provider) except Exception as exc: # noqa: BLE001 -- metering must never break proxy startup verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc) return None _ACTIVE_RECORDER.set(recorder) - # The only positive signal that this component meters. Without it, a silent - # return above is indistinguishable from a working exporter in the logs, and - # a component that carries the cert but no license would look healthy. verbose_proxy_logger.info( "Enterprise billing metrics enabled: exporting to %s every %d ms", config.endpoint, diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index bedb6f4e98b..6f0450c0091 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -212,33 +212,23 @@ orthogonal to the OTLP variables above; their credentials still go in ### Enterprise billing metrics -License-gated request metering is opt-in and gated entirely on -`billing_metrics_endpoint`. Empty (default) and no billing env is added to -the container, so existing deployments are unchanged. Set it and both -gateway and backend export billable-request counts over OTLP/HTTP, -authenticating to the collector with the mTLS client certificate issued for -your deployment. - -The proxy accepts the certificate, key, and CA bundle as either a file path -or literal PEM content. This stack takes the PEM, writes each one to its own -Secrets Manager entry, grants the task-execution role -`secretsmanager:GetSecretValue` on them, and injects them as -`LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and `_CA_CERT` when -set), so no volume mount is needed on Fargate. +License-gated request metering uses Cognito bearer authentication. Set the +endpoint, client ID, and token URL as plain values, and provide the client +secret through the sensitive Terraform variable. The client secret is stored +in Secrets Manager and injected into both gateway and backend tasks ```hcl billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +billing_metrics_client_id = "cognito-client-id" +billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" ``` ```bash -export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" -export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +export TF_VAR_billing_metrics_client_secret="cognito-client-secret" ``` -`billing_metrics_ca_cert_pem` is only for private or test collectors whose -CA is not in the system trust store; leave it empty against -`telemetry.litellm.ai`. Metering requires an enterprise license, so pair -this with `litellm_license`. To tune the export cadence, set +The client ID and client secret are required when the endpoint is set +Metering also requires an enterprise license. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 412e6c423f5..dec7f599ced 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -79,32 +79,19 @@ locals { { name = "OTEL_HEADERS", valueFrom = var.otel_headers_secret_arn }, ] : [] - # Enterprise request metering, gated on billing_metrics_endpoint. The - # endpoint rides in as a plain env var; the mTLS material is stored in - # Secrets Manager (secrets.tf) and injected as PEM-valued env vars, which - # the proxy accepts in place of file paths. Each PEM is wired only when the - # operator supplied it, so an empty ca_cert_pem falls back to the system - # trust store. - billing_metrics_enabled = var.billing_metrics_endpoint != "" - billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" - billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" - billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + # Enterprise request metering, gated on billing_metrics_endpoint. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != "" billing_metrics_env = local.billing_metrics_enabled ? [ { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id }, + { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url }, ] : [] - billing_metrics_secrets = concat( - local.billing_metrics_client_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_client_cert[0].arn }, - ] : [], - local.billing_metrics_client_key_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", valueFrom = aws_secretsmanager_secret.billing_metrics_client_key[0].arn }, - ] : [], - local.billing_metrics_ca_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CA_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_ca_cert[0].arn }, - ] : [], - ) + billing_metrics_secrets = local.billing_metrics_client_secret_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", valueFrom = aws_secretsmanager_secret.billing_metrics_client_secret[0].arn }, + ] : [] managed_db_env = var.create_database ? [ { name = "IAM_TOKEN_DB_AUTH", value = "true" }, @@ -338,27 +325,12 @@ locals { # ---------- Gateway ---------- resource "aws_ecs_task_definition" "gateway" { - # Metering needs a client certificate AND its key. Each secret is created only - # when its own PEM is supplied, so an endpoint set with a missing key would - # otherwise apply cleanly and leave the proxy logging "missing config" and - # never exporting. ca_cert_pem stays optional: empty means fall back to the - # system trust store. - # - # The guard lives here, on an unconditional resource, rather than on the cert - # secret: that secret is count-gated on the cert itself, so it has zero - # instances in exactly the case this must catch. Adding count or for_each to - # this resource would silently stop the guard from evaluating. - # - # endpoint cert key -> result - # "" any any -> metering off, no secrets created - # set set set -> metering on - # set any-missing -> plan fails here lifecycle { precondition { condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set." } precondition { @@ -465,9 +437,7 @@ resource "aws_ecs_service" "gateway" { aws_secretsmanager_secret_version.license, aws_secretsmanager_secret_version.database_url, aws_secretsmanager_secret_version.redis_url, - aws_secretsmanager_secret_version.billing_metrics_client_cert, - aws_secretsmanager_secret_version.billing_metrics_client_key, - aws_secretsmanager_secret_version.billing_metrics_ca_cert, + aws_secretsmanager_secret_version.billing_metrics_client_secret, ] tags = local.tags @@ -481,9 +451,9 @@ resource "aws_ecs_task_definition" "backend" { lifecycle { precondition { condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set." } } @@ -566,9 +536,7 @@ resource "aws_ecs_service" "backend" { aws_secretsmanager_secret_version.ui_password, aws_secretsmanager_secret_version.database_url, aws_secretsmanager_secret_version.redis_url, - aws_secretsmanager_secret_version.billing_metrics_client_cert, - aws_secretsmanager_secret_version.billing_metrics_client_key, - aws_secretsmanager_secret_version.billing_metrics_ca_cert, + aws_secretsmanager_secret_version.billing_metrics_client_secret, ] tags = local.tags diff --git a/terraform/litellm/aws/iam.tf b/terraform/litellm/aws/iam.tf index 3c55f07b02a..a653bd82d31 100644 --- a/terraform/litellm/aws/iam.tf +++ b/terraform/litellm/aws/iam.tf @@ -53,9 +53,7 @@ data "aws_iam_policy_document" "secrets_access" { [aws_secretsmanager_secret.master_key.arn], aws_secretsmanager_secret.license[*].arn, aws_secretsmanager_secret.ui_password[*].arn, - aws_secretsmanager_secret.billing_metrics_client_cert[*].arn, - aws_secretsmanager_secret.billing_metrics_client_key[*].arn, - aws_secretsmanager_secret.billing_metrics_ca_cert[*].arn, + aws_secretsmanager_secret.billing_metrics_client_secret[*].arn, aws_secretsmanager_secret.database_url[*].arn, aws_secretsmanager_secret.redis_url[*].arn, local.extra_secret_arns, diff --git a/terraform/litellm/aws/secrets.tf b/terraform/litellm/aws/secrets.tf index 921bae4d827..6821b8b43dc 100644 --- a/terraform/litellm/aws/secrets.tf +++ b/terraform/litellm/aws/secrets.tf @@ -75,59 +75,21 @@ resource "aws_secretsmanager_secret_version" "ui_password" { secret_string = var.ui_password } -# Billing-metrics mTLS material — only created when metering is enabled -# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. -# The task-execution role gets GetSecretValue via iam.tf, and gateway + -# backend pick the env vars up through shared_secrets in ecs.tf. -resource "aws_secretsmanager_secret" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "aws_secretsmanager_secret" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - name = "${local.name}-billing-metrics-client-cert" - description = "LITELLM_BILLING_METRICS_CLIENT_CERT for gateway + backend." + name = "${local.name}-billing-metrics-client-secret" + description = "LITELLM_BILLING_METRICS_CLIENT_SECRET for gateway + backend." recovery_window_in_days = 0 tags = local.tags } -resource "aws_secretsmanager_secret_version" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "aws_secretsmanager_secret_version" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = aws_secretsmanager_secret.billing_metrics_client_cert[0].id - secret_string = var.billing_metrics_client_cert_pem -} - -resource "aws_secretsmanager_secret" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - name = "${local.name}-billing-metrics-client-key" - description = "LITELLM_BILLING_METRICS_CLIENT_KEY for gateway + backend." - recovery_window_in_days = 0 - - tags = local.tags -} - -resource "aws_secretsmanager_secret_version" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = aws_secretsmanager_secret.billing_metrics_client_key[0].id - secret_string = var.billing_metrics_client_key_pem -} - -resource "aws_secretsmanager_secret" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - name = "${local.name}-billing-metrics-ca-cert" - description = "LITELLM_BILLING_METRICS_CA_CERT for gateway + backend." - recovery_window_in_days = 0 - - tags = local.tags -} - -resource "aws_secretsmanager_secret_version" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = aws_secretsmanager_secret.billing_metrics_ca_cert[0].id - secret_string = var.billing_metrics_ca_cert_pem + secret_id = aws_secretsmanager_secret.billing_metrics_client_secret[0].id + secret_string = var.billing_metrics_client_secret } resource "aws_secretsmanager_secret" "db_master_password" { diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 9b7a279a907..be7656c2d48 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -764,13 +764,7 @@ variable "otel_headers_secret_arn" { # ---------- Enterprise billing metrics ---------- # -# License-gated request metering. Opt-in and gated entirely on -# billing_metrics_endpoint: leave it empty (the default) and nothing -# metering-related lands in the container env. Set it and gateway + backend -# export billable-request counts over OTLP/HTTP, authenticating to the -# collector with an mTLS client cert. The proxy accepts the cert, key, and CA -# as either a file path or literal PEM content, so on Fargate they are -# injected straight from Secrets Manager as env vars and no volume is needed. +# License-gated request metering. The client secret is stored in Secrets Manager. variable "billing_metrics_endpoint" { description = <<-EOT @@ -784,44 +778,34 @@ variable "billing_metrics_endpoint" { default = "" } -variable "billing_metrics_client_cert_pem" { +variable "billing_metrics_client_id" { description = <<-EOT - PEM content of the mTLS client certificate issued for this deployment. - When billing_metrics_endpoint is set, the stack stores this in a - `-litellm--billing-metrics-client-cert` Secrets Manager - entry, grants the task-execution role GetSecretValue on it, and exposes - it to gateway + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required - whenever metering is enabled. + Cognito OAuth2 client ID exposed to the gateway and backend as + LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint + is set. + EOT + type = string + default = "" +} + +variable "billing_metrics_client_secret" { + description = <<-EOT + Cognito OAuth2 client secret stored in Secrets Manager and exposed to + the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET. + Required when billing_metrics_endpoint is set. EOT type = string default = "" sensitive = true } -variable "billing_metrics_client_key_pem" { +variable "billing_metrics_token_url" { description = <<-EOT - PEM content of the private key matching - billing_metrics_client_cert_pem. Stored in a - `-litellm--billing-metrics-client-key` Secrets Manager - entry and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required - whenever metering is enabled. + Cognito OAuth2 client-credentials token endpoint exposed to gateway and + backend as LITELLM_BILLING_METRICS_TOKEN_URL. EOT type = string - default = "" - sensitive = true -} - -variable "billing_metrics_ca_cert_pem" { - description = <<-EOT - PEM content of the CA bundle used to verify the metering collector. - Only needed for private or test collectors whose CA is not in the - system trust store; telemetry.litellm.ai is publicly trusted, so leave - this empty for production. When set, it is exposed as - LITELLM_BILLING_METRICS_CA_CERT. - EOT - type = string - default = "" - sensitive = true + default = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" } # ---------- Collector sidecar ---------- diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 3227dcac179..9e21ba14543 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -210,33 +210,23 @@ Behavior matches the AWS stack 1:1; the only naming differences are ### Enterprise billing metrics -License-gated request metering is opt-in and gated entirely on -`billing_metrics_endpoint`. Empty (default) and no billing env is added to -the container, so existing deployments are unchanged. Set it and both -gateway and backend export billable-request counts over OTLP/HTTP, -authenticating to the collector with the mTLS client certificate issued for -your deployment. - -The proxy accepts the certificate, key, and CA bundle as either a file path -or literal PEM content. This stack takes the PEM, writes each one to its own -Secret Manager entry, grants the runtime service account -`roles/secretmanager.secretAccessor` on them, and injects them as Cloud Run -secret env vars `LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and -`_CA_CERT` when set), so no volume mount is needed. +License-gated request metering uses Cognito bearer authentication. Set the +endpoint, client ID, and token URL as plain values, and provide the client +secret through the sensitive Terraform variable. The client secret is stored +in Secret Manager and injected into both gateway and backend services ```hcl billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +billing_metrics_client_id = "cognito-client-id" +billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" ``` ```bash -export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" -export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +export TF_VAR_billing_metrics_client_secret="cognito-client-secret" ``` -`billing_metrics_ca_cert_pem` is only for private or test collectors whose -CA is not in the system trust store; leave it empty against -`telemetry.litellm.ai`. Metering requires an enterprise license, so pair -this with `litellm_license`. To tune the export cadence, set +The client ID and client secret are required when the endpoint is set +Metering also requires an enterprise license. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index d0b32a367d6..07bf68b3a62 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -62,32 +62,19 @@ locals { { name = "OTEL_HEADERS", secret = var.otel_headers_secret, version = "latest" }, ] : [] - # Enterprise request metering, gated on billing_metrics_endpoint. The - # endpoint rides in as a plain env var; the mTLS material lives in Secret - # Manager (secrets.tf) and is injected as PEM-valued env vars, which the - # proxy accepts in place of file paths. Each PEM is wired only when the - # operator supplied it, so an empty ca_cert_pem falls back to the system - # trust store. - billing_metrics_enabled = var.billing_metrics_endpoint != "" - billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" - billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" - billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + # Enterprise request metering, gated on billing_metrics_endpoint. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != "" billing_metrics_env_kv = local.billing_metrics_enabled ? [ { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id }, + { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url }, ] : [] - billing_metrics_env_secrets = concat( - local.billing_metrics_client_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", secret = google_secret_manager_secret.billing_metrics_client_cert[0].id, version = "latest" }, - ] : [], - local.billing_metrics_client_key_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", secret = google_secret_manager_secret.billing_metrics_client_key[0].id, version = "latest" }, - ] : [], - local.billing_metrics_ca_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CA_CERT", secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id, version = "latest" }, - ] : [], - ) + billing_metrics_env_secrets = local.billing_metrics_client_secret_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", secret = google_secret_manager_secret.billing_metrics_client_secret[0].id, version = "latest" }, + ] : [] # Cloud Run v2 secret env vars use value_source.secret_key_ref pointing at a # secret resource ID. Shared between gateway and backend (the migrations @@ -481,9 +468,7 @@ resource "google_cloud_run_v2_service" "gateway" { google_secret_manager_secret_iam_member.license, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, - google_secret_manager_secret_iam_member.billing_metrics_client_cert, - google_secret_manager_secret_iam_member.billing_metrics_client_key, - google_secret_manager_secret_iam_member.billing_metrics_ca_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_secret, google_secret_manager_secret_iam_member.metrics_run_monitoring, google_project_iam_member.runtime_metric_writer, google_storage_bucket_iam_member.proxy_config_runtime, @@ -603,9 +588,7 @@ resource "google_cloud_run_v2_service" "backend" { google_secret_manager_secret_iam_member.ui_password, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, - google_secret_manager_secret_iam_member.billing_metrics_client_cert, - google_secret_manager_secret_iam_member.billing_metrics_client_key, - google_secret_manager_secret_iam_member.billing_metrics_ca_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_secret, google_storage_bucket_iam_member.proxy_config_runtime, google_sql_user.app, terraform_data.migration, diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index 5c2bc38bc12..1e0a5fe9626 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -10,9 +10,9 @@ resource "google_service_account" "runtime" { lifecycle { precondition { condition = !var.create_runtime || var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set and create_runtime is true." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set and create_runtime is true." } } } @@ -91,28 +91,10 @@ resource "google_secret_manager_secret_iam_member" "otel_headers" { member = "serviceAccount:${google_service_account.runtime.email}" } -# Billing-metrics mTLS accessors — only created when request metering is -# enabled and the matching PEM was supplied. -resource "google_secret_manager_secret_iam_member" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret_iam_member" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = google_secret_manager_secret.billing_metrics_client_cert[0].id - role = "roles/secretmanager.secretAccessor" - member = "serviceAccount:${google_service_account.runtime.email}" -} - -resource "google_secret_manager_secret_iam_member" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = google_secret_manager_secret.billing_metrics_client_key[0].id - role = "roles/secretmanager.secretAccessor" - member = "serviceAccount:${google_service_account.runtime.email}" -} - -resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = google_secret_manager_secret.billing_metrics_ca_cert[0].id + secret_id = google_secret_manager_secret.billing_metrics_client_secret[0].id role = "roles/secretmanager.secretAccessor" member = "serviceAccount:${google_service_account.runtime.email}" } diff --git a/terraform/litellm/gcp/secrets.tf b/terraform/litellm/gcp/secrets.tf index b3b656721ab..f3636b6e8e8 100644 --- a/terraform/litellm/gcp/secrets.tf +++ b/terraform/litellm/gcp/secrets.tf @@ -64,59 +64,21 @@ resource "google_secret_manager_secret_version" "ui_password" { secret_data = var.ui_password } -# Billing-metrics mTLS material — only created when metering is enabled -# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. -# The runtime SA gets accessor permission via iam.tf, and gateway + backend -# pick the env vars up through billing_metrics_env_secrets in cloudrun.tf. -resource "google_secret_manager_secret" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = "${local.name}-billing-metrics-client-cert" + secret_id = "${local.name}-billing-metrics-client-secret" labels = local.labels replication { auto {} } } -resource "google_secret_manager_secret_version" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret_version" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret = google_secret_manager_secret.billing_metrics_client_cert[0].id - secret_data = var.billing_metrics_client_cert_pem -} - -resource "google_secret_manager_secret" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = "${local.name}-billing-metrics-client-key" - labels = local.labels - replication { - auto {} - } -} - -resource "google_secret_manager_secret_version" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret = google_secret_manager_secret.billing_metrics_client_key[0].id - secret_data = var.billing_metrics_client_key_pem -} - -resource "google_secret_manager_secret" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = "${local.name}-billing-metrics-ca-cert" - labels = local.labels - replication { - auto {} - } -} - -resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id - secret_data = var.billing_metrics_ca_cert_pem + secret = google_secret_manager_secret.billing_metrics_client_secret[0].id + secret_data = var.billing_metrics_client_secret } resource "google_secret_manager_secret" "metrics_run_monitoring" { diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 6161e893e02..e40ee36fa18 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -602,14 +602,7 @@ variable "gateway_metrics_collector_image" { # ---------- Enterprise billing metrics ---------- # -# License-gated request metering. Opt-in and gated entirely on -# billing_metrics_endpoint: leave it empty (the default) and nothing -# metering-related is added to the container env. Set it and gateway + -# backend export billable-request counts over OTLP/HTTP, authenticating to -# the collector with an mTLS client cert. The proxy accepts the cert, key, -# and CA as either a file path or literal PEM content, so on Cloud Run they -# are injected straight from Secret Manager as env vars and no volume is -# needed. +# License-gated request metering. The client secret is stored in Secret Manager. variable "billing_metrics_endpoint" { description = <<-EOT @@ -623,44 +616,34 @@ variable "billing_metrics_endpoint" { default = "" } -variable "billing_metrics_client_cert_pem" { +variable "billing_metrics_client_id" { description = <<-EOT - PEM content of the mTLS client certificate issued for this deployment. - When billing_metrics_endpoint is set, the stack stores this in a - `-litellm--billing-metrics-client-cert` Secret Manager - entry, grants the runtime SA accessor on it, and exposes it to gateway + - backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required whenever - metering is enabled. + Cognito OAuth2 client ID exposed to the gateway and backend as + LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint + is set. + EOT + type = string + default = "" +} + +variable "billing_metrics_client_secret" { + description = <<-EOT + Cognito OAuth2 client secret stored in Secret Manager and exposed to + the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET. + Required when billing_metrics_endpoint is set. EOT type = string default = "" sensitive = true } -variable "billing_metrics_client_key_pem" { +variable "billing_metrics_token_url" { description = <<-EOT - PEM content of the private key matching - billing_metrics_client_cert_pem. Stored in a - `-litellm--billing-metrics-client-key` Secret Manager entry - and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required whenever - metering is enabled. + Cognito OAuth2 client-credentials token endpoint exposed to gateway and + backend as LITELLM_BILLING_METRICS_TOKEN_URL. EOT type = string - default = "" - sensitive = true -} - -variable "billing_metrics_ca_cert_pem" { - description = <<-EOT - PEM content of the CA bundle used to verify the metering collector. - Only needed for private or test collectors whose CA is not in the - system trust store; telemetry.litellm.ai is publicly trusted, so leave - this empty for production. When set, it is exposed as - LITELLM_BILLING_METRICS_CA_CERT. - EOT - type = string - default = "" - sensitive = true + default = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" } # ---------- Collector sidecar ---------- diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index 9b29eac208e..958960af27a 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/tests/unit/litellm_core_utils/test_url_utils.py @@ -16,11 +16,25 @@ from litellm.litellm_core_utils.url_utils import ( encode_url_path_segment, encode_url_path_segments, ensure_api_route, + is_https_url_with_host, join_api_base, validate_url, ) +@pytest.mark.parametrize( + ("url", "expected"), + ( + ("https://example.com/path", True), + ("http://example.com/path", False), + ("https:///path", False), + ("https://[", False), + ), +) +def test_is_https_url_with_host(url: str, expected: bool) -> None: + assert is_https_url_with_host(url) is expected + + @pytest.fixture def mock_dns_public(monkeypatch): """Resolve any hostname to 93.184.216.34 (public).""" diff --git a/tests/unit/proxy/enterprise_billing/test_billing_metrics.py b/tests/unit/proxy/enterprise_billing/test_billing_metrics.py index 0446cfeeab0..14a8fea04f7 100644 --- a/tests/unit/proxy/enterprise_billing/test_billing_metrics.py +++ b/tests/unit/proxy/enterprise_billing/test_billing_metrics.py @@ -1,398 +1,379 @@ -""" -Tests for the enterprise billing-metrics recorder and its factory. - -These verify the license gate, the missing-config and missing-cert disable -paths, the OTLP/HTTP exporter wiring (client cert+key authenticate us to the -collector's mTLS-terminating front end; CA override optional for private -collectors), the metric attribute mapping, and that recording produces the -expected OTLP counter via an in-memory reader. -""" - -import os -import socket -import stat -from pathlib import Path -from typing import Dict, List, Optional +from collections.abc import Generator, Mapping +from typing import Final import pytest +import requests +import requests_mock from opentelemetry.sdk.metrics import MeterProvider -from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Sum +from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.proxy._types import EnterpriseLicenseData from litellm.proxy.enterprise_billing import billing_metrics as bm from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory -_ENV_VARS = ( + +class _RecordedRequest(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + url: str | None + body: str | bytes | None + headers: dict[str, str] + + +def _request_history(http_mock: requests_mock.Mocker) -> tuple[_RecordedRequest, ...]: + return TypeAdapter(tuple[_RecordedRequest, ...]).validate_python(http_mock.request_history) + + +def _metric_category(attributes: Mapping[str, object] | None) -> str: + if attributes is None: + return "" + return str(attributes.get("litellm.endpoint.category")) + + +_ENV_VARS: Final = ( bm.ENDPOINT_ENV, - bm.CLIENT_CERT_ENV, - bm.CLIENT_KEY_ENV, - bm.CA_CERT_ENV, + bm.CLIENT_ID_ENV, + bm.CLIENT_SECRET_ENV, + bm.TOKEN_URL_ENV, bm.EXPORT_INTERVAL_ENV, ) @pytest.fixture(autouse=True) -def clear_env(monkeypatch): +def clear_env(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: for name in _ENV_VARS: monkeypatch.delenv(name, raising=False) yield bm.shutdown_billing_metrics_recorder() -def _write_certs(tmp_path: Path) -> Dict[str, str]: - files = { - bm.CA_CERT_ENV: ("ca.pem", b"ca-bytes"), - bm.CLIENT_CERT_ENV: ("client.pem", b"client-cert-bytes"), - bm.CLIENT_KEY_ENV: ("client.key", b"client-key-bytes"), - } - paths = {} - for env_name, (filename, content) in files.items(): - path = tmp_path / filename - path.write_bytes(content) - paths[env_name] = str(path) - return paths +def _set_full_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4318") + monkeypatch.setenv(bm.CLIENT_ID_ENV, "client-id") + monkeypatch.setenv(bm.CLIENT_SECRET_ENV, "client-secret") + monkeypatch.setenv(bm.TOKEN_URL_ENV, "https://auth.example/oauth2/token") -def _set_full_env(monkeypatch, tmp_path: Path) -> Dict[str, str]: - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - return paths - - -def _config(tmp_path: Path, license_id: Optional[str] = "org-1") -> bm.BillingMetricsConfig: - paths = _write_certs(tmp_path) +def _config(license_id: str | None = "org-1") -> bm.BillingMetricsConfig: return bm.BillingMetricsConfig( - endpoint="https://collector.example:4317", - client_cert_path=paths[bm.CLIENT_CERT_ENV], - client_key_path=paths[bm.CLIENT_KEY_ENV], - ca_cert_path=paths[bm.CA_CERT_ENV], + endpoint="https://collector.example:4318", + token_url="https://auth.example/oauth2/token", + client_id="client-id", + client_secret="client-secret", export_interval_ms=60_000, litellm_version="1.2.3", license_id=license_id, ) -# ── Factory gating ──────────────────────────────────────────────────────────── +@pytest.mark.parametrize("missing_env", (bm.ENDPOINT_ENV, bm.CLIENT_ID_ENV, bm.CLIENT_SECRET_ENV, bm.TOKEN_URL_ENV)) +def test_missing_required_env_disables_metrics_without_logging_values( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + missing_env: str, +) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(bm.CLIENT_SECRET_ENV, "secret-for-test") + monkeypatch.delenv(missing_env) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert f"licensed deployment missing config ({missing_env})" in caplog.text + assert "secret-for-test" not in caplog.text -def test_not_premium_returns_none(tmp_path, monkeypatch): - _set_full_env(monkeypatch, tmp_path) +def test_empty_required_env_disables_metrics(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(bm.ENDPOINT_ENV, "") + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert bm.ENDPOINT_ENV in caplog.text + + +@pytest.mark.parametrize( + ("env_name", "url"), + ( + (bm.ENDPOINT_ENV, "http://collector.example:4318"), + (bm.TOKEN_URL_ENV, "http://auth.example/oauth2/token"), + ), +) +def test_http_urls_disable_metrics( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + env_name: str, + url: str, +) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(env_name, url) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert f"{env_name} must use HTTPS" in caplog.text + + +def test_valid_config_contains_bearer_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + _set_full_env(monkeypatch) + + license_data: Final[EnterpriseLicenseData] = {"user_id": "org-1"} + config: Final = bm.load_billing_metrics_config(license_data=license_data, litellm_version="1.0") + assert config is not None + + assert config == bm.BillingMetricsConfig( + endpoint="https://collector.example:4318", + token_url="https://auth.example/oauth2/token", + client_id="client-id", + client_secret="client-secret", + export_interval_ms=bm.DEFAULT_EXPORT_INTERVAL_MS, + litellm_version="1.0", + license_id="org-1", + ) + assert "client-secret" not in repr(config) + assert "license-secret" not in repr(config) + + +def test_not_premium_returns_none_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + _set_full_env(monkeypatch) + assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None + assert "Enterprise billing metrics disabled" not in caplog.text -def test_premium_without_config_returns_none(monkeypatch): +def test_premium_without_config_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: assert bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "x"}, litellm_version="1.0") is None -def test_premium_with_missing_cert_files_returns_none(monkeypatch, tmp_path): - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CA_CERT_ENV, str(tmp_path / "missing-ca.pem")) - monkeypatch.setenv(bm.CLIENT_CERT_ENV, str(tmp_path / "missing-cert.pem")) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, str(tmp_path / "missing-key.pem")) - assert bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") is None +def test_premium_builds_recorder_and_logs_enabled( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + _set_full_env(monkeypatch) + caplog.set_level("INFO", logger=bm.verbose_proxy_logger.name) + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example:4318/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "access-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + recorder: Final = bm.build_billing_metrics_recorder( + premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" + ) + assert recorder is not None + assert "Enterprise billing metrics enabled: exporting to https://collector.example:4318 every 60000 ms" in caplog.text + bm.shutdown_billing_metrics_recorder() -def test_premium_with_full_config_builds_recorder(monkeypatch, tmp_path): - """Builds a real MeterProvider, so the exporter is stubbed: the live one - resolves the collector and opens a TLS connection during the shutdown flush. - The getaddrinfo spy keeps that stub from being quietly dropped later.""" - _set_full_env(monkeypatch, tmp_path) - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) +class _ShutdownCountingMeterProvider(MeterProvider): + def __init__(self) -> None: + super().__init__() + self.shutdown_calls = 0 - resolved: List[str] = [] - real_getaddrinfo = socket.getaddrinfo + def shutdown(self, timeout_millis: float = 30_000) -> None: + self.shutdown_calls += 1 + super().shutdown(timeout_millis=timeout_millis) - def _spy_getaddrinfo(host, port, *args, **kwargs): - resolved.append(str(host)) - return real_getaddrinfo(host, port, *args, **kwargs) - monkeypatch.setattr(socket, "getaddrinfo", _spy_getaddrinfo) +def test_shutdown_flushes_active_recorder_once(monkeypatch: pytest.MonkeyPatch) -> None: + provider: Final = _ShutdownCountingMeterProvider() + _set_full_env(monkeypatch) - recorder = bm.build_billing_metrics_recorder( - premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" + def create_provider(config: bm.BillingMetricsConfig) -> MeterProvider: + return provider + + monkeypatch.setattr(bm, "build_bearer_meter_provider", create_provider) + recorder: Final = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") + assert recorder is not None + + bm.shutdown_billing_metrics_recorder() + bm.shutdown_billing_metrics_recorder() + + assert provider.shutdown_calls == 1 + + +class _FakeClock: + def __init__(self, now: float) -> None: + self._now = now + + def __call__(self) -> float: + return self._now + + def set(self, now: float) -> None: + self._now = now + + +def test_bearer_auth_fetches_and_caches_client_credentials_token() -> None: + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example/v1/metrics" + auth: Final = bm.BillingMetricsBearerAuth( + client_id="client-id", + client_secret="client-secret", + token_url=token_url, + clock=lambda: 100.0, ) - assert isinstance(recorder, bm.BillingMetricsRecorder) - recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id=None) - bm.shutdown_billing_metrics_recorder() + session: Final = requests.Session() + session.auth = auth + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "access-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + session.post(export_url, data=b"metrics") + session.post(export_url, data=b"metrics") - assert [host for host in resolved if "collector.example" in host] == [] + request_records: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in request_records if request.url == token_url) + export_requests: Final = tuple(request for request in request_records if request.url == export_url) + assert len(token_requests) == 1 + assert token_requests[0].headers["Authorization"].startswith("Basic ") + assert token_requests[0].body == "grant_type=client_credentials&scope=telemetry%2Fwrite" + assert tuple(request.headers["Authorization"] for request in export_requests) == ( + "Bearer access-token", + "Bearer access-token", + ) -def test_building_the_recorder_logs_an_affirmative_line(monkeypatch, tmp_path): - """ - Every disable path logs; a successful build must log too. Otherwise an - operator cannot tell a metering component from one that silently returned - None, which is how an unlicensed component looks healthy while exporting - nothing. - """ - _set_full_env(monkeypatch, tmp_path) +def test_bearer_auth_refreshes_before_token_expiry() -> None: + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example/v1/metrics" + clock: Final = _FakeClock(1000.0) + auth: Final = bm.BillingMetricsBearerAuth( + client_id="client-id", + client_secret="client-secret", + token_url=token_url, + clock=clock, + ) + session: Final = requests.Session() + session.auth = auth + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post( + token_url, + [ + {"json": {"access_token": "first-token", "expires_in": 120}}, + {"json": {"access_token": "second-token", "expires_in": 120}}, + ], + ) + http_mock.post(export_url, status_code=200) + session.post(export_url, data=b"metrics") + clock.set(1059.0) + session.post(export_url, data=b"metrics") + clock.set(1060.0) + session.post(export_url, data=b"metrics") + + request_records: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in request_records if request.url == token_url) + export_requests: Final = tuple(request for request in request_records if request.url == export_url) + assert len(token_requests) == 2 + assert tuple(request.headers["Authorization"] for request in export_requests) == ( + "Bearer first-token", + "Bearer first-token", + "Bearer second-token", + ) + + +def test_bearer_export_sets_authorization_header() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "export-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + provider: Final = bm.build_bearer_meter_provider(config) + try: + provider.get_meter("test").create_counter("requests").add(1) + assert provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + export_requests: Final = tuple(request for request in _request_history(http_mock) if request.url == export_url) + assert export_requests + assert all(request.headers["Authorization"] == "Bearer export-token" for request in export_requests) + + +def test_unauthorized_export_refetches_token_before_next_export() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post( + token_url, + [ + {"json": {"access_token": "first-token", "expires_in": 3600}}, + {"json": {"access_token": "second-token", "expires_in": 3600}}, + ], + ) + http_mock.post(export_url, [{"status_code": 401}, {"status_code": 200}]) + provider: Final = bm.build_bearer_meter_provider(config) + try: + counter: Final = provider.get_meter("test").create_counter("requests") + counter.add(1) + provider.force_flush(timeout_millis=1000) + counter.add(1) + provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + requests: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in requests if request.url == token_url) + export_requests: Final = tuple(request for request in requests if request.url == export_url) + assert len(token_requests) == 2 + assert tuple(request.headers["Authorization"] for request in export_requests[:2]) == ( + "Bearer first-token", + "Bearer second-token", + ) + + +def test_failed_bearer_token_fetch_does_not_raise_from_export() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, status_code=400) + provider: Final = bm.build_bearer_meter_provider(config) + try: + provider.get_meter("test").create_counter("requests").add(1) + assert provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + request_urls: Final = tuple(request.url for request in _request_history(http_mock)) + assert token_url in request_urls + assert export_url not in request_urls + + +def test_export_interval_default_and_override(monkeypatch: pytest.MonkeyPatch) -> None: + _set_full_env(monkeypatch) + default_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert default_config is not None + assert default_config.export_interval_ms == bm.DEFAULT_EXPORT_INTERVAL_MS monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) - - infos: List[str] = [] - monkeypatch.setattr(bm.verbose_proxy_logger, "info", lambda msg, *args: infos.append(msg % args if args else msg)) - - recorder = bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0") - - assert recorder is not None - joined = "\n".join(infos) - assert "https://collector.example:4317" in joined - assert "5000" in joined - - -def test_unlicensed_build_does_not_warn(monkeypatch, tmp_path): - """Unlicensed is the common OSS case; warning there would be pure noise.""" - _set_full_env(monkeypatch, tmp_path) - - warnings: List[str] = [] - monkeypatch.setattr(bm.verbose_proxy_logger, "warning", lambda msg, *args: warnings.append(str(msg))) - - assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None - assert warnings == [] - - -def test_shutdown_flushes_active_recorder_once(monkeypatch, tmp_path): - """The shutdown hook must flush the recorder the factory built (buffered - counts are lost on restart otherwise) and be idempotent for repeat calls.""" - _set_full_env(monkeypatch, tmp_path) - shutdowns = [] - - class _SpyProvider: - def get_meter(self, name): - return MeterProvider().get_meter(name) - - def shutdown(self, timeout_millis=None): - shutdowns.append(timeout_millis) - - monkeypatch.setattr(bm, "build_mtls_meter_provider", lambda config: _SpyProvider()) - recorder = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") - assert recorder is not None - - bm.shutdown_billing_metrics_recorder() - bm.shutdown_billing_metrics_recorder() - assert shutdowns == [bm.SHUTDOWN_FLUSH_TIMEOUT_MS] - - -def test_shutdown_without_active_recorder_is_noop(): - bm.shutdown_billing_metrics_recorder() - - -# ── Config loading ──────────────────────────────────────────────────────────── - - -def test_load_config_carries_license_id(monkeypatch, tmp_path): - _set_full_env(monkeypatch, tmp_path) - config = bm.load_billing_metrics_config(license_data={"user_id": "org-42"}, litellm_version="9.9") - assert config is not None and config.license_id == "org-42" and config.litellm_version == "9.9" - - -def test_load_config_with_empty_string_env_is_disabled(monkeypatch, tmp_path): - """An env var set to the empty string is as unusable as an unset one and - must disable metering rather than produce a config with a blank endpoint.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -_CLIENT_CERT_PEM = "-----BEGIN CERTIFICATE-----\nclient-cert-body\n-----END CERTIFICATE-----" -_CLIENT_KEY_PEM = "-----BEGIN PRIVATE KEY-----\nclient-key-body\n-----END PRIVATE KEY-----" -_CA_CERT_PEM = "-----BEGIN CERTIFICATE-----\nca-body\n-----END CERTIFICATE-----" - - -def test_load_config_materializes_inline_pem_content(monkeypatch): - """ - ECS and Cloud Run inject secrets as env content, not as mounted files, so the - cert env vars must accept PEM directly. The exporter takes paths, so the PEM - is written to disk and the config points at those files. - """ - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - monkeypatch.setenv(bm.CA_CERT_ENV, _CA_CERT_PEM) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.ca_cert_path is not None - written = { - config.client_cert_path: _CLIENT_CERT_PEM, - config.client_key_path: _CLIENT_KEY_PEM, - config.ca_cert_path: _CA_CERT_PEM, - } - for path, pem in written.items(): - assert path != pem, "config must carry a file path, not the PEM itself" - assert os.path.isfile(path) - assert Path(path).read_text(encoding="utf-8") == f"{pem}\n" - - # The private key must not be world- or group-readable. - assert stat.S_IMODE(os.stat(config.client_key_path).st_mode) == 0o600 - - -def test_load_config_accepts_a_mix_of_pem_content_and_file_paths(monkeypatch, tmp_path): - """A deployment may mount the CA but inject the client credentials inline.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.ca_cert_path == paths[bm.CA_CERT_ENV] - assert Path(config.client_cert_path).read_text(encoding="utf-8") == f"{_CLIENT_CERT_PEM}\n" - - -def test_load_config_leaves_file_paths_untouched(monkeypatch, tmp_path): - """Path-valued env vars keep working; nothing is copied or rewritten.""" - paths = _set_full_env(monkeypatch, tmp_path) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.client_cert_path == paths[bm.CLIENT_CERT_ENV] - assert config.client_key_path == paths[bm.CLIENT_KEY_ENV] - assert config.ca_cert_path == paths[bm.CA_CERT_ENV] - - -def test_load_config_with_inline_pem_disabled_when_unwritable(monkeypatch): - """A failure to materialize the PEM disables metering instead of raising.""" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - - def _explode(prefix=None): - raise OSError("read-only filesystem") - - monkeypatch.setattr(bm.tempfile, "mkdtemp", _explode) - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -def test_load_config_never_logs_credential_values(monkeypatch): - """ - A value that is neither a readable path nor `-----BEGIN`-prefixed PEM is - still secret material. The disable warning must name the env vars, never - echo their contents, or a malformed key lands in the proxy logs. - """ - secret_material = "MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQ-not-pem-prefixed" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, secret_material) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, secret_material) - - logged: List[str] = [] - - def _capture(msg, *args): - logged.append(msg % args if args else msg) - - monkeypatch.setattr(bm.verbose_proxy_logger, "warning", _capture) - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - joined = "\n".join(logged) - assert secret_material not in joined - assert bm.CLIENT_CERT_ENV in joined and bm.CLIENT_KEY_ENV in joined - - -def test_load_config_with_empty_pem_env_is_disabled(monkeypatch): - """Empty stays empty: an unset secret must not be mistaken for inline PEM.""" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, "") - monkeypatch.setenv(bm.CLIENT_KEY_ENV, "") - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -def test_export_interval_default_and_override(monkeypatch): - assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS - monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") - assert bm._export_interval_ms() == 5000 + override_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert override_config is not None + assert override_config.export_interval_ms == 5000 monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "not-a-number") - assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS + invalid_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert invalid_config is not None + assert invalid_config.export_interval_ms == bm.DEFAULT_EXPORT_INTERVAL_MS -# ── OTLP/HTTP exporter wiring ───────────────────────────────────────────────── +def test_resource_attributes_include_license_id() -> None: + attributes: Final = bm.resource_attributes(_config(license_id="org-7")) + assert attributes == { + "service.name": "litellm-proxy", + "litellm.version": "1.2.3", + "litellm.license.id": "org-7", + } -def test_metrics_endpoint_appends_signal_path(): - assert bm._metrics_endpoint("https://telemetry.example.com") == "https://telemetry.example.com/v1/metrics" - assert bm._metrics_endpoint("https://telemetry.example.com/") == "https://telemetry.example.com/v1/metrics" - assert bm._metrics_endpoint("https://telemetry.example.com/v1/metrics") == "https://telemetry.example.com/v1/metrics" +def test_resource_attributes_omit_license_id_when_absent() -> None: + assert "litellm.license.id" not in bm.resource_attributes(_config(license_id=None)) -def _fake_exporter_class(captured: Dict[str, object]) -> type: - """A no-network stand-in for OTLPMetricExporter. Tests that build a real - MeterProvider must install this: the real exporter resolves the collector - host and opens a TLS connection on the reader's first export and on the - shutdown flush.""" - - class _FakeExporter: - # PeriodicExportingMetricReader probes these on the exporter it wraps. - _preferred_temporality: dict = {} - _preferred_aggregation: dict = {} - - def __init__(self, **kwargs): - captured.update(kwargs) - - def export(self, *args, **kwargs): - return None - - def shutdown(self, *args, **kwargs): - return None - - def force_flush(self, *args, **kwargs): - return True - - return _FakeExporter - - -def test_meter_provider_wires_client_cert_into_http_exporter(tmp_path, monkeypatch): - """Client cert+key authenticate us at the collector's mTLS front end; CA override rides certificate_file.""" - captured: Dict[str, object] = {} - - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class(captured)) - config = _config(tmp_path) - provider = bm.build_mtls_meter_provider(config) - provider.shutdown() - - assert captured["endpoint"] == "https://collector.example:4317/v1/metrics" - assert captured["client_certificate_file"] == config.client_cert_path - assert captured["client_key_file"] == config.client_key_path - assert captured["certificate_file"] == config.ca_cert_path - - -def test_load_config_without_ca_is_valid(monkeypatch, tmp_path): - """The production collector presents a public web-PKI cert: no CA override required.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://telemetry.example.com") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - assert config is not None and config.ca_cert_path is None - - -# ── Resource and metric attributes ──────────────────────────────────────────── - - -def test_resource_attributes_include_license_id(tmp_path): - attrs = bm._resource_attributes(_config(tmp_path, license_id="org-7")) - assert attrs["service.name"] == "litellm-proxy" - assert attrs["litellm.version"] == "1.2.3" - assert attrs["litellm.license.id"] == "org-7" - - -def test_resource_attributes_omit_license_id_when_absent(tmp_path): - attrs = bm._resource_attributes(_config(tmp_path, license_id=None)) - assert "litellm.license.id" not in attrs - - -def test_billable_attributes_with_model_id(): - attrs = bm._billable_attributes(BillableCategory.LLM, "/chat/completions", 200, "deploy-3") - assert attrs == { +def test_billable_attributes_include_model_id() -> None: + assert bm.billable_attributes( + BillableCategory.LLM, "/chat/completions", 200, "deploy-3" + ) == { "litellm.endpoint.category": "llm", "http.route": "/chat/completions", "http.response.status_code": 200, @@ -400,33 +381,35 @@ def test_billable_attributes_with_model_id(): } -def test_billable_attributes_omit_model_id_when_none(): - attrs = bm._billable_attributes(BillableCategory.MCP, "/mcp", 200, None) - assert "litellm.model_id" not in attrs +def test_billable_attributes_omit_model_id_when_none() -> None: + attributes: Final = bm.billable_attributes( + BillableCategory.MCP, "/mcp", 200, None + ) + assert "litellm.model_id" not in attributes -# ── End-to-end recording via in-memory reader ───────────────────────────────── +def _counter_points(reader: InMemoryMetricReader) -> tuple[tuple[str, int], ...]: + data: Final = reader.get_metrics_data() + if data is None: + return () + return tuple( + (_metric_category(point.attributes), int(point.value)) + for resource_metric in data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + if metric.name == bm.METRIC_NAME and isinstance(metric.data, Sum) + for point in metric.data.data_points + ) -def _counter_points(reader: InMemoryMetricReader): - data = reader.get_metrics_data() - for resource_metric in data.resource_metrics: - for scope_metric in resource_metric.scope_metrics: - for metric in scope_metric.metrics: - if metric.name == bm.METRIC_NAME: - return list(metric.data.data_points) - return [] - - -def test_record_increments_counter_with_attributes(): - reader = InMemoryMetricReader() - recorder = bm.BillingMetricsRecorder(MeterProvider(metric_readers=[reader])) +def test_record_increments_counter_with_attributes() -> None: + reader: Final = InMemoryMetricReader() + recorder: Final = bm.BillingMetricsRecorder(MeterProvider(metric_readers=[reader])) recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") recorder.record(category=BillableCategory.MCP, route="/mcp", status_code=200, model_id=None) - points = _counter_points(reader) - by_category = {point.attributes["litellm.endpoint.category"]: point.value for point in points} - assert by_category["llm"] == 2 - assert by_category["mcp"] == 1 + points: Final = _counter_points(reader) + by_category: Final = dict(points) + assert by_category == {"llm": 2, "mcp": 1} From d38c5143ae0973acba796b434d02df6626d88b41 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 14:12:55 -0800 Subject: [PATCH 08/19] test(skills): drop the archive rejection test whose zip-byte ids differ across xdist workers (#50118) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_skill_upload_archive.py | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py b/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py index 531d97efe64..04b81605e29 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py +++ b/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py @@ -8,10 +8,8 @@ import pytest from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import ( MAX_SKILL_ARCHIVE_COMPRESSION_RATIO, - MAX_SKILL_UPLOAD_BYTES, build_uploaded_skill_archive, ) -from litellm.proxy.discovery_endpoints.agent_skills_archive import MAX_ARCHIVE_ENTRIES, MAX_ARCHIVE_UNPACKED_BYTES MANIFEST: Final = b"---\nname: skill\ndescription: Notes\n---\nRead guide.md" @@ -64,29 +62,3 @@ def test_installable_archive_preserves_files_and_safe_modes( assert {entry.compress_type for entry in archive.infolist()} == {compression} assert sum(entry.file_size for entry in archive.infolist()) <= len(result) * MAX_SKILL_ARCHIVE_COMPRESSION_RATIO assert build_uploaded_skill_archive(upload, "registered-name") == result - - -_VALID: Final = _zip((("SKILL.md", MANIFEST),)) -_ENCRYPTED: Final = _zip((("__MACOSX/encrypted", b"secret"), ("SKILL.md", MANIFEST))) -_CENTRAL: Final = _ENCRYPTED.index(b"PK\x01\x02") - - -@pytest.mark.parametrize("content, name", [ - (b"not a zip", "skill"), - (_zip((("SKILL.md", MANIFEST),), zipfile.ZIP_STORED).replace(b"Read guide.md", b"Dead guide.md"), "skill"), - (_ENCRYPTED[:_CENTRAL + 8] + b"\x01\x00" + _ENCRYPTED[_CENTRAL + 10:], "skill"), - *[(_VALID, name) for name in ("../escape", "", "Not a slug", "x" * 65536)], - *[(_zip(((path, MANIFEST),)), "skill") for path in ("readme.md", "wrapper/nested/SKILL.md", "skills/example/SKILL.md")], - *[(_zip((("SKILL.md", MANIFEST), (path, b"unsafe"))), "skill") for path in ( - "../escape", "/absolute", "C:/escape", "a\\escape", "__MACOSX/../escape", "bad/./file", - )], - *[(_zip((("SKILL.md", MANIFEST), (_entry(path, stat.S_IFLNK | 0o777), b"../../outside"))), "skill") - for path in ("__MACOSX/link", "scripts/link")], - (_zip(((".claude-plugin/plugin.json", b'{"name":"plugin"}'), ("skills/example/SKILL.md", MANIFEST))), "skill"), - (_zip((("SKILL.md", MANIFEST), ("__MACOSX/a", b"one"), ("__MACOSX//a", b"two"))), "skill"), - (_zip((("SKILL.md", MANIFEST), ("asset.bin", bytes(MAX_SKILL_UPLOAD_BYTES))), zipfile.ZIP_STORED), "skill"), - (_zip((("SKILL.md", MANIFEST), ("__MACOSX/large", bytes(MAX_ARCHIVE_UNPACKED_BYTES)))), "skill"), - (_zip((("SKILL.md", MANIFEST),) + tuple((f"__MACOSX/{i}", b"") for i in range(MAX_ARCHIVE_ENTRIES))), "skill"), -]) -def test_rejects_unsafe_invalid_or_oversized_archives(content: bytes, name: str) -> None: - assert build_uploaded_skill_archive(content, name) is None From 97aae6c65e8661964fc6a28a155d92277fb580e3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 22:13:07 +0000 Subject: [PATCH 09/19] fix(bedrock): remove GovCloud Kimi K3 bedrock_mantle rows (#50117) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 70 ------------------- model_prices_and_context_window.json | 70 ------------------- 2 files changed, 140 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2ebf3dd0ae1..5d91941a953 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -84635,41 +84635,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "bedrock_mantle/us-gov-east-1/moonshotai.kimi-k3": { - "cache_creation_input_token_cost": 4.95e-06, - "cache_creation_input_token_cost_flex": 2.475e-06, - "cache_creation_input_token_cost_priority": 8.6625e-06, - "cache_read_input_token_cost": 3.96e-07, - "cache_read_input_token_cost_flex": 1.98e-07, - "cache_read_input_token_cost_priority": 6.93e-07, - "input_cost_per_token": 3.96e-06, - "input_cost_per_token_batches": 1.98e-06, - "input_cost_per_token_flex": 1.98e-06, - "input_cost_per_token_priority": 6.93e-06, - "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, - "max_output_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1.98e-05, - "output_cost_per_token_batches": 9.9e-06, - "output_cost_per_token_flex": 9.9e-06, - "output_cost_per_token_priority": 3.465e-05, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_audio_input": false, - "supports_function_calling": true, - "supports_prompt_cache_breakpoint": false, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "token_count_framing": 64, - "token_count_tokenizer": "kimi_k3" - }, "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2.4e-07, "input_cost_per_token_batches": 1.2e-07, @@ -84766,41 +84731,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "bedrock_mantle/us-gov-west-1/moonshotai.kimi-k3": { - "cache_creation_input_token_cost": 4.95e-06, - "cache_creation_input_token_cost_flex": 2.475e-06, - "cache_creation_input_token_cost_priority": 8.6625e-06, - "cache_read_input_token_cost": 3.96e-07, - "cache_read_input_token_cost_flex": 1.98e-07, - "cache_read_input_token_cost_priority": 6.93e-07, - "input_cost_per_token": 3.96e-06, - "input_cost_per_token_batches": 1.98e-06, - "input_cost_per_token_flex": 1.98e-06, - "input_cost_per_token_priority": 6.93e-06, - "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, - "max_output_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1.98e-05, - "output_cost_per_token_batches": 9.9e-06, - "output_cost_per_token_flex": 9.9e-06, - "output_cost_per_token_priority": 3.465e-05, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_audio_input": false, - "supports_function_calling": true, - "supports_prompt_cache_breakpoint": false, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "token_count_framing": 64, - "token_count_tokenizer": "kimi_k3" - }, "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2.4e-07, "input_cost_per_token_batches": 1.2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2ebf3dd0ae1..5d91941a953 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -84635,41 +84635,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "bedrock_mantle/us-gov-east-1/moonshotai.kimi-k3": { - "cache_creation_input_token_cost": 4.95e-06, - "cache_creation_input_token_cost_flex": 2.475e-06, - "cache_creation_input_token_cost_priority": 8.6625e-06, - "cache_read_input_token_cost": 3.96e-07, - "cache_read_input_token_cost_flex": 1.98e-07, - "cache_read_input_token_cost_priority": 6.93e-07, - "input_cost_per_token": 3.96e-06, - "input_cost_per_token_batches": 1.98e-06, - "input_cost_per_token_flex": 1.98e-06, - "input_cost_per_token_priority": 6.93e-06, - "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, - "max_output_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1.98e-05, - "output_cost_per_token_batches": 9.9e-06, - "output_cost_per_token_flex": 9.9e-06, - "output_cost_per_token_priority": 3.465e-05, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_audio_input": false, - "supports_function_calling": true, - "supports_prompt_cache_breakpoint": false, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "token_count_framing": 64, - "token_count_tokenizer": "kimi_k3" - }, "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2.4e-07, "input_cost_per_token_batches": 1.2e-07, @@ -84766,41 +84731,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "bedrock_mantle/us-gov-west-1/moonshotai.kimi-k3": { - "cache_creation_input_token_cost": 4.95e-06, - "cache_creation_input_token_cost_flex": 2.475e-06, - "cache_creation_input_token_cost_priority": 8.6625e-06, - "cache_read_input_token_cost": 3.96e-07, - "cache_read_input_token_cost_flex": 1.98e-07, - "cache_read_input_token_cost_priority": 6.93e-07, - "input_cost_per_token": 3.96e-06, - "input_cost_per_token_batches": 1.98e-06, - "input_cost_per_token_flex": 1.98e-06, - "input_cost_per_token_priority": 6.93e-06, - "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, - "max_output_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1.98e-05, - "output_cost_per_token_batches": 9.9e-06, - "output_cost_per_token_flex": 9.9e-06, - "output_cost_per_token_priority": 3.465e-05, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_audio_input": false, - "supports_function_calling": true, - "supports_prompt_cache_breakpoint": false, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "token_count_framing": 64, - "token_count_tokenizer": "kimi_k3" - }, "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2.4e-07, "input_cost_per_token_batches": 1.2e-07, From 4e769d153a6f9a36fb5990669b7547380570b10c Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 5 Nov 2026 07:18:19 +0900 Subject: [PATCH 10/19] fix(auth): align cloud credential loading and expiry with the Python SDKs (#49589) * fix(auth): align cloud credential loading and expiry with the Python SDKs * feat(auth): resolve consumer-scoped default identities for AWS and Azure AwsAuthService::default_chain_credentials and AzureAuthService::identity_token give a consumer the SDK's own default identity without LiteLLM's Bedrock or Azure OpenAI settings. Moved here from #49510 so every shared auth crate change lives in one PR Co-Authored-By: Claude Opus 5.5 * fix(auth): include identity hints in secret snapshots --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 9 + litellm-rust/crates/auth-aws/Cargo.toml | 9 +- litellm-rust/crates/auth-aws/src/aws.rs | 1159 ++++++++++++----- litellm-rust/crates/auth-aws/src/constants.rs | 4 +- litellm-rust/crates/auth-aws/src/error.rs | 2 + litellm-rust/crates/auth-aws/src/signer.rs | 3 +- litellm-rust/crates/auth-azure/Cargo.toml | 6 +- litellm-rust/crates/auth-azure/src/lib.rs | 2 +- litellm-rust/crates/auth-azure/src/native.rs | 433 +++++- litellm-rust/crates/auth-azure/src/resolve.rs | 524 +++++++- litellm-rust/crates/auth-gcp/Cargo.toml | 8 +- litellm-rust/crates/auth-gcp/src/auth.rs | 428 +++--- litellm-rust/crates/auth-gcp/src/config.rs | 39 +- litellm-rust/crates/auth-gcp/src/identity.rs | 445 +++++++ litellm-rust/crates/auth-gcp/src/lib.rs | 6 +- .../crates/auth-gcp/tests/vertex_auth.rs | 120 -- litellm-rust/crates/auth-types/src/error.rs | 5 + .../crates/auth-types/src/params/aws.rs | 7 + .../crates/auth-types/tests/aws_params.rs | 12 +- litellm-rust/crates/cache-gcs/src/cache.rs | 4 +- litellm-rust/crates/cache-gcs/src/lib.rs | 2 +- litellm-rust/crates/cache-gcs/src/token.rs | 25 +- litellm-rust/crates/cache-s3/src/auth.rs | 9 +- litellm-rust/crates/cache-s3/tests/auth.rs | 6 +- .../crates/cache-s3/tests/support/mod.rs | 4 +- .../crates/inference-ocr/tests/resources.rs | 16 +- .../crates/secret-value/src/secret_value.rs | 12 + litellm-rust/crates/secrets-aws/src/auth.rs | 25 +- 28 files changed, 2467 insertions(+), 857 deletions(-) create mode 100644 litellm-rust/crates/auth-gcp/src/identity.rs delete mode 100644 litellm-rust/crates/auth-gcp/tests/vertex_auth.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b93a8827428..ba1131d7a92 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3389,15 +3389,19 @@ version = "0.1.0" dependencies = [ "aws-config", "aws-credential-types", + "aws-runtime", "aws-sdk-sts", "aws-sigv4", "aws-smithy-runtime-api", + "aws-smithy-types", "aws-types", "bytes", "futures-util", "litellm-auth-types", "litellm-error", "litellm-http", + "litellm-secret-value", + "litellm-secrets", "moka", "reqwest", "rstest", @@ -3406,6 +3410,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", ] @@ -3418,6 +3423,7 @@ dependencies = [ "litellm-auth-types", "litellm-error", "litellm-secret-value", + "litellm-secrets", "moka", "rstest", "serde_json", @@ -3431,6 +3437,7 @@ dependencies = [ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ + "axum", "base64 0.22.1", "gcp_auth", "google-cloud-auth", @@ -3440,10 +3447,12 @@ dependencies = [ "litellm-secret-value", "moka", "rstest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", "tempfile", + "time", "tokio", "veil", ] diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index e4578777a08..c8ed7bb0f14 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -11,14 +11,16 @@ litellm-error.workspace = true litellm-http.workspace = true reqwest.workspace = true -moka = { workspace = true, features = ["sync"] } +moka.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true -aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] } aws-sigv4 = "1.5.1" @@ -26,6 +28,9 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +litellm-secrets.workspace = true +aws-smithy-types = "1.6.4" +aws-runtime = { version = "1.9.4", default-features = false } rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } bytes.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 81aa4b068bc..e6e32754b4f 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,13 +1,16 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; -use moka::sync::Cache; -use sha2::{Digest, Sha256}; +use aws_credential_types::provider::{SharedCredentialsProvider, error::CredentialsError}; +use litellm_secret_value::SecretValue; +use moka::{Expiry, future::Cache}; +use std::{future::Future, pin::Pin, sync::Arc}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + SignableBody, SignableRequest, SignatureLocation, SigningInstructions, SigningParams, + SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -23,72 +26,106 @@ use super::constants::{ }; use crate::error::{CredentialFlow, Error, SigningStage}; -const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); -const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(600); +const EXPIRY_BUFFER: Duration = Duration::from_secs(60); #[derive(Clone)] pub struct AwsAuthService { - static_credentials: Cache, - ambient_credentials: Cache, + providers: Cache, + credentials: Cache, + loader: Arc, } impl Default for AwsAuthService { fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl AwsAuthService { + fn new(loader: Arc) -> Self { Self { - static_credentials: Cache::builder() + providers: Cache::builder().max_capacity(200).build(), + credentials: Cache::builder() .max_capacity(200) - .time_to_live(STATIC_CREDENTIALS_TTL) - .build(), - ambient_credentials: Cache::builder() - .max_capacity(200) - .time_to_live(AMBIENT_CREDENTIALS_TTL) + .expire_after(CredentialExpiry) .build(), + loader, } } } -fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { - match flow { - AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), - AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), - AwsAuthFlow::WebIdentity { .. } - | AwsAuthFlow::AssumeRole { .. } - | AwsAuthFlow::Profile { .. } - | AwsAuthFlow::SessionToken { .. } => None, +struct CredentialExpiry; + +impl Expiry for CredentialExpiry { + fn expire_after_create( + &self, + _: &AwsAuthConfig, + credentials: &Credentials, + _: std::time::Instant, + ) -> Option { + Some(credentials.expiry().map_or(CREDENTIAL_CACHE_TTL, |expiry| { + expiry + .duration_since(SystemTime::now()) + .unwrap_or_default() + .saturating_sub(EXPIRY_BUFFER) + .min(CREDENTIAL_CACHE_TTL) + })) } } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +type ProviderFuture<'a> = + Pin> + Send + 'a>>; + +trait ProviderLoader: Send + Sync { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_>; +} + +struct NativeProviderLoader; + +impl ProviderLoader for NativeProviderLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + Box::pin(build_provider(config, flow)) + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct AwsAuthConfig { - pub access_key_id: Option, - pub secret_access_key: Option, - pub session_token: Option, + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, pub region_name: Option, pub session_name: Option, pub profile_name: Option, pub role_name: Option, - pub web_identity_token: Option, + pub web_identity_token: Option, pub sts_endpoint: Option, - pub external_id: Option, + pub external_id: Option, } impl AwsAuthConfig { fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { - access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + access_key_id: self + .access_key_id + .or_else(|| env_lookup(AWS_ACCESS_KEY_ID).map(SecretValue::new)), secret_access_key: self .secret_access_key - .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), - session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY).map(SecretValue::new)), + session_token: self + .session_token + .or_else(|| env_lookup(AWS_SESSION_TOKEN).map(SecretValue::new)), region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), web_identity_token: self .web_identity_token - .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN).map(SecretValue::new)), sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), - external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + external_id: self + .external_id + .or_else(|| env_lookup(AWS_EXTERNAL_ID).map(SecretValue::new)), } } } @@ -96,52 +133,31 @@ impl AwsAuthConfig { #[derive(Clone, Debug, PartialEq, Eq)] pub enum AwsAuthFlow { WebIdentity { - token: String, + token: SecretValue, role: String, session_name: String, }, AssumeRole { role: String, session_name: Option, + web_identity_role: Option, }, Profile { name: String, }, SessionToken { - access_key_id: String, - secret_access_key: String, - session_token: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, + session_token: SecretValue, }, StaticKeys { - access_key_id: String, - secret_access_key: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, region_name: String, }, DefaultChain, } -fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { - let mut hasher = Sha256::new(); - hasher.update(format!("{config:?}:{flow:?}")); - format!("{:x}", hasher.finalize()) -} - -impl AwsAuthService { - fn get_cached_credentials(&self, key: &str) -> Option { - self.static_credentials - .get(key) - .or_else(|| self.ambient_credentials.get(key)) - } - - fn set_cached_credentials(&self, key: String, credentials: Credentials, ttl: Duration) { - if ttl == STATIC_CREDENTIALS_TTL { - self.static_credentials.insert(key, credentials); - } else { - self.ambient_credentials.insert(key, credentials); - } - } -} - fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { let mut parts = arn.splitn(6, ':'); let ("arn", partition, _, _, account, resource) = ( @@ -170,6 +186,9 @@ pub fn classify_auth( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { + let web_identity_role = env_lookup(AWS_WEB_IDENTITY_TOKEN_FILE) + .filter(|token_file| !token_file.is_empty()) + .and_then(|_| env_lookup(AWS_ROLE_ARN)); let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( config.web_identity_token.clone(), @@ -186,6 +205,7 @@ pub fn classify_auth( return AwsAuthFlow::AssumeRole { role, session_name: config.session_name.clone(), + web_identity_role, }; } if let Some(name) = config.profile_name { @@ -216,6 +236,30 @@ pub fn classify_auth( AwsAuthFlow::DefaultChain } +async fn profile_provider( + name: &str, + region: Option, + config: aws_config::provider_config::ProviderConfig, + builder: aws_config::profile::credentials::Builder, +) -> Result { + let region = match region { + Some(region) => Some(aws_types::region::Region::new(region)), + None => { + aws_config::default_provider::region::DefaultRegionChain::builder() + .profile_name(name) + .build() + .region() + .await + } + }; + Ok(SharedCredentialsProvider::new( + builder + .configure(&config.with_region(region)) + .profile_name(name) + .build(), + )) +} + impl AwsAuthService { pub async fn resolve_credentials( &self, @@ -224,183 +268,240 @@ impl AwsAuthService { ) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); - match flow { - AwsAuthFlow::SessionToken { - access_key_id, - secret_access_key, - session_token, - } => Ok(Credentials::new( - access_key_id, - secret_access_key, - Some(session_token), - None, - "litellm-static-session", - )), - AwsAuthFlow::StaticKeys { - access_key_id, - secret_access_key, - region_name, - } => { - let flow = AwsAuthFlow::StaticKeys { - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - region_name, - }; - let key = cache_key(&resolved, &flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let credentials = Credentials::new( - access_key_id, - secret_access_key, - None, - None, - "litellm-static", - ); - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), - ); - Ok(credentials) - } - AwsAuthFlow::Profile { name } => { - let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() - .profile_name(name) - .build(); + self.fresh_credentials(&resolved, &flow).await + } + + /// The credentials botocore's default chain finds in `region`, without LiteLLM's Bedrock settings + /// (`AWS_ROLE_NAME`, `AWS_PROFILE_NAME`, `AWS_SESSION_NAME`) that `resolve_credentials` also reads. + pub async fn default_chain_credentials(&self, region: &str) -> Result { + let config = AwsAuthConfig { + region_name: Some(region.to_owned()), + ..Default::default() + }; + self.fresh_credentials(&config, &AwsAuthFlow::DefaultChain) + .await + } + + async fn fresh_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + let cached = self.cached_credentials(resolved, flow).await?; + if fresh(&cached) { + return Ok(cached); + } + self.credentials.invalidate(resolved).await; + let refetched = self.cached_credentials(resolved, flow).await?; + if fresh(&refetched) { + return Ok(refetched); + } + self.credentials.invalidate(resolved).await; + Err(Error::ExpiredCredentials) + } + + async fn cached_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + self.credentials + .try_get_with(resolved.clone(), async { + let provider = self + .providers + .try_get_with( + resolved.clone(), + self.loader.load(resolved.clone(), flow.clone()), + ) + .await + .map_err(|error| (*error).clone())?; provider .provide_credentials() .await - .map_err(Error::credentials(CredentialFlow::Profile)) - } - AwsAuthFlow::AssumeRole { role, session_name } => { - if is_already_running_as_role(&role, &resolved).await? { - let ambient_flow = AwsAuthFlow::DefaultChain; - let key = cache_key(&resolved, &ambient_flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - return Ok(credentials); - } - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - if let (Some(access_key_id), Some(secret_access_key)) = - (resolved.access_key_id, resolved.secret_access_key) - { - loader = loader.credentials_provider(Credentials::new( - access_key_id, - secret_access_key, - resolved.session_token, - None, - "litellm-role-source", - )); - } - let sdk_config = loader.load().await; - let builder = aws_config::sts::AssumeRoleProvider::builder(role); - let builder = match session_name { - Some(name) => builder.session_name(name), - None => builder.session_name(default_session_name()), - }; - let builder = match resolved.external_id { - Some(id) => builder.external_id(id), - None => builder, - }; - let provider = builder.configure(&sdk_config).build().await; - provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::AssumeRole)) - } - AwsAuthFlow::WebIdentity { - token, - role, - session_name, - } => { - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let client = aws_sdk_sts::Client::new(&sdk_config); - let response = client - .assume_role_with_web_identity() - .role_arn(role) - .role_session_name(session_name) - .web_identity_token(token) - .send() - .await - .map_err(Error::credentials(CredentialFlow::WebIdentity))?; - let credentials = response - .credentials() - .ok_or(Error::MissingWebIdentityCredentials)?; - let expiration = - SystemTime::try_from(*credentials.expiration()).map_err(|source| { - Error::WebIdentityExpiration(litellm_error::Source::new(source)) - })?; - Ok(Credentials::new( - credentials.access_key_id(), - credentials.secret_access_key(), - Some(credentials.session_token().to_string()), - Some(expiration), - "litellm-web-identity", - )) - } - AwsAuthFlow::DefaultChain => { - let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&AwsAuthFlow::DefaultChain) - .unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - Ok(credentials) + .map_err(Error::credentials(flow.credential_flow())) + }) + .await + .map_err(|error| (*error).clone()) + } +} + +fn fresh(credentials: &Credentials) -> bool { + credentials + .expiry() + .is_none_or(|expiry| expiry > SystemTime::now() + EXPIRY_BUFFER) +} + +impl AwsAuthFlow { + fn credential_flow(&self) -> CredentialFlow { + match self { + Self::Profile { .. } => CredentialFlow::Profile, + Self::AssumeRole { .. } => CredentialFlow::AssumeRole, + Self::WebIdentity { .. } => CredentialFlow::WebIdentity, + Self::DefaultChain | Self::StaticKeys { .. } | Self::SessionToken { .. } => { + CredentialFlow::DefaultChain } } } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { +async fn build_provider( + resolved: AwsAuthConfig, + flow: AwsAuthFlow, +) -> Result { + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + Some(session_token.expose().to_owned()), + None, + "litellm-static-session", + ))), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + .. + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + None, + None, + "litellm-static", + ))), + AwsAuthFlow::Profile { name } => { + profile_provider( + &name, + resolved.region_name, + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder(), + ) + .await + } + AwsAuthFlow::DefaultChain => { + let builder = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder(); + let builder = match resolved.region_name { + Some(region) => builder.region(aws_types::region::Region::new(region)), + None => builder, + }; + Ok(SharedCredentialsProvider::new(builder.build().await)) + } + AwsAuthFlow::AssumeRole { + role, + session_name, + web_identity_role, + } => { + if is_already_running_as_role(&role, web_identity_role.as_deref(), &resolved).await? { + return Ok(SharedCredentialsProvider::new( + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await, + )); + } + let sdk_config = sdk_config(resolved.clone()).await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role) + .session_name(session_name.unwrap_or_else(default_session_name)); + let builder = match resolved.external_id { + Some(id) => builder.external_id(id.expose()), + None => builder, + }; + Ok(SharedCredentialsProvider::new( + builder.configure(&sdk_config).build().await, + )) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => Ok(SharedCredentialsProvider::new(WebIdentityProvider { + client: aws_sdk_sts::Client::new(&sdk_config(resolved).await), + token, + role, + session_name, + })), + } +} + +async fn sdk_config(config: AwsAuthConfig) -> aws_types::SdkConfig { + let loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + let loader = match config.region_name { + Some(region) => loader.region(aws_types::region::Region::new(region)), + None => loader, + }; + let loader = match config.sts_endpoint { + Some(endpoint) => loader.endpoint_url(endpoint), + None => loader, + }; + let loader = match (config.access_key_id, config.secret_access_key) { + (Some(key), Some(secret)) => loader.credentials_provider(Credentials::new( + key.expose(), + secret.expose(), + config.session_token.map(|token| token.expose().to_owned()), + None, + "litellm-role-source", + )), + _ => loader, + }; + loader.load().await +} + +#[derive(Debug)] +struct WebIdentityProvider { + client: aws_sdk_sts::Client, + token: SecretValue, + role: String, + session_name: String, +} + +impl ProvideCredentials for WebIdentityProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async { + let response = self + .client + .assume_role_with_web_identity() + .role_arn(&self.role) + .role_session_name(&self.session_name) + .web_identity_token(self.token.expose()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let credentials = response.credentials().ok_or_else(|| { + CredentialsError::provider_error(Error::MissingWebIdentityCredentials) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()) + .map_err(CredentialsError::provider_error)?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_owned()), + Some(expiration), + "litellm-web-identity", + )) + }) + } +} + +async fn is_already_running_as_role( + role: &str, + web_identity_role: Option<&str>, + config: &AwsAuthConfig, +) -> Result { + if config.access_key_id.is_some() || config.secret_access_key.is_some() { + return Ok(false); + } if role_identity(role).is_none() { return Ok(false); } - if let (Ok(current_role), Ok(token_file)) = ( - std::env::var(AWS_ROLE_ARN), - std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) && !token_file.is_empty() - { - return Ok(same_role_arns(role, ¤t_role)); + if let Some(current_role) = web_identity_role { + return Ok(same_role_arns(role, current_role)); } let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); @@ -473,16 +574,6 @@ pub fn sign_headers( headers: &HeaderMap, scope: SigningScope<'_>, ) -> Result { - let identity: Identity = scope.credentials.clone().into(); - let params = v4::SigningParams::builder() - .identity(&identity) - .region(scope.region) - .name(scope.service) - .time(scope.time) - .settings(SigningSettings::default()) - .build() - .map(SigningParams::from) - .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() .map(|(name, value)| { @@ -494,16 +585,14 @@ pub fn sign_headers( }) }) .collect::, _>>()?; - let request = SignableRequest::new( + let instructions = sign_request( method, url, header_refs.into_iter(), - SignableBody::Bytes(body), - ) - .map_err(Error::signing(SigningStage::Request))?; - let (instructions, _) = sign(request, ¶ms) - .map_err(Error::signing(SigningStage::Signature))? - .into_parts(); + body, + SigningSettings::default(), + scope, + )?; instructions .headers() .map(|(name, value)| { @@ -516,6 +605,60 @@ pub fn sign_headers( .collect() } +pub fn presign_url( + url: &str, + expires_in: Duration, + scope: SigningScope<'_>, +) -> Result { + let mut signed = url::Url::parse(url).map_err(Error::signing(SigningStage::Request))?; + let mut settings = SigningSettings::default(); + settings.signature_location = SignatureLocation::QueryParams; + settings.expires_in = Some(expires_in); + let instructions = sign_request( + "GET", + signed.as_str(), + std::iter::empty(), + &[], + settings, + scope, + )?; + let (signature, auth): (Vec<_>, Vec<_>) = instructions + .params() + .iter() + .map(|(name, value)| (*name, value.as_ref())) + .partition(|(name, _)| *name == "X-Amz-Signature"); + signed + .query_pairs_mut() + .extend_pairs(auth.into_iter().chain(signature)); + Ok(signed.into()) +} + +fn sign_request<'a>( + method: &'a str, + url: &'a str, + headers: impl Iterator, + body: &'a [u8], + settings: SigningSettings, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(scope.region) + .name(scope.service) + .time(scope.time) + .settings(settings) + .build() + .map(SigningParams::from) + .map_err(Error::signing(SigningStage::Parameters))?; + let request = SignableRequest::new(method, url, headers, SignableBody::Bytes(body)) + .map_err(Error::signing(SigningStage::Request))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(Error::signing(SigningStage::Signature))? + .into_parts(); + Ok(instructions) +} + /// Model-id and region parsing shared by every Bedrock route. pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { let mut stripped = model; @@ -586,16 +729,26 @@ pub fn aws_auth_config( env_lookup: &dyn Fn(&str) -> Option, ) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: params.resolve(&AwsParams::ACCESS_KEY_ID, env_lookup), - secret_access_key: params.resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup), - session_token: params.resolve(&AwsParams::SESSION_TOKEN, env_lookup), + access_key_id: params + .resolve(&AwsParams::ACCESS_KEY_ID, env_lookup) + .map(SecretValue::new), + secret_access_key: params + .resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup) + .map(SecretValue::new), + session_token: params + .resolve(&AwsParams::SESSION_TOKEN, env_lookup) + .map(SecretValue::new), region_name: params.resolve(&AwsParams::REGION, env_lookup), session_name: params.resolve(&AwsParams::SESSION_NAME, env_lookup), profile_name: params.resolve(&AwsParams::PROFILE_NAME, env_lookup), role_name: params.resolve(&AwsParams::ROLE_NAME, env_lookup), - web_identity_token: params.resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup), + web_identity_token: params + .resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup) + .map(SecretValue::new), sts_endpoint: params.resolve(&AwsParams::STS_ENDPOINT, env_lookup), - external_id: params.resolve(&AwsParams::EXTERNAL_ID, env_lookup), + external_id: params + .resolve(&AwsParams::EXTERNAL_ID, env_lookup) + .map(SecretValue::new), } } @@ -660,10 +813,133 @@ mod tests { use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; use reqwest::header::CONTENT_TYPE; + #[cfg(unix)] + #[rstest::fixture] + fn process_profile() -> aws_runtime::env_config::file::EnvConfigFiles { + let credentials = serde_json::json!({ + "Version": 1, + "AccessKeyId": "process-access", + "SecretAccessKey": "process-secret", + "SessionToken": "process-session", + }); + aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + format!( + "[profile smoke]\ncredential_process = /usr/bin/printf '%s' '{credentials}'\n" + ), + ) + .build() + } + + #[cfg(unix)] + #[rstest::rstest] + #[tokio::test] + async fn profile_process_credentials_can_be_loaded( + process_profile: aws_runtime::env_config::file::EnvConfigFiles, + ) { + let provider = profile_provider( + "smoke", + Some("us-west-2".into()), + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_files(process_profile), + ) + .await + .unwrap(); + let credentials = provider.provide_credentials().await.unwrap(); + + assert_eq!(credentials.access_key_id(), "process-access"); + assert_eq!(credentials.secret_access_key(), "process-secret"); + assert_eq!(credentials.session_token(), Some("process-session")); + } + + #[derive(Debug)] + struct StsConnector; + + impl aws_smithy_runtime_api::client::http::HttpConnector for StsConnector { + fn call( + &self, + request: aws_smithy_runtime_api::client::orchestrator::HttpRequest, + ) -> aws_smithy_runtime_api::client::http::HttpConnectorFuture { + assert!(request.uri().contains("sts.us-west-2.")); + aws_smithy_runtime_api::client::http::HttpConnectorFuture::new(async { + Ok( + aws_smithy_runtime_api::client::orchestrator::HttpResponse::new( + 200_u16.try_into().unwrap(), + aws_smithy_types::body::SdkBody::from( + r#" + + role-accessrole-secret + role-session2100-01-01T00:00:00Z + arn:aws:sts::123456789012:assumed-role/smoke/session + smoke:session + smoke"#, + ), + ), + ) + }) + } + } + + #[rstest::rstest] + #[tokio::test] + async fn role_profile_uses_region_when_acquiring_sts_credentials() { + use aws_smithy_runtime_api::client::http::{SharedHttpConnector, http_client_fn}; + let files = aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + "[profile smoke]\nrole_arn = arn:aws:iam::123456789012:role/smoke\nsource_profile = source\n\n[profile source]\naws_access_key_id = source-access\naws_secret_access_key = source-secret\n", + ).build(); + let config = aws_config::provider_config::ProviderConfig::without_region() + .with_http_client(http_client_fn(|_, _| { + SharedHttpConnector::new(StsConnector) + })); + let credentials = profile_provider( + "smoke", + Some("us-west-2".into()), + config, + aws_config::profile::ProfileFileCredentialsProvider::builder().profile_files(files), + ) + .await + .unwrap(); + let credentials = credentials.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "role-access"); + assert_eq!(credentials.secret_access_key(), "role-secret"); + assert_eq!(credentials.session_token(), Some("role-session")); + assert!(credentials.expiry().is_some()); + } + fn no_env(_: &str) -> Option { None } + #[rstest::rstest] + #[case::access_key(Some("source-access"), None)] + #[case::secret_key(None, Some("source-secret"))] + #[case::source_pair(Some("source-access"), Some("source-secret"))] + #[tokio::test] + async fn explicit_source_keys_cannot_reuse_an_ambient_role( + #[case] access_key: Option<&str>, + #[case] secret_key: Option<&str>, + ) { + let role = "arn:aws:iam::123456789012:role/target"; + let config = AwsAuthConfig { + access_key_id: access_key.map(SecretValue::new), + secret_access_key: secret_key.map(SecretValue::new), + ..Default::default() + }; + assert!( + !is_already_running_as_role(role, Some(role), &config) + .await + .unwrap() + ); + } + #[test] fn secret_names_cover_environment_reads() { let seen = std::sync::Arc::new(std::sync::Mutex::new( @@ -706,7 +982,10 @@ mod tests { }; assert_eq!( - aws_auth_config(¶ms, &env).secret_access_key.as_deref(), + aws_auth_config(¶ms, &env) + .secret_access_key + .as_ref() + .map(SecretValue::expose), expected ); assert_eq!( @@ -831,6 +1110,53 @@ mod tests { ); } + #[rstest::rstest] + #[case::running_with_web_identity("/var/run/token", Some("arn:aws:iam::1:role/current"))] + #[case::empty_token_file("", None)] + #[tokio::test] + async fn assume_role_reads_the_ambient_role_from_the_secret_snapshot( + #[case] token_file: &str, + #[case] expected: Option<&str>, + ) { + use litellm_secrets::source::SecretSource; + + struct WebIdentitySecrets<'a>(&'a str); + + impl SecretSource for WebIdentitySecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> futures_util::future::BoxFuture< + 'a, + Result, litellm_secrets::Error>, + > { + Box::pin(async move { + Ok(match name { + AWS_WEB_IDENTITY_TOKEN_FILE => Some(SecretValue::new(self.0)), + AWS_ROLE_ARN => Some(SecretValue::new("arn:aws:iam::1:role/current")), + _ => None, + }) + }) + } + } + + let snapshot = WebIdentitySecrets(token_file) + .resolve(AwsParams::secret_names()) + .await + .unwrap(); + let config = AwsAuthConfig { + role_name: Some("arn:aws:iam::1:role/target".into()), + ..Default::default() + }; + let AwsAuthFlow::AssumeRole { + web_identity_role, .. + } = classify_auth(config, &|name| snapshot.get(name)) + else { + panic!("a role name selects AssumeRole"); + }; + assert_eq!(web_identity_role.as_deref(), expected); + } + #[tokio::test] async fn static_credentials_do_not_use_network() { let credentials = AwsAuthService::default() @@ -849,110 +1175,231 @@ mod tests { assert_eq!(credentials.session_token(), None); } - #[test] - fn cache_policy_matches_python_flows() { - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::StaticKeys { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - region_name: "us-east-1".into(), - }), - Some(STATIC_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::DefaultChain), - Some(AMBIENT_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::SessionToken { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - session_token: "token".into(), - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::Profile { - name: "profile".into() - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::AssumeRole { - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: None, - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::WebIdentity { - token: "token".into(), - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: "session".into(), - }), - None - ); + struct CountingLoader { + loads: Arc, + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + } + + impl ProviderLoader for CountingLoader { + fn load(&self, config: AwsAuthConfig, _: AwsAuthFlow) -> ProviderFuture<'_> { + let loads = self.loads.clone(); + let provider = CountingProvider { + calls: self.calls.clone(), + short_lifetime: self.short_lifetime, + short_calls: self.short_calls, + key: config.access_key_id.unwrap(), + }; + Box::pin(async move { + loads.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(SharedCredentialsProvider::new(provider)) + }) + } + } + + #[derive(Debug)] + struct CountingProvider { + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + key: SecretValue, + } + + impl ProvideCredentials for CountingProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async move { + tokio::task::yield_now().await; + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let lifetime = if call < self.short_calls { + self.short_lifetime + } else { + Duration::from_secs(3600) + }; + Ok(Credentials::new( + self.key.expose(), + "credential-value", + Some(format!("session-{call}")), + Some(SystemTime::now() + lifetime), + "counting", + )) + }) + } + } + + #[derive(Default)] + struct RecordingLoader { + loads: std::sync::Mutex>, + } + + impl ProviderLoader for RecordingLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + self.loads.lock().unwrap().push((config, flow)); + Box::pin(async { + Ok(SharedCredentialsProvider::new(Credentials::new( + "ak", + "sk", + None, + None, + "recording", + ))) + }) + } } #[rstest::rstest] - fn cache_round_trip_preserves_credentials() { - let auth = AwsAuthService::default(); - let key = format!("cache-test-{}", std::process::id()); - let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); - auth.set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + #[tokio::test] + async fn default_chain_credentials_do_not_share_a_configured_role_in_the_same_region() { + let loader = Arc::new(RecordingLoader::default()); + let auth = AwsAuthService::new(loader.clone()); + let bedrock_env = |name: &str| match name { + AWS_ROLE_NAME => Some("arn:aws:iam::123456789012:role/bedrock".to_string()), + AWS_REGION_NAME => Some("us-east-1".to_string()), + _ => None, + }; + auth.resolve_credentials(AwsAuthConfig::default(), &bedrock_env) + .await + .unwrap(); + + auth.default_chain_credentials("us-east-1").await.unwrap(); + auth.default_chain_credentials("us-east-1").await.unwrap(); + + let loads = loader.loads.lock().unwrap(); + assert_eq!(loads.len(), 2); assert_eq!( - auth.get_cached_credentials(&key) - .map(|value| value.access_key_id().to_string()), - Some("cache-ak".to_string()) + loads[1], + ( + AwsAuthConfig { + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + AwsAuthFlow::DefaultChain + ) ); } #[rstest::rstest] #[tokio::test] - async fn cloned_services_reuse_credentials_but_independent_services_do_not() { - let auth = AwsAuthService::default(); + async fn near_expiry_credentials_are_refreshed_without_rebuilding_the_provider() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(30), + short_calls: 1, + })); let config = AwsAuthConfig { - access_key_id: Some("configured-key".into()), - secret_access_key: Some("configured-secret".into()), - region_name: Some("us-east-1".into()), - ..AwsAuthConfig::default() + access_key_id: Some("identity".into()), + ..Default::default() }; - let flow = classify_auth(config.clone(), &no_env); - let cached = Credentials::new("cached-key", "cached-secret", None, None, "test"); - auth.set_cached_credentials( - cache_key(&config, &flow), - cached.clone(), - STATIC_CREDENTIALS_TTL, - ); - - let reused = auth + let first = auth + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let second = auth .clone() .resolve_credentials(config.clone(), &no_env) .await .unwrap(); - let independent = AwsAuthService::default() - .resolve_credentials(config.clone(), &no_env) - .await - .unwrap(); - let different = AwsAuthConfig { - access_key_id: Some("different-key".into()), - ..config.clone() - }; - let other_identity = auth - .resolve_credentials(different.clone(), &no_env) - .await - .unwrap(); + assert_eq!(first.session_token(), Some("session-1")); + assert_eq!(first.session_token(), second.session_token()); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } - assert_eq!(reused.access_key_id(), cached.access_key_id()); - assert_eq!(reused.secret_access_key(), cached.secret_access_key()); + #[rstest::rstest] + #[case::recovered_by_one_refetch(1, Some("session-1"))] + #[case::still_expired_after_refetch(2, None)] + #[tokio::test] + async fn expired_credentials_are_refetched_once( + #[case] short_calls: usize, + #[case] session: Option<&str>, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: Arc::new(AtomicUsize::new(0)), + calls: calls.clone(), + short_lifetime: Duration::ZERO, + short_calls, + })); + let config = AwsAuthConfig { + access_key_id: Some("identity".into()), + ..Default::default() + }; + let result = auth.resolve_credentials(config, &no_env).await; assert_eq!( - Some(independent.access_key_id()), - config.access_key_id.as_deref() + result.as_ref().ok().and_then(Credentials::session_token), + session ); - assert_eq!( - Some(other_identity.access_key_id()), - different.access_key_id.as_deref() + assert!(session.is_some() || matches!(result, Err(Error::ExpiredCredentials))); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test] + async fn concurrent_calls_share_acquisition_and_distinct_identities_stay_isolated() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(3600), + short_calls: 0, + })); + let first = AwsAuthConfig { + access_key_id: Some("first".into()), + ..Default::default() + }; + let second = AwsAuthConfig { + access_key_id: Some("second".into()), + ..Default::default() + }; + let (a, b, c) = tokio::join!( + auth.resolve_credentials(first.clone(), &no_env), + auth.resolve_credentials(first, &no_env), + auth.resolve_credentials(second, &no_env) ); + let a = a.unwrap(); + let b = b.unwrap(); + let c = c.unwrap(); + assert_eq!(a.session_token(), b.session_token()); + assert_ne!(a.access_key_id(), c.access_key_id()); + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + fn auth_config_debug_never_discloses_credentials() { + let config = AwsAuthConfig { + access_key_id: Some("credential-value".into()), + secret_access_key: Some("credential-value".into()), + session_token: Some("credential-value".into()), + web_identity_token: Some("credential-value".into()), + external_id: Some("credential-value".into()), + ..Default::default() + }; + let output = format!("{config:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); + } + + #[rstest::rstest] + #[case::session(AwsAuthFlow::SessionToken { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), session_token: "credential-value".into() })] + #[case::static_keys(AwsAuthFlow::StaticKeys { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), region_name: "region".into() })] + #[case::web_identity(AwsAuthFlow::WebIdentity { token: "credential-value".into(), role: "role".into(), session_name: "session".into() })] + fn auth_flow_debug_never_discloses_credentials(#[case] flow: AwsAuthFlow) { + let output = format!("{flow:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); } #[test] @@ -1142,8 +1589,8 @@ mod tests { let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { - access_key_id: Some(access_key_id), - secret_access_key: Some(secret_access_key), + access_key_id: Some(access_key_id.into()), + secret_access_key: Some(secret_access_key.into()), region_name: Some("us-west-2".to_string()), ..Default::default() }, @@ -1191,4 +1638,44 @@ mod tests { failures.join("; ") ); } + + #[rstest::rstest] + #[case::provisioned( + None, + "https://cache-1/?Action=connect&User=iam-user", + "cache-1/?Action=connect&User=iam-user&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Signature=dfa58492a6ae9ba065088724170a52437329d03f3e9bc9a4c90a80cc8a18e2f6" + )] + #[case::serverless_with_a_session_token( + Some("session-token"), + "https://cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache", + "cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Security-Token=session-token&X-Amz-Signature=e48bdbb3bf9d5e8f1d325bf33d5260e333df1f6b8868e8edd1e4c3cc01b4c3f9" + )] + fn presigned_urls_match_botocore_sigv4_query_auth( + #[case] session_token: Option<&str>, + #[case] url: &str, + #[case] expected: &str, + ) { + // Expected values from botocore 1.43.117 SigV4QueryAuth with get_current_datetime patched to + // 2026-10-29T12:00:00Z (generated 2026-10-29), the signer Python's ElastiCache IAM provider uses + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + session_token.map(str::to_owned), + None, + "test", + ); + let signed = presign_url( + url, + Duration::from_secs(900), + SigningScope { + region: "us-east-1", + service: "elasticache", + credentials: &credentials, + time: UNIX_EPOCH + Duration::from_secs(1_793_275_200), + }, + ) + .unwrap(); + + assert_eq!(signed.strip_prefix("https://"), Some(expected)); + } } diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 789fde156bd..0e4662eda9e 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0]; pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0]; pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0]; pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0]; -pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; -pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV; +pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV; pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0]; pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0]; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index 5f88fc8a171..5445de1e6c3 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -2,6 +2,8 @@ use litellm_error::Source; #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { + #[error("AWS credential provider returned expired credentials")] + ExpiredCredentials, #[error("could not load AWS {flow} credentials")] Credentials { flow: CredentialFlow, diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index dff917d7c13..2d9ad9ba123 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer { ) .map_err(|error| match error { Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), - Error::Credentials { .. } + Error::ExpiredCredentials + | Error::Credentials { .. } | Error::MissingWebIdentityCredentials | Error::WebIdentityExpiration(_) | Error::Signing { .. } => litellm_http::Error::signature(error), diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 7dcb0e513a7..36c10a74c80 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -8,17 +8,19 @@ repository.workspace = true [dependencies] litellm-secret-value.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true moka.workspace = true serde_json.workspace = true sha2.workspace = true strum.workspace = true +tokio.workspace = true url.workspace = true azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] -litellm-error.workspace = true +litellm-secrets.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 9ed505e4160..6319406ce9c 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::{AzureAuthService, SECRET_NAMES}; +pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index f40745ab7a2..03021428ea7 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest { pub(crate) struct ValidatedAzureRequest { request: NativeAzureRequest, credential_source: InputSource, + probe_imds: bool, } impl ValidatedAzureRequest { @@ -68,9 +69,18 @@ impl ValidatedAzureRequest { Ok(Self { request, credential_source, + probe_imds: false, }) } + pub(crate) fn with_imds_probe(self, enabled: bool) -> Self { + Self { + probe_imds: enabled + && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }), + ..self + } + } + pub(crate) fn credential_source(&self) -> InputSource { self.credential_source } @@ -87,9 +97,34 @@ impl ValidatedAzureRequest { } } +pub(crate) enum CredentialAttempt { + Resolved(ResolvedCredential), + Unavailable(Error), + Failed(Error), +} + +impl CredentialAttempt { + pub(crate) fn into_result(self) -> Result { + match self { + Self::Resolved(credential) => Ok(credential), + Self::Unavailable(error) | Self::Failed(error) => Err(error), + } + } +} + pub(crate) struct NativeAzureTokenAcquirer { cache: AzureCredentialProviderCache, transport: Option, + imds_available: tokio::sync::Mutex, +} + +enum ImdsAvailability { + Unknown, + Available, + Unavailable { + retry_at: tokio::time::Instant, + error: Error, + }, } impl Default for NativeAzureTokenAcquirer { @@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: None, + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } @@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: Some(transport), + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } - pub(crate) async fn acquire( + pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt { + let discover_imds = request.probe_imds; + if discover_imds && let Err(error) = self.discover_imds().await { + return CredentialAttempt::Unavailable(error); + } + match self.acquire_token(request).await { + Ok(credential) => CredentialAttempt::Resolved(credential), + Err(error) => { + let missing_identity = discover_imds && unassigned_identity(&error); + if missing_identity { + CredentialAttempt::Unavailable(error) + } else { + CredentialAttempt::Failed(error) + } + } + } + } + + async fn discover_imds(&self) -> Result<(), Error> { + let mut availability = self.imds_available.lock().await; + match &*availability { + ImdsAvailability::Available => return Ok(()), + ImdsAvailability::Unavailable { retry_at, error } + if tokio::time::Instant::now() < *retry_at => + { + return Err(error.clone()); + } + ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {} + } + let result = probe_imds(self.transport.clone().unwrap_or_default()).await; + *availability = match &result { + Ok(()) => ImdsAvailability::Available, + Err(error) => ImdsAvailability::Unavailable { + retry_at: tokio::time::Instant::now() + Duration::from_secs(30), + error: error.clone(), + }, + }; + result + } + + async fn acquire_token( &self, request: ValidatedAzureRequest, ) -> Result { @@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer { } } +fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool { + std::iter::successors(Some(error), |cause| cause.source()).any(|cause| { + cause + .downcast_ref::() + .or_else(|| { + cause + .downcast_ref::()? + .downcast_ref::() + }) + .is_some_and(|error| { + error.http_status() == Some(azure_core::http::StatusCode::BadRequest) + }) + }) +} + +async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> { + use azure_core::http::{Context, Method, Request, Url}; + + let endpoint = + Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01") + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error)) + })?; + let mut request = Request::new(endpoint, Method::Get); + let response = tokio::time::timeout( + Duration::from_secs(1), + transport.send(&Context::new(), &mut request), + ) + .await + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })? + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })?; + // Headerless IMDS discovery expects 400, checked 2026-11-01 + // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging + if response.status() != azure_core::http::StatusCode::BadRequest { + return Err(Error::CredentialAcquisition( + ErrorDetail::UnexpectedStatus { + operation: "Azure IMDS availability probe", + status: u16::from(response.status()), + }, + )); + } + Ok(()) +} + impl NativeAzureRequest { fn scope(&self) -> &str { match self { @@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion { #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; + use std::time::Duration; use azure_core::http::headers::Headers; use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; @@ -569,6 +695,296 @@ mod tests { } } + #[derive(Debug)] + struct AvailableImdsClient { + probe_status: StatusCode, + probe_delay: Duration, + recover_after: Option, + response_status: Option, + probes: std::sync::atomic::AtomicUsize, + } + + impl Default for AvailableImdsClient { + fn default() -> Self { + Self { + probe_status: StatusCode::BadRequest, + probe_delay: Duration::ZERO, + recover_after: None, + response_status: None, + probes: std::sync::atomic::AtomicUsize::new(0), + } + } + } + + impl HttpClient for AvailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if request + .headers() + .get_optional_str(&azure_core::http::headers::HeaderName::from_static( + "metadata", + )) + .is_none() + { + let probe = self + .probes + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if self.recover_after.is_some_and(|failures| probe >= failures) { + return Ok(AsyncRawResponse::from_bytes( + StatusCode::BadRequest, + Headers::new(), + "metadata header required", + )); + } + tokio::time::sleep(self.probe_delay).await; + return Ok(AsyncRawResponse::from_bytes( + self.probe_status, + Headers::new(), + "metadata header required", + )); + } + if let Some(status) = self.response_status { + return Ok(AsyncRawResponse::from_bytes( + status, + Headers::new(), + "identity endpoint failure", + )); + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let expires = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + serde_json::json!({ + "access_token": "managed-token", "expires_on": expires.to_string(), + "token_type": "Bearer", "resource": "https://service.test", + }) + .to_string(), + )) + }) + } + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn available_imds_retains_full_sdk_token_acquisition() { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: None, + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + let credential = acquirer.acquire(request).await.into_result().unwrap(); + assert_eq!(credential.secret().expose(), "managed-token"); + } + + #[rstest::rstest] + #[case::another_metadata_service(StatusCode::NotFound)] + #[case::token_required(StatusCode::Unauthorized)] + #[case::proxy_rejection(StatusCode::Forbidden)] + #[case::unavailable(StatusCode::ServiceUnavailable)] + #[tokio::test(start_paused = true)] + async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition( + #[case] status: StatusCode, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + assert!(matches!( + acquirer.acquire(request).await, + super::CredentialAttempt::Unavailable(_) + )); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn an_imds_timeout_is_not_repeated_on_the_next_request() { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request)) + .await + .unwrap(); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::fixture] + fn imds_request() -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true) + } + + #[rstest::rstest] + #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))] + #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)] + #[tokio::test(start_paused = true)] + async fn failed_imds_discovery_recovers_without_restarting_the_acquirer( + imds_request: ValidatedAzureRequest, + #[case] status: StatusCode, + #[case] delay: Duration, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + probe_delay: delay, + recover_after: Some(1), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert_eq!( + first.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!( + second.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!( + acquirer + .acquire(imds_request) + .await + .into_result() + .unwrap() + .secret() + .expose(), + "managed-token", + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn a_persistent_imds_outage_reuses_each_failed_probe( + imds_request: ValidatedAzureRequest, + ) { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + assert!(matches!( + acquirer.acquire(imds_request.clone()).await, + super::CredentialAttempt::Unavailable(_), + )); + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)] + #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)] + #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)] + #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)] + #[tokio::test(start_paused = true)] + async fn only_unavailable_default_identities_allow_fallback( + #[case] discovery: bool, + #[case] status: StatusCode, + #[case] unavailable: bool, + ) { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: Some(status), + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(discovery); + let outcome = acquirer.acquire(request).await; + let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_)); + let error = outcome.into_result().unwrap_err(); + assert_eq!( + is_unavailable, + unavailable, + "{}", + litellm_error::Report::new(&error) + ); + } + + #[rstest::rstest] #[tokio::test] async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { let transport = Arc::new(RecordingTokenClient::default()); @@ -582,8 +998,12 @@ mod tests { "https://login.test", ); - let first = acquirer.acquire(request.clone()).await.unwrap(); - let second = acquirer.acquire(request).await.unwrap(); + let first = acquirer + .acquire(request.clone()) + .await + .into_result() + .unwrap(); + let second = acquirer.acquire(request).await.into_result().unwrap(); assert_eq!(first.secret().expose(), "native-token"); assert_eq!(second.secret().expose(), "native-token"); @@ -599,6 +1019,7 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] async fn credential_provider_cache_isolates_every_client_secret_identity_field() { let transport = Arc::new(RecordingTokenClient::default()); @@ -645,10 +1066,10 @@ mod tests { ), ]; - acquirer.acquire(base.clone()).await.unwrap(); - acquirer.acquire(base).await.unwrap(); + acquirer.acquire(base.clone()).await.into_result().unwrap(); + acquirer.acquire(base).await.into_result().unwrap(); for request in variants.into_iter().skip(1) { - acquirer.acquire(request).await.unwrap(); + acquirer.acquire(request).await.into_result().unwrap(); } assert_eq!(transport.requests.lock().unwrap().len(), 6); diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index a92630af26a..9ff4eda540a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,3 +1,4 @@ +use crate::native::CredentialAttempt; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, TokenProviderHandle, @@ -12,25 +13,31 @@ use std::sync::Arc; use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; -const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; -const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; -const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID"; -const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET"; -const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE"; -const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; -const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; -const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)] +enum AzureSecretName { + #[strum(serialize = "AZURE_AD_TOKEN")] + AdToken, + #[strum(serialize = "AZURE_TENANT_ID")] + TenantId, + #[strum(serialize = "AZURE_CLIENT_ID")] + ClientId, + #[strum(serialize = "AZURE_CLIENT_SECRET")] + ClientSecret, + #[strum(serialize = "AZURE_SCOPE")] + Scope, + #[strum(serialize = "AZURE_AUTHORITY_HOST")] + AuthorityHost, + #[strum(serialize = "AZURE_CREDENTIAL")] + Credential, + #[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")] + FederatedTokenFile, + #[strum(serialize = "IDENTITY_ENDPOINT")] + IdentityEndpoint, + #[strum(serialize = "MSI_ENDPOINT")] + MsiEndpoint, +} -pub const SECRET_NAMES: &[&str] = &[ - AZURE_AD_TOKEN_ENV, - AZURE_TENANT_ID_ENV, - AZURE_CLIENT_ID_ENV, - AZURE_CLIENT_SECRET_ENV, - AZURE_SCOPE_ENV, - AZURE_AUTHORITY_HOST_ENV, - AZURE_CREDENTIAL_ENV, - AZURE_FEDERATED_TOKEN_FILE_ENV, -]; +pub const SECRET_NAMES: &[&str] = ::VARIANTS; #[derive(Clone, Debug)] pub(crate) enum AzureCredentialPlan { @@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>>; + ) -> Pin + Send + '_>>; } impl AzureTokenAcquirer for NativeAzureTokenAcquirer { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>> { + ) -> Pin + Send + '_>> { Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) } } @@ -83,12 +90,36 @@ impl AzureAuthService { Self { native } } + pub async fn identity_token( + &self, + identity: &AzureIdentity, + scope: &str, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let plan = identity_plan(identity, scope, env_lookup)?; + let inputs = AzureAuthInputs::default_credential_for_scope(scope); + Ok(self + .run_plan(plan, &inputs, env_lookup) + .await? + .map(Sourced::into_value)) + } + pub async fn get_azure_ad_token( &self, inputs: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result>, Error> { - match select_auth_plan(inputs, env_lookup)? { + self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup) + .await + } + + async fn run_plan( + &self, + plan: AzureCredentialPlan, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, Error> { + match plan { AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; @@ -123,6 +154,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Sourced::new(credential, source)) .map(Some) } @@ -131,6 +163,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Some(Sourced::new(credential, source))) } AzureCredentialPlan::Chain(requests) => { @@ -138,8 +171,16 @@ impl AzureAuthService { for request in requests { let source = request.credential_source(); match self.native.acquire(request).await { - Ok(credential) => return Ok(Some(Sourced::new(credential, source))), - Err(error) => failures.push(error), + CredentialAttempt::Resolved(credential) => { + return Ok(Some(Sourced::new(credential, source))); + } + CredentialAttempt::Unavailable(error) => failures.push(error), + CredentialAttempt::Failed(error) => { + return Err(Error::CredentialChain { + last: Box::new(error), + earlier: failures, + }); + } } } match failures.pop() { @@ -155,39 +196,161 @@ impl AzureAuthService { } } +/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings +/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads. +#[derive(Clone, Debug)] +pub enum AzureIdentity { + ClientSecret { + tenant_id: String, + client_id: String, + client_secret: SecretValue, + }, + ManagedIdentity { + client_id: String, + }, + Default, +} + +fn identity_plan( + identity: &AzureIdentity, + scope: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment); + let environment = |name: &str| { + env_lookup(name) + .filter(|value| !value.trim().is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }; + let scope = deployment(scope); + let authority = environment(AzureSecretName::AuthorityHost.into()); + match identity { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant_id), + client_id: deployment(client_id), + client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment), + scope, + authority, + }, + )?)), + AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: Some(deployment(client_id)), + scope, + selection_source: InputSource::Deployment, + })?, + )), + AzureIdentity::Default => { + let tenant_id = environment(AzureSecretName::TenantId.into()); + let client_id = environment(AzureSecretName::ClientId.into()); + let from_environment = match ( + tenant_id.clone(), + client_id.clone(), + environment(AzureSecretName::ClientSecret.into()), + ) { + (Some(tenant_id), Some(client_id), Some(client_secret)) => { + Some(NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret: client_secret.map(SecretValue::new), + scope: scope.clone(), + authority: authority.clone(), + }) + } + _ => None, + }; + let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + let AzureCredentialPlan::Chain(chain) = select_native_plan( + Some(Sourced::new( + AzureCredentialType::DefaultAzureCredential, + InputSource::Deployment, + )), + tenant_id, + client_id, + environment(AzureSecretName::FederatedTokenFile.into()), + scope, + authority, + InputSource::Environment, + )? + else { + unreachable!("DefaultAzureCredential always plans a chain") + }; + let environment_request = from_environment + .map(ValidatedAzureRequest::new) + .transpose()?; + Ok(AzureCredentialPlan::Chain( + environment_request + .into_iter() + .chain(chain) + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )) + } + } +} + pub(crate) fn select_auth_plan( inputs: &AzureAuthInputs, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); - let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); - let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); - let client_secret = - configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup); - let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup) - .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); - let authority = configured_string( - &inputs.azure_authority_host, - AZURE_AUTHORITY_HOST_ENV, + let token = configured_secret( + &inputs.azure_ad_token, + AzureSecretName::AdToken.into(), env_lookup, ); - let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup) - .map(|value| { - value - .value() - .parse::() - .map(|selector| Sourced::new(selector, value.source())) - }) - .transpose() - .map_err(|error| { - Error::InvalidConfiguration(ErrorDetail::failed( - "Azure credential selector parsing", - error, - )) - })?; + let tenant_id = configured_string( + &inputs.tenant_id, + AzureSecretName::TenantId.into(), + env_lookup, + ); + let client_id = configured_string( + &inputs.client_id, + AzureSecretName::ClientId.into(), + env_lookup, + ); + let client_secret = configured_secret( + &inputs.client_secret, + AzureSecretName::ClientSecret.into(), + env_lookup, + ); + let scope = configured_string( + &inputs.azure_scope, + AzureSecretName::Scope.into(), + env_lookup, + ) + .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); + let authority = configured_string( + &inputs.azure_authority_host, + AzureSecretName::AuthorityHost.into(), + env_lookup, + ); + let selector = configured_string( + &inputs.azure_credential, + AzureSecretName::Credential.into(), + env_lookup, + ) + .map(|value| { + value + .value() + .parse::() + .map(|selector| Sourced::new(selector, value.source())) + }) + .transpose() + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "Azure credential selector parsing", + error, + )) + })?; let federated_token_file = configured_string( &inputs.federated_token_file, - AZURE_FEDERATED_TOKEN_FILE_ENV, + AzureSecretName::FederatedTokenFile.into(), env_lookup, ); @@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan( return Ok(AzureCredentialPlan::Missing); } - select_native_plan( + let probe_imds = selector + .as_ref() + .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential) + && env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + match select_native_plan( selector, tenant_id, client_id, @@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan( scope, authority, inputs.enable_azure_ad_token_refresh.source(), - ) + )? { + AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain( + requests + .into_iter() + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )), + plan => Ok(plan), + } } fn select_native_plan( @@ -509,7 +685,10 @@ mod tests { #[derive(Debug)] struct FileResolver; + use crate::native::CredentialAttempt; + struct ChainAcquirer { + terminal_on: Option<&'static str>, requests: Mutex>, succeed_on: Option<&'static str>, } @@ -518,20 +697,25 @@ mod tests { fn acquire( &self, request: ValidatedAzureRequest, - ) -> std::pin::Pin> + Send + '_>> - { + ) -> std::pin::Pin + Send + '_>> { let kind = request.kind(); self.requests.lock().unwrap().push(kind); Box::pin(async move { if self.succeed_on == Some(kind) { - Ok(ResolvedCredential::AccessToken { + CredentialAttempt::Resolved(ResolvedCredential::AccessToken { token: SecretValue::new("chain-token"), expires_on: None, }) } else { - Err(Error::CredentialAcquisition( - litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, - )) + let error = + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved { + subject: kind, + }); + if self.terminal_on == Some(kind) { + CredentialAttempt::Failed(error) + } else { + CredentialAttempt::Unavailable(error) + } } }) } @@ -690,9 +874,210 @@ mod tests { assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); } + #[derive(Debug)] + struct UnavailableImdsClient; + + impl azure_core::http::HttpClient for UnavailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + _request: &'life1 azure_core::http::Request, + ) -> std::pin::Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(std::future::pending()) + } + } + + struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer); + + impl AzureTokenAcquirer for OfflineChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin + Send + '_>> { + Box::pin(async move { + if request.kind() == "developer-tools" { + return CredentialAttempt::Resolved(ResolvedCredential::AccessToken { + token: SecretValue::new("cli-token"), + expires_on: None, + }); + } + self.0.acquire(request).await + }) + } + } + + #[rstest::rstest] + #[case::default_off_azure("DefaultAzureCredential", None, true)] + #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)] + #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)] + #[tokio::test(start_paused = true)] + async fn only_default_imds_discovery_is_bounded( + #[case] selector: &str, + #[case] identity_endpoint: Option<&str>, + #[case] expect_cli: bool, + ) { + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let inputs = AzureAuthInputs { + azure_credential: crate::types::ConfigValue::Value(Sourced::new( + selector.into(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let environment = |name: &str| { + (name == "IDENTITY_ENDPOINT") + .then_some(identity_endpoint) + .flatten() + .map(str::to_string) + }; + let outcome = tokio::time::timeout( + std::time::Duration::from_secs(2), + service.get_azure_ad_token(&inputs, &environment), + ) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().value().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::app_service(Some("IDENTITY_ENDPOINT"), false)] + #[case::legacy_app_service(Some("MSI_ENDPOINT"), false)] + #[case::off_azure(None, true)] + #[tokio::test(start_paused = true)] + async fn default_chain_uses_endpoints_from_the_secret_snapshot( + #[case] endpoint: Option<&str>, + #[case] expect_cli: bool, + #[values(false, true)] consumer_identity: bool, + ) { + use litellm_secrets::source::SecretSource; + + struct EndpointSecrets<'a>(Option<&'a str>); + + impl SecretSource for EndpointSecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> std::pin::Pin< + Box< + dyn Future, litellm_secrets::Error>> + + Send + + 'a, + >, + > { + Box::pin(async move { + Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test"))) + }) + } + } + + let snapshot = EndpointSecrets(endpoint) + .resolve(SECRET_NAMES) + .await + .unwrap(); + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let lookup = |name: &str| snapshot.get(name); + let scope = "https://service.test/.default"; + let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async { + if consumer_identity { + service + .identity_token(&super::AzureIdentity::Default, scope, &lookup) + .await + } else { + service + .get_azure_ad_token( + &AzureAuthInputs::default_credential_for_scope(scope), + &lookup, + ) + .await + .map(|token| token.map(Sourced::into_value)) + } + }) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::litellm_provider_settings_are_ignored( + &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")], + &["managed-identity", "developer-tools"] + )] + #[case::environment_client_secret_comes_first( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")], + &["client-secret", "managed-identity", "developer-tools"] + )] + #[case::aks_workload_identity( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")], + &["workload-identity", "managed-identity", "developer-tools"] + )] + #[tokio::test] + async fn the_default_identity_follows_azure_identity_defaults( + #[case] environment: &[(&str, &str)], + #[case] expected: &[&str], + ) { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let credential = service + .identity_token( + &super::AzureIdentity::Default, + "https://redis.azure.com/.default", + &lookup, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.secret().expose(), "chain-token"); + assert_eq!(*acquirer.requests.lock().unwrap(), expected); + } + + #[rstest::rstest] #[tokio::test] async fn default_chain_uses_declared_order_and_stops_after_success() { let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, requests: Mutex::new(Vec::new()), succeed_on: Some("developer-tools"), }); @@ -715,9 +1100,13 @@ mod tests { ); } + #[rstest::rstest] + #[case::all_unavailable(false)] + #[case::terminal_failure_after_unavailable(true)] #[tokio::test] - async fn chain_reports_each_acquisition_failure() { + async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) { let acquirer = Arc::new(ChainAcquirer { + terminal_on: terminal.then_some("developer-tools"), requests: Mutex::new(Vec::new()), succeed_on: None, }); @@ -744,6 +1133,27 @@ mod tests { ); } + #[rstest::rstest] + #[tokio::test] + async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: Some("managed-identity"), + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]); + assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty())); + } + #[derive(Debug)] struct CallerToken(&'static str); diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index ad60f0db817..90aab8d281c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -google-sdk = ["dep:google-cloud-auth", "dep:http"] +google-sdk = ["dep:http"] [dependencies] litellm-secret-value.workspace = true @@ -21,10 +21,14 @@ tokio.workspace = true veil.workspace = true gcp_auth = "0.12.7" -google-cloud-auth = { workspace = true, optional = true } +google-cloud-auth.workspace = true http = { workspace = true, optional = true } [dev-dependencies] +time = { workspace = true, features = ["formatting"] } +axum = { workspace = true, features = ["json"] } litellm-error.workspace = true rstest.workspace = true +rustls.workspace = true tempfile.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 650d5831796..5976ef82344 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -1,20 +1,11 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::sync::Arc; -use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential}; -use moka::future::Cache; -use serde_json::Value; -use sha2::{Digest, Sha256}; +use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential}; use crate::{ - config::{ - CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, - is_base64_encoded_json, names_a_credentials_file, non_empty_env, - }, - constants::{ - CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV, - VERTEXAI_API_KEY_ENV, - }, + config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env}, + constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV}, + identity::{GcpAuth, GcpProviderLoader, GcpTokenSource}, }; pub struct VertexEnvironment { @@ -22,28 +13,15 @@ pub struct VertexEnvironment { pub project_id: String, } -struct VertexAccessToken { - token: String, - project_id: String, -} - -#[derive(Clone)] +#[derive(Clone, Default)] pub struct VertexAuth { - providers: Cache>, - loader: Arc, -} - -impl Default for VertexAuth { - fn default() -> Self { - Self::new(Arc::new(GcpProviderLoader)) - } + identity: GcpAuth, } impl VertexAuth { - pub fn new(loader: Arc) -> Self { + pub fn new(loader: Arc) -> Self { Self { - providers: Cache::builder().max_capacity(64).build(), - loader, + identity: GcpAuth::new(loader), } } @@ -52,7 +30,10 @@ impl VertexAuth { config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { - self.load_provider(config, env_lookup).await?.token().await + self.identity + .access_token(credential_source(config, env_lookup)) + .await + .map(|token| token.expose().to_owned()) } pub async fn validate_environment( @@ -77,10 +58,18 @@ impl VertexAuth { (true, _) => None, (false, Some(token)) => Some(token), (false, None) => { - let access = self.get_access_token(config, env_lookup).await?; + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = match project_id { + Some(project_id) => (provider.token().await?, project_id), + None => tokio::try_join!(provider.token(), provider.project_id())?, + }; return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), + headers: apply_credential( + headers, + token.expose(), + CredentialPlacement::Bearer, + )?, + project_id, }); } }; @@ -104,248 +93,171 @@ impl VertexAuth { }) } - async fn get_access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let provider = self.load_provider(config, env_lookup).await?; - let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; - Ok(VertexAccessToken { token, project_id }) - } - async fn load_provider( &self, config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, Error> { - let source = credential_source(config, env_lookup); - let key = source.cache_key(); - self.providers - .try_get_with(key, self.loader.load(source)) + ) -> Result, Error> { + self.identity + .load_provider(credential_source(config, env_lookup)) .await - .map_err(|error| (*error).clone()) } } -pub trait VertexTokenSource: Send + Sync { - fn project_id(&self) -> VertexAuthFuture<'_, String>; - fn token(&self) -> VertexAuthFuture<'_, String>; -} - -pub trait VertexProviderLoader: Send + Sync { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; -} - -pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; - -struct GcpTokenSource(Arc); - -impl VertexTokenSource for GcpTokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .project_id() - .await - .map(|project| project.to_string()) - .map_err(auth_acquisition_error) - }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .token(&[CLOUD_PLATFORM_SCOPE]) - .await - .map(|token| token.as_str().to_string()) - .map_err(auth_acquisition_error) - }) - } -} - -struct GcpProviderLoader; - -impl VertexProviderLoader for GcpProviderLoader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - Box::pin(async move { - let provider: Arc = match source { - CredentialSource::Inline(configured) => Arc::new( - CustomServiceAccount::from_json(validate_request_credentials( - configured.expose(), - )?) - .map_err(auth_acquisition_error)?, - ), - CredentialSource::Trusted(configured) => { - Arc::new(load_trusted_service_account(configured.expose())?) - } - CredentialSource::ApplicationCredentials(path) => { - Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) - } - CredentialSource::Adc => { - gcp_auth::provider().await.map_err(auth_acquisition_error)? - } - }; - Ok(Arc::new(GcpTokenSource(provider)) as Arc) - }) - } -} - -fn validate_request_credentials(configured: &str) -> Result<&str, Error> { - let token_uri = serde_json::from_str::(configured) - .ok() - .and_then(|credentials| { - credentials - .get("token_uri") - .and_then(Value::as_str) - .map(str::to_string) - }); - if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { - return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { - field: "request-controlled vertex_credentials token_uri".into(), - expected: "the canonical Google OAuth token endpoint", - })); - } - Ok(configured) -} - -impl CredentialSource { - fn cache_key(&self) -> CredentialCacheKey { - match self { - Self::Inline(configured) => { - CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) - } - Self::Trusted(configured) => { - CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) - } - Self::ApplicationCredentials(path) => { - CredentialCacheKey::ApplicationCredentials(path.clone()) - } - Self::Adc => CredentialCacheKey::Adc, - } - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -enum CredentialCacheKey { - Inline([u8; 32]), - Trusted([u8; 32]), - ApplicationCredentials(String), - Adc, -} - -fn load_trusted_service_account(configured: &str) -> Result { - if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration( - ErrorDetail::Base64EncodedJson { - field: "vertex_credentials", - }, - )); - } - if names_a_credentials_file(configured) { - CustomServiceAccount::from_file(configured) - } else { - CustomServiceAccount::from_json(configured) - } - .map_err(auth_acquisition_error) -} - -fn auth_acquisition_error(error: gcp_auth::Error) -> Error { - Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( - "Vertex AI credentials", - error, - )) -} - #[cfg(test)] mod tests { + use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use crate::{ + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, + VertexConfig, get_vertex_ai_location, secret_names, + }; use litellm_secret_value::SecretValue; + use rstest::rstest; + use serde_json::{Value, json}; - use super::*; - - #[rstest::rstest] - #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] - #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] - #[case::missing_endpoint("{}", false)] - fn request_credentials_require_canonical_token_endpoint( - #[case] credentials: &str, - #[case] accepted: bool, - ) { - assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + struct FakeProvider { + project: Option<&'static str>, + calls: Arc, } - #[rstest::rstest] - #[case::missing_file( - "/nonexistent/vertex-key.json", - "failed to open application credentials file" - )] - #[case::malformed_json( - " {\"type\": \"service_account\",", - "failed to deserialize ApplicationCredentials" - )] - #[case::quoted_json( - "'{\"type\": \"service_account\"}'", - "failed to deserialize ApplicationCredentials" - )] - #[tokio::test] - async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( - #[case] configured: &str, - #[case] expected: &str, - ) { - let Err(error) = GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - else { - panic!("{configured:?} should not load"); - }; - let report = litellm_error::Report::new(&error).to_string(); - assert!(report.contains(expected), "{report}"); - } + impl GcpTokenSource for FakeProvider { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + self.project.map(str::to_owned).ok_or( + litellm_auth_types::Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { + subject: "credential project", + }, + ), + ) + }) + } - #[rstest::rstest] - #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] - #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] - #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] - #[tokio::test] - async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { - assert!( - trusted_load_error(configured) - .await - .contains("is base64-encoded JSON, not the JSON itself") - ); - } - - async fn trusted_load_error(configured: &str) -> String { - match GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - { - Ok(_) => panic!("{configured:?} should not load"), - Err(error) => litellm_error::Report::new(&error).to_string(), + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(SecretValue::new("adc-token")) }) } } - #[rstest::rstest] - #[tokio::test] - async fn a_relative_key_file_named_like_json_is_read_as_a_file() { - let key = tempfile::Builder::new() - .prefix("{vertex") - .suffix(".json") - .tempfile_in(".") - .unwrap(); - std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); - let relative_name = key.path().file_name().unwrap().to_str().unwrap(); - let absolute_path = std::fs::canonicalize(key.path()).unwrap(); - assert_eq!( - trusted_load_error(relative_name).await, - trusted_load_error(absolute_path.to_str().unwrap()).await - ); + struct FakeLoader { + loads: Arc, + provider: Arc, } - #[rstest::rstest] - fn inline_and_trusted_credentials_never_share_a_cache_entry() { - assert_ne!( - CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), - CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() - ); + impl GcpProviderLoader for FakeLoader { + fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { + calls, + project: Some("adc-project"), + }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[rstest] + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + let allowed = secret_names() + .iter() + .map(|name| (*name).to_owned()) + .collect::>(); + assert!(seen.lock().unwrap().is_subset(&allowed)); + } + + #[rstest] + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[rstest] + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + let config = VertexConfig::default(); + let environment = |_: &str| None; + let (first, second) = tokio::try_join!( + auth.validate_environment(Vec::new(), None, &config, &environment), + auth.validate_environment(Vec::new(), None, &config, &environment), + ) + .unwrap(); + assert_eq!(first.project_id, "adc-project"); + assert_eq!(first.headers[0].1, "Bearer adc-token"); + assert_eq!(first.project_id, second.project_id); + assert_eq!(first.headers, second.headers); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } + + #[rstest] + #[tokio::test] + async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() { + let auth = VertexAuth::new(Arc::new(FakeLoader { + loads: Arc::new(AtomicUsize::new(0)), + provider: Arc::new(FakeProvider { + calls: Arc::new(AtomicUsize::new(0)), + project: None, + }), + })); + let configured = VertexConfig::new(None, Some("configured-project".into()), None); + let environment = auth + .validate_environment(Vec::new(), None, &configured, &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "configured-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); } } diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 133e1c26903..048e40b0a2f 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -146,6 +146,20 @@ pub enum CredentialSource { Adc, } +impl CredentialSource { + pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option) -> Self { + VertexParams::CREDENTIALS + .resolve(&|_| None, env_lookup) + .map(|configured| Self::Trusted(SecretValue::new(configured))) + .unwrap_or_else(|| Self::application_default(env_lookup)) + } + + pub fn application_default(env_lookup: &dyn Fn(&str) -> Option) -> Self { + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map_or(Self::Adc, Self::ApplicationCredentials) + } +} + pub(crate) fn credential_source( config: &VertexConfig, env_lookup: &dyn Fn(&str) -> Option, @@ -158,12 +172,7 @@ pub(crate) fn credential_source( } }; } - if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) { - return CredentialSource::Trusted(SecretValue::new(configured)); - } - non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) - .map(CredentialSource::ApplicationCredentials) - .unwrap_or(CredentialSource::Adc) + CredentialSource::from_vertex_environment(env_lookup) } fn optional_credentials( @@ -319,6 +328,24 @@ mod tests { )); } + #[rstest::rstest] + fn application_default_ignores_vertex_credentials() { + let env = |name: &str| match name { + VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()), + GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()), + _ => None, + }; + assert!( + matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + CredentialSource::application_default( + &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into()) + ), + CredentialSource::Adc + )); + } + #[rstest::rstest] fn params_become_a_deployment_sourced_config() { let params = VertexParams { diff --git a/litellm-rust/crates/auth-gcp/src/identity.rs b/litellm-rust/crates/auth-gcp/src/identity.rs new file mode 100644 index 00000000000..5144221b218 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/identity.rs @@ -0,0 +1,445 @@ +use crate::{ + config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file}, + constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT}, +}; +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth_types::{Error, ErrorDetail}; +use litellm_secret_value::SecretValue; +use moka::future::Cache; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::{future::Future, pin::Pin, sync::Arc}; + +#[derive(Clone)] +pub struct GcpAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for GcpAuth { + fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl GcpAuth { + pub fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + pub async fn access_token(&self, source: CredentialSource) -> Result { + self.load_provider(source).await?.token().await + } + + pub(crate) async fn load_provider( + &self, + source: CredentialSource, + ) -> Result, Error> { + self.providers + .try_get_with(source.cache_key(), self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +pub trait GcpTokenSource: Send + Sync { + fn project_id(&self) -> GcpAuthFuture<'_, String>; + fn token(&self) -> GcpAuthFuture<'_, SecretValue>; +} + +pub trait GcpProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc>; +} + +pub type GcpAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct ServiceAccountTokenSource(Arc); + +impl GcpTokenSource for ServiceAccountTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| SecretValue::new(token.as_str())) + .map_err(auth_acquisition_error) + }) + } +} + +struct SdkTokenSource { + credentials: google_cloud_auth::credentials::AccessTokenCredentials, + project_id: Option, +} + +impl GcpTokenSource for SdkTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.project_id + .clone() + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project", + })) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.credentials + .access_token() + .await + .map(|token| SecretValue::new(token.token)) + .map_err(auth_acquisition_error) + }) + } +} + +async fn load_application_credentials(path: &str) -> Result, Error> { + let contents = tokio::fs::read_to_string(path) + .await + .map_err(auth_acquisition_error)?; + load_credentials_json(&contents) +} + +fn load_credentials_json(contents: &str) -> Result, Error> { + let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?; + let kind = configured.get("type").and_then(Value::as_str); + let project_id = configured + .get("project_id") + .and_then(Value::as_str) + .map(str::to_string); + let credentials = match kind { + Some("authorized_user") => { + google_cloud_auth::credentials::user_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("impersonated_service_account") => { + google_cloud_auth::credentials::impersonated::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("external_account") => { + google_cloud_auth::credentials::external_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("service_account") | None => { + let provider = + CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?; + return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider)))); + } + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Google application credentials type".into(), + expected: "a supported Google credential type", + })); + } + } + .map_err(auth_acquisition_error)?; + Ok(Arc::new(SdkTokenSource { + credentials, + project_id, + })) +} + +struct NativeProviderLoader; + +impl GcpProviderLoader for NativeProviderLoader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + return load_trusted_credentials(configured.expose()).await; + } + CredentialSource::ApplicationCredentials(path) => { + return load_application_credentials(&path).await; + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, Error> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "request-controlled vertex_credentials token_uri".into(), + expected: "the canonical Google OAuth token endpoint", + })); + } + Ok(configured) +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +async fn load_trusted_credentials(configured: &str) -> Result, Error> { + if is_base64_encoded_json(configured) { + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); + } + if names_a_credentials_file(configured) { + load_application_credentials(configured).await + } else { + load_credentials_json(configured) + } +} + +fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( + "Google credentials", + error, + )) +} + +#[cfg(test)] +mod tests { + use litellm_secret_value::SecretValue; + + use super::*; + + #[rstest::rstest] + #[case::authorized_user("authorized_user", 2)] + #[case::external_account("external_account", 2)] + #[case::impersonated_service_account("impersonated_service_account", 4)] + #[tokio::test] + async fn explicit_application_credentials_refresh_expired_tokens( + #[case] kind: &str, + #[case] expected_requests: usize, + #[values(false, true)] inline: bool, + ) { + use axum::{Json, Router, routing::post}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let _ = rustls::crypto::ring::default_provider().install_default(); + let requests = Arc::new(AtomicUsize::new(0)); + let recorded = requests.clone(); + let impersonations = requests.clone(); + let app = Router::new().route("/token", post(move || { + let recorded = recorded.clone(); + async move { + Json(serde_json::json!({ + "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expires_in": 3600, + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + })) + } + })); + let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || { + let recorded = impersonations.clone(); + async move { + let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap(); + Json(serde_json::json!({ + "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expireTime": expiration, + })) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let file = tempfile::NamedTempFile::new().unwrap(); + let subject_token = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(subject_token.path(), "subject-token").unwrap(); + let configured = if kind == "external_account" { + serde_json::json!({ + "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"), + "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project" + }) + } else { + let source = serde_json::json!({ + "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret", + "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token") + }); + if kind == "impersonated_service_account" { + serde_json::json!({ + "type": kind, "source_credentials": source, "quota_project_id": "quota-project", + "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken") + }) + } else { + source + } + }; + std::fs::write(file.path(), configured.to_string()).unwrap(); + let source = if inline { + CredentialSource::Trusted(SecretValue::new(configured.to_string())) + } else { + CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into()) + }; + let provider = NativeProviderLoader.load(source).await.unwrap(); + + assert!(matches!( + provider.project_id().await, + Err(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project" + })) + )); + let original = provider.token().await.unwrap(); + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(3601)).await; + tokio::time::resume(); + let refreshed = provider.token().await.unwrap(); + assert_ne!(original, refreshed); + assert_eq!(requests.load(Ordering::SeqCst), expected_requests); + server.abort(); + } + + #[rstest::rstest] + #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] + #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] + #[case::missing_endpoint("{}", false)] + fn request_credentials_require_canonical_token_endpoint( + #[case] credentials: &str, + #[case] accepted: bool, + ) { + assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + } + + #[rstest::rstest] + #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")] + #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")] + #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")] + #[tokio::test] + async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( + #[case] configured: &str, + #[case] expected: &str, + ) { + let Err(error) = NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + else { + panic!("{configured:?} should not load"); + }; + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + } + + #[rstest::rstest] + #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] + #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] + #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] + #[tokio::test] + async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { + assert!( + trusted_load_error(configured) + .await + .contains("is base64-encoded JSON, not the JSON itself") + ); + } + + async fn trusted_load_error(configured: &str) -> String { + match NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + { + Ok(_) => panic!("{configured:?} should not load"), + Err(error) => litellm_error::Report::new(&error).to_string(), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_relative_key_file_named_like_json_is_read_as_a_file() { + let key = tempfile::Builder::new() + .prefix("{vertex") + .suffix(".json") + .tempfile_in(".") + .unwrap(); + std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); + let relative_name = key.path().file_name().unwrap().to_str().unwrap(); + let absolute_path = std::fs::canonicalize(key.path()).unwrap(); + assert_eq!( + trusted_load_error(relative_name).await, + trusted_load_error(absolute_path.to_str().unwrap()).await + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#"{"type":"unsupported","private_key":"credential-value"}"#, + ) + .unwrap(); + let result = GcpAuth::default() + .access_token(CredentialSource::ApplicationCredentials( + file.path().to_str().unwrap().into(), + )) + .await; + let error = result.unwrap_err(); + assert!(matches!(error, Error::InvalidConfiguration(_))); + assert!( + !litellm_error::Report::new(&error) + .to_string() + .contains("credential-value") + ); + } + + #[rstest::rstest] + fn inline_and_trusted_credentials_never_share_a_cache_entry() { + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 2b389d45926..03c660e6eb8 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,16 +1,16 @@ mod auth; mod config; pub mod constants; +mod identity; #[cfg(feature = "google-sdk")] mod sdk; -pub use auth::{ - VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource, -}; +pub use auth::{VertexAuth, VertexEnvironment}; pub use config::{ CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, }; pub use constants::secret_names; +pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource}; #[cfg(feature = "google-sdk")] pub use sdk::GoogleCredentials; diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs deleted file mode 100644 index 5991749ab8a..00000000000 --- a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs +++ /dev/null @@ -1,120 +0,0 @@ -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; - -use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader, - VertexTokenSource, get_vertex_ai_location, secret_names, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -struct FakeProvider { - calls: Arc, -} - -impl VertexTokenSource for FakeProvider { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-project".into()) }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-token".into()) }) - } -} - -struct FakeLoader { - loads: Arc, - provider: Arc, -} - -impl VertexProviderLoader for FakeLoader { - fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - let loads = self.loads.clone(); - let provider = self.provider.clone(); - Box::pin(async move { - loads.fetch_add(1, Ordering::SeqCst); - Ok(provider) - }) - } -} - -fn config(value: Value) -> VertexConfig { - VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) - .unwrap() -} - -fn auth(calls: Arc, loads: Arc) -> VertexAuth { - let provider: Arc = Arc::new(FakeProvider { calls }); - VertexAuth::new(Arc::new(FakeLoader { loads, provider })) -} - -#[rstest] -#[tokio::test] -async fn secret_names_cover_environment_reads() { - let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); - let recorded = seen.clone(); - let env = |name: &str| { - recorded.lock().unwrap().insert(name.to_string()); - None - }; - let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); - auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) - .await - .unwrap(); - get_vertex_ai_location(&VertexConfig::default(), &env); - assert!( - seen.lock() - .unwrap() - .iter() - .all(|name| secret_names().contains(&name.as_str())) - ); -} - -#[rstest] -#[tokio::test] -async fn explicit_token_and_header_do_not_acquire_adc() { - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); - let configured = config(json!({"vertex_project":"project-1"})); - let explicit = auth - .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) - .await - .unwrap(); - assert_eq!(explicit.headers[0].1, "Bearer access-token"); - let existing = auth - .validate_environment( - vec![("authorization".into(), "Bearer existing".into())], - None, - &configured, - &|_| None, - ) - .await - .unwrap(); - assert_eq!(existing.headers[0].1, "Bearer existing"); - assert_eq!(loads.load(Ordering::SeqCst), 0); -} - -#[rstest] -#[tokio::test] -async fn provider_is_reused_across_authentication_calls() { - let calls = Arc::new(AtomicUsize::new(0)); - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(calls.clone(), loads.clone()); - for _ in 0..2 { - let environment = auth - .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) - .await - .unwrap(); - assert_eq!(environment.project_id, "adc-project"); - assert_eq!(environment.headers[0].1, "Bearer adc-token"); - } - assert_eq!(loads.load(Ordering::SeqCst), 1); - assert_eq!(calls.load(Ordering::SeqCst), 4); -} diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index a8ba7b64e03..a2e5629bf3f 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -92,6 +92,11 @@ impl Error { #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { + #[error("{operation} returned HTTP {status}")] + UnexpectedStatus { + operation: &'static str, + status: u16, + }, #[error("{field} must be {expected}")] InvalidType { field: String, diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index 8075fc15e9f..de21fbf94c9 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -53,6 +53,9 @@ pub struct AwsParams { } impl AwsParams { + pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN"; + pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; + pub const ACCESS_KEY_ID: ParamSpec = spec( "access key id", ParamKind::Credential, @@ -149,6 +152,10 @@ impl AwsParams { AwsParams::SPECS .iter() .flat_map(|spec| spec.env.iter().copied()) + .chain([ + AwsParams::ROLE_ARN_ENV, + AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV, + ]) .fold(Vec::new(), |mut names, name| { if !names.contains(&name) { names.push(name); diff --git a/litellm-rust/crates/auth-types/tests/aws_params.rs b/litellm-rust/crates/auth-types/tests/aws_params.rs index ed9a4f78ec2..b438984c199 100644 --- a/litellm-rust/crates/auth-types/tests/aws_params.rs +++ b/litellm-rust/crates/auth-types/tests/aws_params.rs @@ -3,21 +3,13 @@ use rstest::rstest; use serde_json::{Map, Value}; #[rstest] -fn fields_name_every_param_once_and_in_declaration_order() { +fn every_named_param_roundtrips_without_loss() { let filled: Map = AwsParams::fields() .map(|name| (name.to_string(), Value::from(format!("value-of-{name}")))) .collect(); let typed = AwsParams::from_optional_params(&filled); let serialized = serde_json::to_value(&typed).unwrap(); - assert_eq!( - serialized - .as_object() - .unwrap() - .keys() - .map(String::as_str) - .collect::>(), - AwsParams::fields().collect::>() - ); + assert_eq!(filled.len(), AwsParams::fields().count()); assert_eq!(serialized, Value::Object(filled.clone())); assert_eq!( serde_json::from_value::(Value::Object(filled)).unwrap(), diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index a6b08445bc8..b8c6720c234 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -8,7 +8,7 @@ use litellm_cache::{ use litellm_http::Client; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; -use crate::{GcpTokenSource, TokenSource}; +use crate::{GcsTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; @@ -62,7 +62,7 @@ pub struct GcsCache { impl GcsCache { pub fn new(config: GcsConfig, client: Client, codec: S) -> Self { - let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone())); Self::with_token_source(config, client, codec, token) } diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs index cbb61cf0685..d40bca081a4 100644 --- a/litellm-rust/crates/cache-gcs/src/lib.rs +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -2,4 +2,4 @@ mod cache; mod token; pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; -pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; +pub use token::{GcsTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 852d8bbf035..4f639e7b865 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,6 @@ use std::{future::Future, pin::Pin}; -use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, Sourced}; +use litellm_auth_gcp::{CredentialSource, GcpAuth}; use litellm_secret_value::SecretValue; use litellm_cache::Error; @@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static { fn bearer_token(&self) -> Pin> + Send + '_>>; } -pub struct GcpTokenSource { - auth: VertexAuth, - config: VertexConfig, +pub struct GcsTokenSource { + auth: GcpAuth, + credentials: Option, } -impl GcpTokenSource { +impl GcsTokenSource { pub fn new(path_service_account: Option) -> Self { - let credentials = path_service_account - .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); Self { - auth: VertexAuth::default(), - config: VertexConfig::new(credentials, None, None), + auth: GcpAuth::default(), + credentials: path_service_account.filter(|path| !path.trim().is_empty()), } } } -impl TokenSource for GcpTokenSource { +impl TokenSource for GcsTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { Box::pin(async move { self.auth - .access_token(&self.config, &|name| std::env::var(name).ok()) + .access_token(match &self.credentials { + Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())), + None => CredentialSource::application_default(&|name| std::env::var(name).ok()), + }) .await + .map(|token| token.expose().to_owned()) .map_err(Error::unavailable) }) } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index 659390a5659..02707c24881 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials { self.config.secret_access_key.clone(), ) { return Ok(AwsCredentials::new( - access_key_id, - secret_access_key, - self.config.session_token.clone(), + access_key_id.expose(), + secret_access_key.expose(), + self.config + .session_token + .as_ref() + .map(|value| value.expose().to_owned()), None, "litellm-s3-cache", )); diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs index 59c5aa2e091..5aa251ac749 100644 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ b/litellm-rust/crates/cache-s3/tests/auth.rs @@ -5,9 +5,9 @@ use rstest::rstest; fn explicit(session_token: Option<&str>) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), - session_token: session_token.map(str::to_string), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + session_token: session_token.map(Into::into), region_name: Some("us-east-1".to_string()), ..Default::default() } diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 7378dde21bf..311d2402e72 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig { url: endpoint.to_string(), }), auth: AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), region_name: Some("us-east-1".to_string()), ..Default::default() }, diff --git a/litellm-rust/crates/inference-ocr/tests/resources.rs b/litellm-rust/crates/inference-ocr/tests/resources.rs index d618647e89b..3132a837eaf 100644 --- a/litellm-rust/crates/inference-ocr/tests/resources.rs +++ b/litellm-rust/crates/inference-ocr/tests/resources.rs @@ -7,7 +7,7 @@ use std::sync::{ use litellm_auth::AuthServices; use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource, + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, }; use litellm_http::{HttpSettings, Resolution}; use litellm_inference::resources::CoreResources; @@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream}; struct TokenSource(String); -impl VertexTokenSource for TokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { +impl GcpTokenSource for TokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { Box::pin(async { Ok(self.0.clone()) }) } - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async { Ok(self.0.clone()) }) + fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> { + Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) }) } } #[derive(Default)] struct Loader(AtomicUsize); -impl VertexProviderLoader for Loader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { +impl GcpProviderLoader for Loader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { Box::pin(async move { self.0.fetch_add(1, Ordering::SeqCst); let identity = match source { CredentialSource::Trusted(secret) => secret.expose().to_string(), other => panic!("unexpected credential source: {other:?}"), }; - Ok(Arc::new(TokenSource(identity)) as Arc) + Ok(Arc::new(TokenSource(identity)) as Arc) }) } } diff --git a/litellm-rust/crates/secret-value/src/secret_value.rs b/litellm-rust/crates/secret-value/src/secret_value.rs index 27e4227080d..33c559b80ec 100644 --- a/litellm-rust/crates/secret-value/src/secret_value.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -16,6 +16,18 @@ impl SecretValue { } } +impl From for SecretValue { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl From<&str> for SecretValue { + fn from(value: &str) -> Self { + Self::new(value) + } +} + impl PartialEq for SecretValue { fn eq(&self, other: &Self) -> bool { subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into() diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 66f4f51bda9..6bec8ad96a1 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -39,30 +39,15 @@ impl Credentials { Self { auth, config: AwsAuthConfig { - access_key_id: context - .access_key_id - .as_ref() - .map(|value| value.expose().to_owned()), - secret_access_key: context - .secret_access_key - .as_ref() - .map(|value| value.expose().to_owned()), - session_token: context - .session_token - .as_ref() - .map(|value| value.expose().to_owned()), + access_key_id: context.access_key_id.as_ref().cloned(), + secret_access_key: context.secret_access_key.as_ref().cloned(), + session_token: context.session_token.as_ref().cloned(), region_name: region(settings, environment.as_ref()).ok(), role_name: settings.aws_role_name.clone(), session_name: settings.aws_session_name.clone(), - external_id: settings - .aws_external_id - .as_ref() - .map(|v| v.expose().to_owned()), + external_id: settings.aws_external_id.as_ref().cloned(), profile_name: settings.aws_profile_name.clone(), - web_identity_token: settings - .aws_web_identity_token - .as_ref() - .map(|v| v.expose().to_owned()), + web_identity_token: settings.aws_web_identity_token.as_ref().cloned(), sts_endpoint: settings.aws_sts_endpoint.clone(), }, environment, From 8fcb016dc6c887c7ae545e86598f69dc63e20fc9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 14:38:44 -0800 Subject: [PATCH 11/19] test: make three flaky unit tests deterministic (#50114) * test: make four flaky unit tests deterministic The skill archive fixtures stamped the current time into zip bytes used as parametrize ids, so xdist workers collected different ids and errored the whole shard. They now use a fixed timestamp. The usage-based v2 override test pinned no clock, so a minute boundary between rpm checks reset the counter. It now pins get_utc_datetime the way #49505 does. The stalled-burst DB lookup test asserted a 2s wall-clock bound. It now asserts every lookup was waiting at once, which fails if lookups are serialized. The long tool_search placement test asserted a 1s wall-clock bound. It keeps the placement assertion and drops the bound, as #49505 did for the json accumulator * test(skills): drop the archive fixture timestamp change now that #50118 removed its test --- ...rompt_templates_mid_conversation_system.py | 5 +---- tests/unit/proxy/db/test_db_lookup_gate.py | 22 ++++++++++++++++--- .../test_router_routing_groups.py | 6 ++++- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index aab57b87a90..d7696c00dee 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -327,7 +327,7 @@ EMPTY_ASSISTANT = pytest.mark.parametrize( ) -def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time(): +def test_flagged_placement_turns_reminders_before_a_long_run_of_unanswered_tool_search_uses_into_user_turns(): search_uses: Final = tuple( { "role": "assistant", @@ -349,12 +349,9 @@ def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_ {"role": "user", "content": "end"}, ] - started: Final = time.perf_counter() placed: Final = place_mid_conversation_system(messages, supports_mid_conversation_system=True) - elapsed: Final = time.perf_counter() - started assert set(_roles(placed)[:21]) == {"user"} - assert elapsed < 1.0 def test_flagged_placement_merges_every_system_run_that_shares_a_user_turn(): diff --git a/tests/unit/proxy/db/test_db_lookup_gate.py b/tests/unit/proxy/db/test_db_lookup_gate.py index 68903170840..3450d3e05b8 100644 --- a/tests/unit/proxy/db/test_db_lookup_gate.py +++ b/tests/unit/proxy/db/test_db_lookup_gate.py @@ -46,21 +46,37 @@ async def test_bounded_db_lookup_returns_a_prompt_answer_without_recording_a_sta assert tracker.stalled_within(30) is False +class _PendingLookups: + def __init__(self) -> None: + self.pending = 0 + self.peak = 0 + + async def never_answers(self) -> None: + self.pending += 1 + self.peak = max(self.peak, self.pending) + try: + await asyncio.Event().wait() + finally: + self.pending -= 1 + + @pytest.mark.asyncio async def test_bounded_db_lookup_fails_a_whole_stalled_burst_within_one_deadline(): tracker: Final = DBLookupStallTracker() burst: Final = 200 - started: Final = time.monotonic() + lookups: Final = _PendingLookups() results: Final = await asyncio.gather( *( - bounded_db_lookup(_never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) + bounded_db_lookup(lookups.never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) for i in range(burst) ), return_exceptions=True, ) - assert time.monotonic() - started < 2 + assert lookups.peak == burst, ( + f"only {lookups.peak} of {burst} stalled lookups were waiting on their deadline together" + ) assert len(results) == burst assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) assert tracker.stalled_within(30) is True diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 5eeb46c3518..afcd1e2a7ff 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -1022,7 +1022,11 @@ async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_fo await _mock_completion(router, routing_strategy="usage-based-routing-v2") -def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): +def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(monkeypatch): + monkeypatch.setattr( + "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", + lambda: datetime.datetime(2026, 1, 1, 12, 0, 30, tzinfo=datetime.timezone.utc), + ) router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) messages = [{"role": "user", "content": "hi"}] From d29c515444b9b7a3d886a72340cff68a06a0afd8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 14:40:38 -0800 Subject: [PATCH 12/19] ci(unit): stop the full history clone on the db-and-spend shard, fetch the schema merge base by sha (#50120) * ci(unit): fetch the merge base by sha for the db-and-spend schema check instead of cloning full history * ci(unit): read the schema check arguments through one typed mapping * ci(unit): give the schema check its own step with the ui-unit base fallback and register it with the merge-base step test * ci(unit): keep the setup ceiling at 35 minutes, the install step never took over 33 seconds --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 28 ++++++++++++++------- scripts/check_schema_compatibility.py | 18 ++++++++----- tests/unit/test_lint_workflow_diff_gates.py | 11 ++++++++ 3 files changed, 42 insertions(+), 15 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a090ec7bdb1..94152aeb7a0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -977,10 +977,10 @@ jobs: dist: loadscope steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }} + timeout-minutes: 3 with: persist-credentials: false - fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }} + fetch-depth: 1 - name: Detect relevant changes id: changes @@ -1013,13 +1013,11 @@ jobs: with: artifact: ${{ needs.rust-bridge.outputs.artifact }} - - name: Install dependencies and check schema + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 8 + timeout-minutes: 5 env: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} - UNIT_SHARD: ${{ matrix.shard }} - SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }} run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then @@ -1031,9 +1029,21 @@ jobs: uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" fi uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - if [ "$UNIT_SHARD" = "db-and-spend" ]; then - uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF" - fi + + - name: Check schema.prisma against the merge base + if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request' + timeout-minutes: 3 + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.pull_request.number }} + EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref) + merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') + test -n "$merge_base" + git fetch --no-tags --depth=1 origin "$merge_base" + uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base" - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 898c65dff2c..b8e9ac3c07a 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -52,14 +52,20 @@ def check_breaking_changes(old_schema: Schema, new_schema: Schema) -> tuple[str, return tuple(changes()) -def main() -> int: - parser: Final = argparse.ArgumentParser() - parser.add_argument("--base-ref", required=True) - options: Final = parser.parse_args() - base_ref: Final = TypeAdapter(str).validate_python(options.base_ref) - base: Final = subprocess.run( +def _merge_base(base_ref: str) -> str: + return subprocess.run( ("git", "merge-base", "HEAD", base_ref), capture_output=True, text=True, check=True ).stdout.strip() + + +def main() -> int: + parser: Final = argparse.ArgumentParser() + source: Final = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--base-ref", help="compare against the merge base of HEAD and this ref (needs history)") + source.add_argument("--base-commit", help="compare against this commit as is (works in a shallow clone)") + options: Final = TypeAdapter(dict[str, str | None]).validate_python(vars(parser.parse_args())) + base_commit: Final = options["base_commit"] + base: Final = base_commit if base_commit else _merge_base(options["base_ref"] or "") previous: Final = subprocess.run( ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True ).stdout diff --git a/tests/unit/test_lint_workflow_diff_gates.py b/tests/unit/test_lint_workflow_diff_gates.py index 816ef6ce707..5e6ae9d285a 100644 --- a/tests/unit/test_lint_workflow_diff_gates.py +++ b/tests/unit/test_lint_workflow_diff_gates.py @@ -383,10 +383,20 @@ def _guard_base_argument_and_checkout(run: StepRun) -> tuple[str, str]: return argv[argv.index("--base") + 1], _sha(run.runner.checkout, "HEAD") +def _schema_base_commit_and_its_type_in_the_checkout(run: StepRun) -> tuple[str, str]: + argv: Final = _argv_recorded(run.runner, "uv") + base_commit: Final = argv[argv.index("--base-commit") + 1] + return base_commit, _git(run.runner.checkout, "cat-file", "-t", base_commit).strip() + + def _merge_base(upstream: StaleBaseUpstream, ref: str) -> str: return upstream.merge_base_by_ref[ref] +def _merge_base_as_a_commit(upstream: StaleBaseUpstream, ref: str) -> tuple[str, str]: + return upstream.merge_base_by_ref[ref], "commit" + + def _ui_files_changed_since_merge_base(upstream: StaleBaseUpstream, ref: str) -> frozenset[str]: merge_base: Final = upstream.merge_base_by_ref[ref] changed: Final = _git(upstream.path, "diff", "--name-only", f"--relative={UI_DASHBOARD}", merge_base, upstream.head) @@ -402,6 +412,7 @@ CONSUMERS: Final = MappingProxyType( "test-linting.yml:python": Consumer(_gate_base_exported, _merge_base), "test-linting.yml:ui": Consumer(_ui_files_queued_for_prettier, _ui_files_changed_since_merge_base), "test-unit.yml:ui-unit": Consumer(_ui_files_passed_to_vitest_related, _ui_files_changed_since_merge_base), + "test-unit.yml:unit": Consumer(_schema_base_commit_and_its_type_in_the_checkout, _merge_base_as_a_commit), "spend-logs-guard.yml:spend-logs-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip), "cost-map-guard.yml:cost-map-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip), } From 6b8788eb2759ef5f6c14ba970df4e4b0c5bfee1e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:45:14 -0800 Subject: [PATCH 13/19] fix(bedrock): backfill GovCloud service-tier prices (#50121) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 228 ++++++++++++++++-- model_prices_and_context_window.json | 228 ++++++++++++++++-- 2 files changed, 416 insertions(+), 40 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5d91941a953..25c7c3593ed 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15715,8 +15715,10 @@ "input_cost_per_image": 7.2e-05, "input_cost_per_video_per_second": 0.00084, "input_cost_per_audio_per_second": 0.000168, + "input_cost_per_token_batches": 8.1e-08, "output_cost_per_token": 0.0, "output_vector_size": 3072, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_embedding_image_input": true, "supports_image_input": true, "supports_video_input": true, @@ -15725,11 +15727,14 @@ "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.8e-08, "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -15740,11 +15745,14 @@ "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 1.05e-08, "input_cost_per_token": 4.2e-08, + "input_cost_per_token_batches": 2.1e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.68e-07, + "output_cost_per_token_batches": 8.4e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -15752,12 +15760,21 @@ }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, "input_cost_per_token": 9.6e-07, + "input_cost_per_token_batches": 4.8e-07, + "input_cost_per_token_flex": 4.8e-07, + "input_cost_per_token_priority": 1.68e-06, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.84e-06, + "output_cost_per_token_batches": 1.92e-06, + "output_cost_per_token_flex": 1.92e-06, + "output_cost_per_token_priority": 6.72e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -69117,12 +69134,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69135,12 +69158,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69151,12 +69180,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69165,12 +69200,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69182,12 +69223,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69201,12 +69248,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69412,12 +69465,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69430,12 +69489,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69446,12 +69511,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69460,12 +69531,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69477,12 +69554,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69496,12 +69579,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69709,12 +69798,28 @@ "input_cost_per_token_above_272k_tokens": 5.28e-06, "cache_creation_input_token_cost": 3.3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-05, + "cache_creation_input_token_cost_flex": 1.65e-06, + "cache_creation_input_token_cost_priority": 6.6e-06, "cache_read_input_token_cost": 2.64e-07, "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-06, + "cache_read_input_token_cost_flex": 1.32e-07, + "cache_read_input_token_cost_priority": 5.28e-07, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-05, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 5.28e-06, "output_cost_per_token": 1.584e-05, "output_cost_per_token_above_272k_tokens": 2.376e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-05, + "output_cost_per_token_flex": 7.92e-06, + "output_cost_per_token_priority": 3.168e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { @@ -69746,12 +69851,28 @@ "input_cost_per_token_above_272k_tokens": 5.28e-07, "cache_creation_input_token_cost": 3.3e-07, "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-06, + "cache_creation_input_token_cost_flex": 1.65e-07, + "cache_creation_input_token_cost_priority": 6.6e-07, "cache_read_input_token_cost": 2.64e-08, "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-07, + "cache_read_input_token_cost_flex": 1.32e-08, + "cache_read_input_token_cost_priority": 5.28e-08, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-07, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-06, + "input_cost_per_token_flex": 1.32e-07, + "input_cost_per_token_priority": 5.28e-07, "output_cost_per_token": 1.584e-06, "output_cost_per_token_above_272k_tokens": 2.376e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-06, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-06, + "output_cost_per_token_flex": 7.92e-07, + "output_cost_per_token_priority": 3.168e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { @@ -69800,20 +69921,36 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "input_cost_per_token": 1.5e-06, "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 2.4e-07 + "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_batches": 1.2e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.625e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06 }, "bedrock_mantle/us-gov-west-1/xai.grok-4.6": { "use_openai_responses_path": true, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69827,11 +69964,17 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": { "input_cost_per_token": 4.8e-08, + "input_cost_per_token_flex": 2.4e-08, + "input_cost_per_token_priority": 8.4e-08, "output_cost_per_token": 9.6e-08, "litellm_provider": "bedrock_mantle", "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4.8e-08, + "output_cost_per_token_flex": 4.8e-08, + "output_cost_per_token_priority": 1.68e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69845,11 +69988,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.56e-07, + "input_cost_per_token_batches": 7.8e-08, + "input_cost_per_token_flex": 7.8e-08, + "input_cost_per_token_priority": 2.73e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69863,11 +70013,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": { "input_cost_per_token": 1.68e-07, + "input_cost_per_token_batches": 8.4e-08, + "input_cost_per_token_flex": 8.4e-08, + "input_cost_per_token_priority": 2.94e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69882,12 +70039,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69901,12 +70064,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70022,10 +70191,17 @@ "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70040,12 +70216,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70059,12 +70241,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5d91941a953..25c7c3593ed 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15715,8 +15715,10 @@ "input_cost_per_image": 7.2e-05, "input_cost_per_video_per_second": 0.00084, "input_cost_per_audio_per_second": 0.000168, + "input_cost_per_token_batches": 8.1e-08, "output_cost_per_token": 0.0, "output_vector_size": 3072, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_embedding_image_input": true, "supports_image_input": true, "supports_video_input": true, @@ -15725,11 +15727,14 @@ "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.8e-08, "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -15740,11 +15745,14 @@ "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 1.05e-08, "input_cost_per_token": 4.2e-08, + "input_cost_per_token_batches": 2.1e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.68e-07, + "output_cost_per_token_batches": 8.4e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -15752,12 +15760,21 @@ }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, "input_cost_per_token": 9.6e-07, + "input_cost_per_token_batches": 4.8e-07, + "input_cost_per_token_flex": 4.8e-07, + "input_cost_per_token_priority": 1.68e-06, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.84e-06, + "output_cost_per_token_batches": 1.92e-06, + "output_cost_per_token_flex": 1.92e-06, + "output_cost_per_token_priority": 6.72e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -69117,12 +69134,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69135,12 +69158,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69151,12 +69180,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69165,12 +69200,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69182,12 +69223,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69201,12 +69248,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69412,12 +69465,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69430,12 +69489,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69446,12 +69511,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69460,12 +69531,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69477,12 +69554,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69496,12 +69579,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69709,12 +69798,28 @@ "input_cost_per_token_above_272k_tokens": 5.28e-06, "cache_creation_input_token_cost": 3.3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-05, + "cache_creation_input_token_cost_flex": 1.65e-06, + "cache_creation_input_token_cost_priority": 6.6e-06, "cache_read_input_token_cost": 2.64e-07, "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-06, + "cache_read_input_token_cost_flex": 1.32e-07, + "cache_read_input_token_cost_priority": 5.28e-07, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-05, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 5.28e-06, "output_cost_per_token": 1.584e-05, "output_cost_per_token_above_272k_tokens": 2.376e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-05, + "output_cost_per_token_flex": 7.92e-06, + "output_cost_per_token_priority": 3.168e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { @@ -69746,12 +69851,28 @@ "input_cost_per_token_above_272k_tokens": 5.28e-07, "cache_creation_input_token_cost": 3.3e-07, "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-06, + "cache_creation_input_token_cost_flex": 1.65e-07, + "cache_creation_input_token_cost_priority": 6.6e-07, "cache_read_input_token_cost": 2.64e-08, "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-07, + "cache_read_input_token_cost_flex": 1.32e-08, + "cache_read_input_token_cost_priority": 5.28e-08, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-07, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-06, + "input_cost_per_token_flex": 1.32e-07, + "input_cost_per_token_priority": 5.28e-07, "output_cost_per_token": 1.584e-06, "output_cost_per_token_above_272k_tokens": 2.376e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-06, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-06, + "output_cost_per_token_flex": 7.92e-07, + "output_cost_per_token_priority": 3.168e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { @@ -69800,20 +69921,36 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "input_cost_per_token": 1.5e-06, "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 2.4e-07 + "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_batches": 1.2e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.625e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06 }, "bedrock_mantle/us-gov-west-1/xai.grok-4.6": { "use_openai_responses_path": true, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69827,11 +69964,17 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": { "input_cost_per_token": 4.8e-08, + "input_cost_per_token_flex": 2.4e-08, + "input_cost_per_token_priority": 8.4e-08, "output_cost_per_token": 9.6e-08, "litellm_provider": "bedrock_mantle", "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4.8e-08, + "output_cost_per_token_flex": 4.8e-08, + "output_cost_per_token_priority": 1.68e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69845,11 +69988,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.56e-07, + "input_cost_per_token_batches": 7.8e-08, + "input_cost_per_token_flex": 7.8e-08, + "input_cost_per_token_priority": 2.73e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69863,11 +70013,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": { "input_cost_per_token": 1.68e-07, + "input_cost_per_token_batches": 8.4e-08, + "input_cost_per_token_flex": 8.4e-08, + "input_cost_per_token_priority": 2.94e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69882,12 +70039,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69901,12 +70064,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70022,10 +70191,17 @@ "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70040,12 +70216,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70059,12 +70241,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" From 74783d38ea7c0520378bc5fcbe127dc3ee36fa56 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 4 Nov 2026 15:09:17 -0800 Subject: [PATCH 14/19] fix(github_copilot): let per-user connections work with a custom GitHub client ID (#50116) GitHub only issues Copilot tokens through copilot_internal/v2/token to its own Copilot apps, so a custom GITHUB_COPILOT_CLIENT_ID signed users in and then failed with no Copilot access. With a custom client ID, look up the account at copilot_internal/user for the seat and Copilot API address and send the user's GitHub token straight to Copilot. The default client ID keeps the token swap. --- litellm/llms/github_copilot/github_host.py | 7 ++ litellm/llms/github_copilot/per_user_auth.py | 58 +++++++++++-- .../llms/github_copilot/test_github_host.py | 1 + .../llms/github_copilot/test_per_user_auth.py | 84 +++++++++++++++++++ 4 files changed, 141 insertions(+), 9 deletions(-) diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py index 4fc921d490d..cfe122ebac4 100644 --- a/litellm/llms/github_copilot/github_host.py +++ b/litellm/llms/github_copilot/github_host.py @@ -21,8 +21,10 @@ class GithubHostEndpoints: device_code_url: str access_token_url: str api_key_url: str + copilot_user_url: str user_url: str client_id: str + exchanges_copilot_token: bool copilot_api_hosts: frozenset[str] copilot_api_host_suffix: str | None default_copilot_api_base: str @@ -49,14 +51,17 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: if normalized is None: raise ValueError("host must be github.com or a single-subdomain GHE.com hostname") client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) + exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID if normalized == GITHUB_COPILOT_DEFAULT_HOST: return GithubHostEndpoints( host=normalized, device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL), access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL), api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL), + copilot_user_url="https://api.github.com/copilot_internal/user", user_url="https://api.github.com/user", client_id=client_id, + exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({"githubcopilot.com"}), copilot_api_host_suffix=".githubcopilot.com", default_copilot_api_base="https://api.githubcopilot.com", @@ -70,8 +75,10 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: device_code_url=f"https://{normalized}/login/device/code", access_token_url=f"https://{normalized}/login/oauth/access_token", api_key_url=f"https://{api_host}/copilot_internal/v2/token", + copilot_user_url=f"https://{api_host}/copilot_internal/user", user_url=f"https://{api_host}/user", client_id=client_id, + exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({copilot_api_host}), copilot_api_host_suffix=None, default_copilot_api_base=f"https://{copilot_api_host}", diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 75a10f2e671..e210414ce8f 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -46,6 +46,7 @@ _LLM_PROVIDER: Final = "github_copilot" _DEVICE_FLOW_SCOPE: Final = "read:user" _SESSION_CACHE_MAX_SIZE: Final = 1000 _SYNC_LOCK_COUNT: Final = 64 +_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800 class _SessionCache(Protocol): @@ -125,12 +126,20 @@ class _CopilotEndpointsPayload(TypedDict): api: ReadOnly[NotRequired[object]] +@with_config(ConfigDict(extra="allow", strict=True)) +class _CopilotAccountPayload(TypedDict): + access_type_sku: ReadOnly[NotRequired[object]] + chat_enabled: ReadOnly[NotRequired[object]] + endpoints: ReadOnly[NotRequired[object]] + + @with_config(ConfigDict(extra="allow", strict=True)) class _GithubUserPayload(TypedDict): login: ReadOnly[str] _COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload) +_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload) _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload) _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload) @@ -273,6 +282,8 @@ def _session_from_response( llm_provider=_LLM_PROVIDER, model="", ) + if not endpoints.exchanges_copilot_token: + return _session_from_account(response, github_token, credential_name, cache_key, endpoints) try: payload: Final = _COPILOT_TOKEN.validate_python(response.json()) except Exception as e: @@ -282,13 +293,7 @@ def _session_from_response( model="", ) from e - payload_endpoints: Final[object] = payload.get("endpoints") - endpoints_api: Final[object] = ( - _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") - if isinstance(payload_endpoints, Mapping) - else None - ) - api_base: Final = validated_copilot_api_base(endpoints_api, endpoints) + api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base) expires_at: Final = _expires_at_seconds(payload.get("expires_at")) if expires_at > 0: @@ -296,6 +301,41 @@ def _session_from_response( return session +def _session_from_account( + response: httpx.Response, + github_token: str, + credential_name: str, + cache_key: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: + try: + payload: Final = _COPILOT_ACCOUNT.validate_python(response.json()) + except Exception as e: + raise ServiceUnavailableError( + message="GitHub Copilot account lookup returned an invalid response", + llm_provider=_LLM_PROVIDER, + model="", + ) from e + if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access": + _SESSION_CACHE.delete_cache(cache_key) + raise _reject_connection(credential_name) + + api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) + session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base) + _cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS) + return session + + +def _endpoints_api(payload_endpoints: object) -> object: + if not isinstance(payload_endpoints, Mapping): + return None + return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") + + +def _session_lookup_url(endpoints: GithubHostEndpoints) -> str: + return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url + + def _to_int(raw: object) -> int: if isinstance(raw, bool): return 0 @@ -337,7 +377,7 @@ def _do_exchange( _SyncGetClient, litellm.module_level_client ) response: Final = client.get( - endpoints.api_key_url, + _session_lookup_url(endpoints), headers=github_api_headers(github_token), ) except Exception as e: @@ -367,7 +407,7 @@ async def _do_aexchange( client: Final = _new_async_github_client() try: response: Final = await client.get( - endpoints.api_key_url, + _session_lookup_url(endpoints), headers=github_api_headers(github_token), ) except Exception as e: diff --git a/tests/unit/llms/github_copilot/test_github_host.py b/tests/unit/llms/github_copilot/test_github_host.py index df2efe4669b..daf5bfaab80 100644 --- a/tests/unit/llms/github_copilot/test_github_host.py +++ b/tests/unit/llms/github_copilot/test_github_host.py @@ -54,6 +54,7 @@ def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.Monke assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code" assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token" assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token" + assert endpoints.copilot_user_url == "https://api.octo.ghe.com/copilot_internal/user" assert endpoints.user_url == "https://api.octo.ghe.com/user" assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com" assert endpoints.client_id == "custom-client" diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index cf6fb9fa71f..5274f0c1cb8 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -307,6 +307,71 @@ async def test_aexchange_401_raises_and_evicts(): ) +ACCOUNT_URL = "https://api.github.com/copilot_internal/user" +_CUSTOM_CLIENT_ID: Final = "Ov23liCustomOrgApp" +_ACCOUNT_API_BASE: Final = "https://api.business.githubcopilot.com" + + +def _account_payload(access_type_sku="copilot_enterprise_seat", chat_enabled=True, api_base=_ACCOUNT_API_BASE): + return {"access_type_sku": access_type_sku, "chat_enabled": chat_enabled, "endpoints": {"api": api_base}} + + +def _account_responder(payload=None, status=200, model_payload=None): + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == ACCOUNT_URL: + return httpx.Response(status, json=payload or _account_payload(), request=request) + if model_payload is not None and request.url.host.endswith("githubcopilot.com"): + return httpx.Response(200, json=model_payload, request=request) + pytest.fail(f"unexpected request to {request.url}") + + return respond + + +@pytest.fixture +def custom_client_id(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", _CUSTOM_CLIENT_ID) + + +@pytest.mark.asyncio +async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_swap(custom_client_id): + endpoints: Final = github_host_endpoints("github.com") + with github_http(_account_responder()) as requests: + first = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) + second = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) + + assert [str(r.url) for r in requests] == [ACCOUNT_URL] + assert requests[0].headers["Authorization"] == "token gho_org_app" + assert first.token == second.token == "gho_org_app" + assert first.api_base == _ACCOUNT_API_BASE + + +@pytest.mark.parametrize( + "payload", + [ + _account_payload(access_type_sku="no_access", chat_enabled=False), + _account_payload(chat_enabled=False), + _account_payload(access_type_sku="no_access"), + ], +) +def test_custom_client_id_without_a_copilot_seat_raises_reconnect(custom_client_id, payload): + with github_http(_account_responder(payload)): + with pytest.raises(CallerCredentialAuthenticationError): + exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + + +@pytest.mark.parametrize("status", [401, 403, 404]) +def test_custom_client_id_rejected_account_lookup_raises_reconnect(custom_client_id, status): + with github_http(_account_responder(status=status)): + with pytest.raises(CallerCredentialAuthenticationError): + exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + + +def test_custom_client_id_account_api_base_must_stay_on_copilot_hosts(custom_client_id): + with github_http(_account_responder(_account_payload(api_base="https://evil.example"))): + session = exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + assert session.api_base == "https://api.githubcopilot.com" + + def _kwargs_with_connection( token: str = "gho_user_token", github_host: str | None = "github.com", @@ -740,6 +805,25 @@ def test_e2e_embedding_per_user_caller_authorization_cannot_override_session(per _assert_session_token_on_the_wire(requests) +def test_e2e_custom_client_id_sends_the_github_token_to_the_accounts_copilot_host( + per_user_credential, custom_client_id +): + with ( + github_http(_account_responder(model_payload=_embedding_response_payload())) as requests, + _no_authenticator(), + ): + litellm.embedding( + model="github_copilot/text-embedding-3-small", + input=["hi"], + litellm_credential_name="copilot-cred", + secret_fields=_secret_fields("user-a", "gho_org_app"), + ) + assert EXCHANGE_URL not in [str(r.url) for r in requests] + upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")] + assert upstream[0].url.host == httpx.URL(_ACCOUNT_API_BASE).host + assert upstream[0].headers["authorization"] == "Bearer gho_org_app" + + def test_e2e_responses_per_user_caller_authorization_cannot_override_session(per_user_credential): with ( github_http(_copilot_surface_responder(_responses_payload())) as requests, From 85fdea0fd1d6e86495a5a79986f1a1e51e37fee0 Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 4 Nov 2026 15:12:50 -0800 Subject: [PATCH 15/19] fix(guardrails): apply guardrail masks to Bedrock InvokeModel Claude requests and replies (#49947) --- .../chat/guardrail_translation/handler.py | 47 ++- .../messages/adapters/transformation.py | 6 +- .../guardrail_translation/handler.py | 100 ++++++ .../guardrail_translation/test_handler.py | 284 +++++++++++++++++- 4 files changed, 415 insertions(+), 22 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..6436cd832c7 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -781,12 +781,17 @@ class AnthropicMessagesHandler(BaseTranslation): for event in AnthropicMessagesHandler._parse_sse_data_line(line) ) - def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: + def _translate_to_openai( + self, + data: dict, # mutable-ok: the request body the adapter reads + *, + keep_cache_control: bool = False, + ) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" ( chat_completion_compatible_request, _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + ) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai( anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)), preserve_midturn_system=True, ) @@ -813,9 +818,13 @@ class AnthropicMessagesHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + *, + request_data: dict[str, object] | None = None, # mutable-ok: the request dict guardrails receive + keep_cache_control: bool = False, ) -> Mapping[str, object]: """ - Process input messages by applying guardrails to text content. + Process input messages by applying guardrails to text content. ``request_data`` is what the guardrail sees as + the request when ``data`` is only its body, e.g. a provider pass-through route that nests the body """ messages: Final = data.get("messages") if messages is None: @@ -827,8 +836,17 @@ class AnthropicMessagesHandler(BaseTranslation): # The top-level prompt is translated on its own below so it can be hoisted in front of # any mid-turn system entries and scanned first, aligned with that structured position. - translation_source: Final = {key: value for key, value in data.items() if key != "system"} - chat_completion_compatible_request: Final = self._translate_to_openai(translation_source) + body_model: Final = data.get("model") + request_model: Final = (request_data or {}).get("model") + model: Final[str] = ( + body_model + if isinstance(body_model, str) and body_model + else (request_model if isinstance(request_model, str) else "") + ) + translation_source: Final = {**{key: value for key, value in data.items() if key != "system"}, "model": model} + chat_completion_compatible_request: Final = self._translate_to_openai( + translation_source, keep_cache_control=keep_cache_control + ) full_structured_messages: Final = cast( list[AllMessageValues], @@ -837,7 +855,9 @@ class AnthropicMessagesHandler(BaseTranslation): has_midturn_system_message: Final = any( str(message.get("role") or "").lower() == "system" for message in full_structured_messages ) - hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data) + hoisted_system_message: Final = ( + None if skip_system else self._hoisted_top_level_system_message(data, model, keep_cache_control) + ) if hoisted_system_message is not None: full_structured_messages.insert(0, hoisted_system_message) # skip_system already excluded the trusted top-level prompt (it is simply not hoisted); @@ -899,7 +919,6 @@ class AnthropicMessagesHandler(BaseTranslation): if structured_messages: inputs["structured_messages"] = structured_messages # Include model information if available - model: Final = data.get("model") if model: inputs["model"] = model original_tool_names: Final = _original_tool_names( @@ -907,7 +926,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=data, + request_data=data if request_data is None else request_data, input_type="request", logging_obj=litellm_logging_obj, ) @@ -955,6 +974,7 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted_system_message=hoisted_system_message, preserve_system_messages=has_midturn_system_message, original_tool_names=original_tool_names.used, + model=model, ) else: if guardrailed_texts and len(guardrailed_texts) != len(scanned): @@ -977,17 +997,20 @@ class AnthropicMessagesHandler(BaseTranslation): return data - def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None: + def _hoisted_top_level_system_message( + self, data: Mapping[str, object], model: str, keep_cache_control: bool + ) -> AllMessageValues | None: """Return the system message produced by translating the top-level prompt.""" system: Final = data.get("system") if not system: return None probe: Final = self._translate_to_openai( { - "model": data.get("model") or "", + "model": model, "messages": [], "system": system, - } + }, + keep_cache_control=keep_cache_control, ) hoisted: Final = probe.get("messages") or [] return hoisted[0] if hoisted else None @@ -1121,6 +1144,7 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages: bool = False, *, original_tool_names: Mapping[str, str], + model: str, ) -> None: """Write a guardrail's structured-message rewrite back without losing corrections.""" from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -1129,7 +1153,6 @@ class AnthropicMessagesHandler(BaseTranslation): ) _is_system: Final = AnthropicMessagesHandler._is_system - model: Final = str(data.get("model") or "") converted: Final[list] = [] # mutable-ok: API message payload def _convert_run(run: list) -> None: # mutable-ok: API message payload diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 4b1b5a8a11a..afcd94c5b2c 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -380,8 +380,8 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va class LiteLLMAnthropicMessagesAdapter: - def __init__(self): - pass + def __init__(self, *, keep_cache_control: bool = False) -> None: + self.keep_cache_control: Final = keep_cache_control ### FOR [BETA] `/v1/messages` endpoint support @@ -437,7 +437,7 @@ class LiteLLMAnthropicMessagesAdapter: cache_control: Final = ( source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) ) - if cache_control and model and self.target_consumes_cache_control(model): + if cache_control and (self.keep_cache_control or (model and self.target_consumes_cache_control(model))): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) if isinstance(target, dict): diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3a65a8928bf..3514a5e6e5c 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,6 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -60,6 +61,8 @@ if TYPE_CHECKING: _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream" +_INVOKE_ACTIONS: Final = frozenset({"invoke", _INVOKE_STREAM_ACTION}) +_BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN: Final = frozenset({"document", "search_result"}) _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream" @@ -73,6 +76,83 @@ def _is_invoke_stream_endpoint(endpoint: str) -> bool: return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION +def _anthropic_invoke_body( + data: Mapping[str, object], +) -> dict[str, object] | None: # mutable-ok: the body guardrails write back into + """The body of an InvokeModel call in Anthropic Messages format, which Bedrock requires to carry anthropic_version""" + endpoint: Final = data.get("endpoint") + body: Final = data.get("data") + if not isinstance(endpoint, str) or endpoint.rstrip("/").split("/")[-1] not in _INVOKE_ACTIONS: + return None + if not isinstance(body, dict): + return None + invoke_body: Final = _BODY_BY_REFERENCE.validate_python(body) + if not isinstance(invoke_body.get("anthropic_version"), str) or not isinstance(invoke_body.get("messages"), list): + return None + messages: Final = _ARRAY_BY_REFERENCE.validate_python(invoke_body["messages"]) + if not all(_is_well_formed_message(message) for message in messages): + return None + blocks: Final = _anthropic_message_blocks(messages) + if any(block.get("type") in _BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN for block in blocks): + return None + return invoke_body + + +def _body_level_request_data( + passthrough_request: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: the request dict guardrails receive + """What `/v1/messages` hands a guardrail as `request_data`: the Anthropic body at the top level, with the proxy's own keys (model, metadata, headers) on top""" + body: Final = as_str_mapping(passthrough_request.get("data")) or {} + proxy_keys: Final = {key: value for key, value in passthrough_request.items() if key != "data"} + return {**body, **proxy_keys} + + +def _is_well_formed_message(message: object) -> bool: + """Whether the Messages translation can read `message`, so a malformed body still reaches Bedrock to be rejected there""" + fields: Final = as_str_mapping(message) + if fields is None or not isinstance(fields.get("role"), str): + return False + content: Final = fields.get("content") + if isinstance(content, str): + return True + return isinstance(content, list) and all( + as_str_mapping(block) is not None for block in _ARRAY_BY_REFERENCE.validate_python(content) + ) + + +def _anthropic_message_blocks(messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + fields: Final = (as_str_mapping(message) for message in messages) + return tuple(chain.from_iterable(_anthropic_content_blocks(field.get("content")) for field in fields if field)) + + +def _anthropic_content_blocks(content: object) -> tuple[Mapping[str, object], ...]: + blocks: Final[list[Mapping[str, object]]] = [] # mutable-ok: collected while walking the content + pending: Final[list[object]] = [content] # mutable-ok: content lists still to walk, nested tool results included + while pending: + current = pending.pop() + if not isinstance(current, list): + continue + for item in _ARRAY_BY_REFERENCE.validate_python(current): + block = as_str_mapping(item) + if block is None: + continue + blocks.append(block) + if block.get("type") == "tool_result": + pending.append(block.get("content")) + return tuple(blocks) + + +def _anthropic_invoke_reply( + response: object, + request_data: object, +) -> dict[str, object] | None: # mutable-ok: the reply guardrails write back into + if not isinstance(response, dict) or not isinstance(request_data, dict): + return None + if _anthropic_invoke_body(_OBJECT_BY_REFERENCE.validate_python(request_data)) is None: + return None + return _BODY_BY_REFERENCE.validate_python(response) + + def _generic_passthrough_handler() -> "PassThroughEndpointHandler": """ Fallback for non-Converse Bedrock routes (e.g. invoke). The generic @@ -1285,6 +1365,17 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): body: Final = data.get("data") if not _is_converse_endpoint(endpoint): + passthrough_request: Final = _BODY_BY_REFERENCE.validate_python(data) + anthropic_body: Final = _anthropic_invoke_body(passthrough_request) + if anthropic_body is not None: + await AnthropicMessagesHandler().process_input_messages( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + data=anthropic_body, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + request_data=_body_level_request_data(passthrough_request), + keep_cache_control=True, + ) + return passthrough_request return await _generic_passthrough_handler().process_input_messages( data=data, guardrail_to_apply=guardrail_to_apply, @@ -1410,6 +1501,15 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): and not _is_converse_endpoint(endpoint) and not _is_invoke_stream_endpoint(endpoint) ): + anthropic_reply: Final = _anthropic_invoke_reply(response, request_data) # pyright: ignore[reportUnknownArgumentType] # the base signature types request_data as a bare dict + if anthropic_reply is not None: + return await AnthropicMessagesHandler().process_output_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + response=anthropic_reply, # pyright: ignore[reportArgumentType] # the Anthropic reply dict, written back in place + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=_body_level_request_data(_OBJECT_BY_REFERENCE.validate_python(request_data or {})), + ) return await _generic_passthrough_handler().process_output_response( response=response, guardrail_to_apply=guardrail_to_apply, diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index aec1672b335..1c0ce5a334a 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1394,6 +1394,7 @@ class _ScriptedGuardrail(CustomGuardrail): self.skip_system_message_in_guardrail = skip_system self.scan_only_tool_results = scan_only_tool_results self.received: tuple[GenericGuardrailAPIInputs, ...] = () + self.request_data: tuple[Mapping[str, object], ...] = () async def apply_guardrail( self, @@ -1403,6 +1404,7 @@ class _ScriptedGuardrail(CustomGuardrail): logging_obj: LiteLLMLoggingObj | None = None, ) -> Mapping[str, object]: self.received = (*self.received, copy.deepcopy(inputs)) + self.request_data = (*self.request_data, request_data) return self.reply(inputs) @@ -1429,16 +1431,30 @@ def _converse_request(body: dict) -> dict: async def _sent_to_bedrock(data: dict) -> dict: - """The Converse body Bedrock receives when the passthrough route sends ``data``.""" + """The body Bedrock receives when the passthrough route sends ``data``.""" received: list[httpx.Request] = [] + action: Final = data["endpoint"].rsplit("/", 1)[-1] + def respond(request: httpx.Request) -> httpx.Response: received.append(request) - reply = { - "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, - } + reply: Final = ( + { + "id": "msg_bdrk_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + if action == "invoke" + else { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + ) return httpx.Response(200, json=reply, request=request) router = litellm.Router( @@ -1457,7 +1473,7 @@ async def _sent_to_bedrock(data: dict) -> dict: await router.allm_passthrough_route( model="bedrock-model", method="POST", - endpoint="/model/bedrock-model/converse", + endpoint=f"/model/bedrock-model/{action}", data=data["data"], client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), ) @@ -3843,6 +3859,260 @@ class TestInvokeStreamGuardrailRelay: assert rewritten_key == received_key +def _anthropic_invoke_data( + action: str = "invoke", +) -> dict[str, object]: # mutable-ok: tests write into the request the handler gets + return { + "endpoint": f"model/anthropic.claude-sonnet/{action}", + "custom_llm_provider": "bedrock", + "model": "anthropic.claude-sonnet", + "litellm_call_id": "call-1", + "data": { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 512, + "system": "You are a travel assistant. Card on file 4111 1111 1111 1111", + "tools": [ + { + "name": "find_trips", + "description": "Search trips", + "input_schema": {"type": "object", "properties": {"port": {"type": "string"}}}, + } + ], + "messages": [ + {"role": "user", "content": "Pay with 4111 1111 1111 1111"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "find_trips", "input": {"port": "Miami"}} + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card 4111 1111 1111 1111 charged"} + ], + }, + ], + }, + } + + +def _masked_anthropic_invoke_body() -> Mapping[str, object]: + body: Final = _anthropic_invoke_data()["data"] + return { + **body, + "system": "You are a travel assistant. Card on file [CARD] 1111 1111 1111", + "messages": [ + {"role": "user", "content": "Pay with [CARD] 1111 1111 1111"}, + body["messages"][1], + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card [CARD] 1111 1111 1111 charged"} + ], + }, + ], + } + + +def _rows_masked(inputs: GenericGuardrailAPIInputs) -> Mapping[str, object]: + rows = [ + {**row, "content": row["content"].replace("4111", "[CARD]")} if isinstance(row.get("content"), str) else row + for row in inputs["structured_messages"] + ] + return {"structured_messages": rows} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_anthropic_invoke_masks_are_written_back_into_the_body(action: str) -> None: + guardrail = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}) + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + _anthropic_invoke_data(action), guardrail + ) + + assert result["data"] == _masked_anthropic_invoke_body() + + +@pytest.mark.asyncio +async def test_anthropic_invoke_masks_reach_bedrock() -> None: + data: Final = _anthropic_invoke_data() + + await BedrockPassthroughGuardrailHandler().process_input_messages( + data, _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}) + ) + + assert await _sent_to_bedrock(data) == _masked_anthropic_invoke_body() + + +@pytest.mark.asyncio +async def test_anthropic_invoke_guardrail_gets_roles_tools_and_the_body_at_the_top_level_of_request_data() -> None: + data: Final = _anthropic_invoke_data() + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + inputs = guardrail.received[-1] + assert [row["role"] for row in inputs["structured_messages"]] == ["system", "user", "assistant", "tool"] + assert [tool["function"]["name"] for tool in inputs["tools"]] == ["find_trips"] + assert inputs["model"] == data["model"] + request_data: Final = guardrail.request_data[-1] + assert request_data["messages"] is data["data"]["messages"] + assert request_data["endpoint"] == data["endpoint"] + assert request_data["litellm_call_id"] == data["litellm_call_id"] + assert "data" not in request_data + + +@pytest.mark.asyncio +async def test_anthropic_invoke_structured_message_rewrites_are_written_back() -> None: + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + _anthropic_invoke_data(), _ScriptedGuardrail(_rows_masked) + ) + + assert result["data"]["messages"][0] == { + "role": "user", + "content": [{"type": "text", "text": "Pay with [CARD] 1111 1111 1111"}], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["anthropic.claude-sonnet", "prod-model", ""]) +async def test_anthropic_invoke_structured_rewrite_keeps_cache_control_whatever_the_model_is_called( + model: str, +) -> None: + cached: Final = {"type": "ephemeral"} + data: Final = _anthropic_invoke_data() + data["model"] = model + data["data"]["system"] = [{"type": "text", "text": "Card on file 4111 1111 1111 1111", "cache_control": cached}] + data["data"]["tools"] = [ + { + "name": "find_trips", + "description": "Search trips", + "input_schema": {"type": "object", "properties": {}}, + "cache_control": cached, + } + ] + data["data"]["messages"] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Pay with 4111 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}} + ], + } + ] + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + data, + _ScriptedGuardrail( + lambda inputs: { + "structured_messages": json.loads(json.dumps(inputs["structured_messages"]).replace("4111", "[CARD]")), + "tools": json.loads(json.dumps(inputs["tools"])), + } + ), + ) + + assert result["data"]["system"] == [ + {"type": "text", "text": "Card on file [CARD] 1111 1111 1111", "cache_control": cached} + ] + assert result["data"]["messages"][0]["content"] == [ + {"type": "text", "text": "Pay with [CARD] 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}} + ] + assert result["data"]["tools"][0]["cache_control"] == cached + + +@pytest.mark.asyncio +async def test_invoke_body_without_anthropic_version_is_scanned_whole() -> None: + data: Final = {**_anthropic_invoke_data(), "data": {"messages": [{"role": "user", "content": "hello"}]}} + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert "structured_messages" not in guardrail.received[-1] + + +_PLAIN_TEXT_DOCUMENT: Final = { + "type": "document", + "source": {"type": "text", "media_type": "text/plain", "data": "Card 4111 1111 1111 1111"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + [ + pytest.param([{"type": "text", "text": "Summarize this"}, _PLAIN_TEXT_DOCUMENT], id="document"), + pytest.param( + [ + { + "type": "search_result", + "source": "https://example.com/cards", + "title": "Cards", + "content": [{"type": "text", "text": "Card 4111 1111 1111 1111"}], + } + ], + id="search_result", + ), + pytest.param( + [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [_PLAIN_TEXT_DOCUMENT]}], + id="document_in_tool_result", + ), + ], +) +async def test_anthropic_invoke_body_with_blocks_the_messages_handler_does_not_scan_is_scanned_whole(content: Sequence[object]) -> None: + invoke: Final = _anthropic_invoke_data() + data: Final = {**invoke, "data": {**invoke["data"], "messages": [{"role": "user", "content": content}]}} + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert "structured_messages" not in guardrail.received[-1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "messages", + [ + pytest.param([{"content": "hi"}], id="message_without_role"), + pytest.param([{"role": "user", "content": ["hi"]}], id="bare_string_in_content"), + pytest.param([{"role": "user", "content": None}], id="null_content"), + pytest.param(["hi"], id="message_is_a_string"), + ], +) +async def test_anthropic_invoke_malformed_messages_are_scanned_whole_and_left_for_bedrock_to_reject( + messages: Sequence[object], +) -> None: + invoke: Final = _anthropic_invoke_data() + data: Final = {**invoke, "data": {**invoke["data"], "messages": messages}} + guardrail = _replying({}) + + result: Final = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert result["data"]["messages"] == messages + + +@pytest.mark.asyncio +async def test_anthropic_invoke_response_masks_are_written_back() -> None: + response = { + "id": "msg_bdrk_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Charged 4111 1111 1111 1111"}], + "stop_reason": "end_turn", + } + + result: Final = await BedrockPassthroughGuardrailHandler().process_output_response( + response=response, + guardrail_to_apply=_ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}), + request_data=_anthropic_invoke_data(), + ) + + assert result == {**response, "content": [{"type": "text", "text": "Charged [CARD] 1111 1111 1111"}]} + + def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None: event: Final = { "type": "content_block_delta", From b401ac6aafd9df4e3242b881291ad52974c3c536 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 15:47:35 -0800 Subject: [PATCH 16/19] fix(cost): bill Batch cached tokens at the Standard cache price when a row has no Batch cache price (#50062) * fix(cost): bill Batch cached tokens at the Standard cache price when a row has no Batch cache price * fix(cost): resolve the Standard cache price for the prompt length in the Batch fallback --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 28 +++++++- tests/unit/test_cost_calculator.py | 101 ++++++++++++++++++++++++++++- 2 files changed, 125 insertions(+), 4 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9154a6fdc30..100c84fc0fc 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3323,15 +3323,35 @@ def batch_cost_calculator( return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift +def _standard_rate(model_info: ModelInfo, key: str, prompt_tokens: int) -> float | None: + return get_cost_per_unit(model_info, get_modality_token_cost_key(model_info, key, prompt_tokens), None) + + +def _cached_audio_batch_fallback( + model_info: ModelInfo, batch_cache_read_rate: float | None, cache_read_rate: float | None, prompt_tokens: int +) -> float: + if batch_cache_read_rate is not None: + return batch_cache_read_rate + standard_audio_rate: Final = _standard_rate(model_info, "cache_read_input_audio_token_cost", prompt_tokens) + if standard_audio_rate is not None: + return standard_audio_rate + return 0.0 if cache_read_rate is None else cache_read_rate + + def _batch_prompt_cost( usage: Usage, model_info: ModelInfo, input_rate: float, - cache_read_rate: float | None, + batch_cache_read_rate: float | None, cache_creation_rate: float | None, cache_creation_rates_by_ttl: Mapping[str, float], ) -> float: details: Final = parse_prompt_tokens_details(usage) + cache_read_rate: Final = ( + _standard_rate(model_info, "cache_read_input_token_cost", usage.prompt_tokens) + if batch_cache_read_rate is None + else batch_cache_read_rate + ) cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0 cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0 bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl) @@ -3352,7 +3372,11 @@ def _batch_prompt_cost( + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate) + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0) + cached_audio_tokens - * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0) + * _batch_rate( + model_info, + "cache_read_input_audio_token_cost_batches", + _cached_audio_batch_fallback(model_info, batch_cache_read_rate, cache_read_rate, usage.prompt_tokens), + ) + ( calculate_cache_writing_cost( written_tokens, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 6a63a502dd3..6de368fbf00 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -4531,6 +4531,103 @@ def test_batch_cost_calculator_falls_back_to_text_batch_rate_for_modalities(): assert prompt_cost == pytest.approx(100 * 1e-7) +def _cached_audio_batch_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=1000, audio_tokens=1000, cached_tokens_details=CachedTokensDetails(audio_tokens=1000) + ), + ) + + +def _cached_text_batch_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=1000), + ) + + +@pytest.mark.parametrize( + ("audio_batch_cache_price", "usage", "expected_rate"), + [ + pytest.param(None, _cached_audio_batch_usage(), 1e-7, id="cached-audio-bills-the-standard-audio-cache-price"), + pytest.param(None, _cached_text_batch_usage(), 3e-8, id="cached-text-bills-the-standard-cache-price"), + pytest.param(8e-8, _cached_audio_batch_usage(), 8e-8, id="an-explicit-audio-batch-cache-price-still-wins"), + ], +) +def test_batch_cost_calculator_bills_cached_tokens_at_the_standard_cache_price_when_the_row_has_no_batch_cache_price( + audio_batch_cache_price: float | None, usage: Usage, expected_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 3e-7, + "input_cost_per_audio_token": 1e-6, + "cache_read_input_token_cost": 3e-8, + "cache_read_input_audio_token_cost": 1e-7, + "input_cost_per_token_batches": 1.5e-7, + "input_cost_per_audio_token_batches": 5e-7, + "cache_read_input_audio_token_cost_batches": audio_batch_cache_price, + } + + prompt_cost, _ = batch_cost_calculator( + usage=usage, model="gemini-2.5-flash", custom_llm_provider="vertex_ai", model_info=model_info + ) + + assert prompt_cost == pytest.approx(1000 * expected_rate) + + +def test_batch_cost_calculator_bills_long_context_cached_tokens_at_the_standard_long_context_cache_price() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1.25e-6, + "input_cost_per_token_above_200k_tokens": 2.5e-6, + "cache_read_input_token_cost": 1.25e-7, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-7, + "input_cost_per_token_batches": 6.25e-7, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-6, + } + usage: Final = Usage( + prompt_tokens=300_000, + completion_tokens=0, + total_tokens=300_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300_000), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, model="gemini-2.5-pro", custom_llm_provider="vertex_ai", model_info=model_info + ) + + assert prompt_cost == pytest.approx(300_000 * 2.5e-7) + + +def _batch_and_standard_cached_audio_cost(model: str, provider: str) -> tuple[str, float, float]: + batch_prompt_cost, _ = batch_cost_calculator(_cached_audio_batch_usage(), model, provider) + standard_prompt_cost, _ = litellm.cost_per_token( + model=model, custom_llm_provider=provider, usage_object=_cached_audio_batch_usage() + ) + return model, batch_prompt_cost, standard_prompt_cost + + +def test_batch_cached_audio_never_costs_more_than_standard_on_any_row_with_a_batch_input_price( + _local_model_cost_map: None, +) -> None: + rows: Final = tuple( + (name, entry["litellm_provider"]) + for name, entry in litellm.model_cost.items() + if isinstance(entry, dict) + and entry.get("input_cost_per_token_batches") is not None + and entry.get("cache_read_input_audio_token_cost") is not None + ) + assert rows, "no cost map row prices both Batch input and Standard cached audio" + + costs: Final = tuple(_batch_and_standard_cached_audio_cost(name, provider) for name, provider in rows) + overpriced: Final = tuple(cost for cost in costs if cost[1] > cost[2]) + + assert overpriced == () + + @pytest.mark.parametrize( ("image_batch_rate", "image_tokens", "expected_completion_cost"), ( @@ -6254,7 +6351,7 @@ def test_batch_cost_calculator_bills_cached_tokens_at_the_flat_batch_cached_rate assert prompt_cost == pytest.approx(100 * 1e-6 + 900 * 1e-7) -def test_batch_cost_calculator_bills_cached_tokens_at_the_batch_input_rate_without_a_cached_batch_rate( +def test_batch_cost_calculator_bills_cached_tokens_at_the_standard_cache_rate_without_a_cached_batch_rate( _tiered_batch_models: None, ) -> None: from litellm.cost_calculator import batch_cost_calculator @@ -6263,7 +6360,7 @@ def test_batch_cost_calculator_bills_cached_tokens_at_the_batch_input_rate_witho usage=_cached_usage(300_048, 300_045, 11), model=_FLAT_CACHE_BATCH_MODEL, custom_llm_provider="openai" ) - assert prompt_cost == pytest.approx(300_048 * 3e-6) + assert prompt_cost == pytest.approx(3 * 3e-6 + 300_045 * 2e-7) def _cache_write_usage(prompt_tokens: int, cache_write_tokens: int, completion_tokens: int) -> Usage: From 44b85d0c5c3093543b1f911c6808836f16dd3a58 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 15:47:57 -0800 Subject: [PATCH 17/19] fix(cost_calculator): keep regional and speed multipliers for per-call custom prices and log the billed per-call entry (#50063) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 6 ++ litellm/litellm_core_utils/litellm_logging.py | 4 ++ .../test_litellm_logging.py | 47 ++++++++++++++ tests/unit/test_cost_calculator.py | 65 +++++++++++++++++++ 4 files changed, 122 insertions(+) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 100c84fc0fc..3910a554905 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -616,11 +616,17 @@ def cost_per_token( ) if per_second is not None: return per_second + if custom_llm_provider == "anthropic": + return anthropic_cost_per_token( + model=model, usage=usage_block, service_tier=service_tier, model_info=requested_model_info_override + ) return generic_cost_per_token( model=model, usage=usage_block, custom_llm_provider=custom_llm_provider or "", service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, model_info=requested_model_info_override, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1fcb06b80e8..ce57a9bd4a5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -7143,6 +7143,7 @@ class StandardLoggingPayloadSetup: init_response_obj: object, api_base: str | None = None, router_model_id: str | None = None, + custom_model_info: ModelInfo | None = None, ) -> StandardLoggingModelInformation: model_cost_name: Final = select_model_name_for_cost_calc( model=base_model if custom_pricing else None, @@ -7161,6 +7162,8 @@ class StandardLoggingPayloadSetup: if router_model_id is not None else model_cost_name ) + if model_cost_name is not None and custom_model_info is not None: + return StandardLoggingModelInformation(model_map_key=model_cost_name, model_map_value=custom_model_info) if model_cost_name is None: model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None) else: @@ -7911,6 +7914,7 @@ def get_standard_logging_object_payload( init_response_obj=init_response_obj, api_base=litellm_params.get("api_base"), router_model_id=logging_obj.get_router_model_id(), + custom_model_info=logging_obj.custom_model_info if custom_pricing else None, ) error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index a014d1d352d..e8ea0c04f6f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14594,3 +14594,50 @@ async def test_a_recovered_attempt_bills_a_failed_agentic_follow_up_but_not_a_fa assert own_cost > 0 assert after_a_dropped_stream == pytest.approx(own_cost) assert after_a_failed_follow_up == pytest.approx(own_cost + _FAILED_ATTEMPT[1]) + + +def test_standard_logging_payload_reports_the_per_call_prices_the_request_was_billed_from() -> None: + """A direct call that sets its own prices is logged at those prices, not at the catalog row's rates.""" + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + from litellm.litellm_core_utils.per_call_pricing import per_call_model_info + + rates: Final = {"input_cost_per_token": 3e-06, "output_cost_per_token": 7e-06} + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + logging_obj: Final = LitellmLogging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=now, + litellm_call_id="per-call-prices", + function_id="fn", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-nano", user="", optional_params={}, litellm_params=rates, custom_llm_provider="openai" + ) + logging_obj.set_per_call_pricing(per_call_model_info("gpt-5.4-nano", "openai", rates, "chat")) + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=6, total_tokens=16), + ) + + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "custom_llm_provider": "openai", + "litellm_params": rates, + "response_cost": logging_obj.response_cost_calculator(result=response), + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(10 * 3e-06 + 6 * 7e-06) + model_map_value: Final = payload["model_map_information"]["model_map_value"] + assert model_map_value is not None + assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (3e-06, 7e-06) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 6de368fbf00..49d51e99790 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -11645,3 +11645,68 @@ def test_upstream_cost_header_total_on_an_unmapped_alias_is_split_by_the_served_ input_share: Final = (10 * 0.001) / (10 * 0.001 + 20 * 0.002) assert breakdown["input_cost"] == pytest.approx(0.25 * input_share) assert breakdown["output_cost"] == pytest.approx(0.25 * (1 - input_share)) + + +_PER_CALL_RATES: Final = MappingProxyType({"input_cost_per_token": 2e-06, "output_cost_per_token": 4e-06}) + + +def _per_call_priced_reply(model: str, usage_extras: Mapping[str, str]) -> ModelResponse: + usage: Final = Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000) + for field, value in usage_extras.items(): + setattr(usage, field, value) + return ModelResponse(model=model, choices=[Choices(message=Message(content="hi"))], usage=usage) + + +@pytest.mark.parametrize( + ("row", "call_kwargs", "usage_extras", "multiplier"), + [ + pytest.param( + {"litellm_provider": "openai", "mode": "chat", "regional_processing_uplift_multiplier_eu": 1.25}, + { + "model": "openai/lit-eu-residency-test", + "api_key": "fake-key", + "api_base": "https://eu.api.openai.com/v1", + }, + {}, + 1.25, + id="openai-data-residency", + ), + pytest.param( + {"litellm_provider": "vertex_ai", "mode": "chat", "regional_endpoint_uplift_multiplier": 1.5}, + { + "model": "vertex_ai/lit-regional-endpoint-test", + "vertex_project": "lit-project", + "vertex_location": "us-east5", + "vertex_credentials": "{}", + }, + {}, + 1.5, + id="vertex-regional-endpoint", + ), + pytest.param( + {"litellm_provider": "anthropic", "mode": "chat", "provider_specific_entry": {"fast": 2.0}}, + {"model": "anthropic/lit-fast-mode-test", "api_key": "fake-key"}, + {"speed": "fast"}, + 2.0, + id="anthropic-fast-mode", + ), + ], +) +def test_per_call_custom_prices_keep_the_regional_and_speed_multipliers( + monkeypatch: pytest.MonkeyPatch, + row: Mapping[str, object], + call_kwargs: Mapping[str, str], + usage_extras: Mapping[str, str], + multiplier: float, +) -> None: + monkeypatch.setitem(litellm.model_cost, call_kwargs["model"], {**row, **_PER_CALL_RATES}) + _invalidate_model_cost_lowercase_map() + + response: Final = litellm.completion( + messages=[{"role": "user", "content": "hello"}], + mock_response=_per_call_priced_reply(call_kwargs["model"].split("/", 1)[1], usage_extras), + **call_kwargs, + **_PER_CALL_RATES, + ) + + assert response._hidden_params["response_cost"] == pytest.approx((1000 * 2e-06 + 1000 * 4e-06) * multiplier) From 513fe4ac8e175cde810d72094a0e71436adc310f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 00:28:25 +0000 Subject: [PATCH 18/19] chore(cost-map): add openai gpt-image-2.6-sunburst and gpt-image-2.6-flare from the pricing page (#50134) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 32 +++++++++++++++++++ model_prices_and_context_window.json | 32 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 25c7c3593ed..1aa2ac27152 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -85880,5 +85880,37 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "gpt-image-2.6-sunburst": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2.6-flare": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 25c7c3593ed..1aa2ac27152 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -85880,5 +85880,37 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "gpt-image-2.6-sunburst": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2.6-flare": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true } } From 733b7764039711cc9f7ee103c8e17b7e9f4a47d3 Mon Sep 17 00:00:00 2001 From: Oliver Yu <57586652+tc3oliver@users.noreply.github.com> Date: Thu, 5 Nov 2026 08:30:12 +0800 Subject: [PATCH 19/19] fix(least-busy): reserve the picked deployment so a burst of requests spreads out (#47484) * fix(least-busy): reserve the picked deployment so a burst of requests spreads out The in-flight count only went up in log_pre_api_call, which runs after selection and the pre-call hooks. Requests selected in that window all saw the same counts and landed on the same deployment (16 simultaneous requests over 4 deployments went as lopsided as 12/4/0/0 in the proxy). When the router selects for a request it is about to send, least-busy now records a reservation for the chosen deployment in the same synchronous step that reads the reservations, and log_pre_api_call turns it into the normal in-flight count. Selections without request kwargs (token counting, probes) do not reserve. A reservation whose request never reaches the API call expires after 30 s. The shared and per-worker counts and their success/failure decrements are unchanged. * fix(least-busy): prune expired and empty selection reservations The reservation map kept a key per model group forever, so a wildcard route with caller-chosen model names grew it without bound. Drop expired and empty entries on every write. Signed-off-by: Oliver Yu * fix(least-busy): keep pick reservations in the shared router cache Selectors that share a router cache dedupe to one registered callback, so only that instance saw log_pre_api_call. Reservations held on the other instances were never released and skewed picks for the full TTL. Storing them in the router cache's in-memory layer lets any selector release them, and the per-key TTL replaces the full-map rebuild on every write. Tests now hold the burst with a barrier instead of a sleep and use a manual clock object. * fix(least-busy): restart the reservation TTL on every write InMemoryCache keeps the TTL a live key was first given, so a reservation key that kept being rewritten still expired 30 s after its first pick and took the newer reservations with it. Drop the key before each write so the TTL counts from the latest one, and leave it deleted once it is empty. Signed-off-by: Oliver Yu * fix(router): count least-busy decisions calls and release reservations on cache hits Least-busy now reads the model group from litellm_metadata before metadata, so /v1/decisions calls that carry caller metadata are counted in flight and release their reservation at pre-call. A response cache hit releases its reservation instead of decrementing another request's in-flight count, and Responses WebSocket turn admission selects without reserving. The reservation store is reached through a typed protocol. * fix(decisions): count proxy /v1/decisions calls in least-busy's in-flight totals * fix(least-busy): keep reservations out of the shared router cache and count requests that start during a Redis read Least-busy now reads a call's deployment from the routing metadata the router sets, so /v1/decisions logging params stay as on main and failed decisions calls no longer start cooldowns. * style(router): wrap the least-busy reservation store line * fix(least-busy): lock the reservation read-and-update so threaded sync calls keep each other's picks Router.completion from a thread pool and sync success callbacks can read and rewrite a deployment's reservations at the same time, so one thread's pick or release could be lost. * fix(least-busy): release a pick whose call ends before reaching the provider A call that fails or returns before the provider handoff (a 400 for an unsupported param, token counting, a proxy-gate rejection without a start log, batch create) now drops its reservation instead of decrementing a count it never added. Router.select picks and batch cancel no longer reserve, since neither makes a call that logs. * test(least-busy): give callback fixtures a fixed call-start time The routing-group tests fed least-busy a finished call without the start time litellm stamps before counting it, so the count stayed up after the fix. Fixtures now use a fixed timestamp instead of the clock. * fix(least-busy): free tier-group picks, stop lookups reserving, and count only started attempts Reservations are keyed by deployment id, so a complexity, auto, adaptive or quality router frees its tier pick when the call starts. Each started attempt is recorded under its call id, and only a recorded start lowers the count when the attempt ends, so a fallback hop rejected before the provider frees its pick. Whether a pick reserves now comes from a context variable set by the routed call sites, so get_available_deployment and the deprecated async_get_available_deployment no longer hold a deployment, and custom routing strategies and Router subclasses keep their method signatures. * fix(least-busy): free a request's unstarted picks when it ends Each pick now carries its call id, so a start frees only its own call's pick, and the router frees every pick a request never started once that request ends. A fallback hop rejected before the provider no longer leaves its pick reserved for 30 seconds. * fix(router): free a least-busy pick when a vector store create ends * fix(router): free least-busy picks of calls made without a call id * fix(router): free the least-busy pick of a sync text completion * test(router): use fixed start times in least-busy fixtures * fix(router): free the least-busy pick of a sync generic call that fails before its provider * fix(router): read in-flight counts and start attempts under one lock so a busy deployment never looks free * fix(router): keep pass-through lookups from reserving least-busy picks * test(least-busy): wait for the helper thread to start before the timed join * fix(least-busy): keep realtime sessions out of in-flight counts --------- Signed-off-by: Oliver Yu Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router_deprecated.py | 20 +- litellm/router_python_core.py | 441 ++++++------ litellm/router_strategy/least_busy.py | 341 +++++++-- .../test_least_busy_concurrent_decisions.py | 90 +++ tests/unit/decisions/test_main.py | 39 + tests/unit/router_strategy/test_least_busy.py | 672 +++++++++++++++++- .../test_router_routing_groups.py | 37 + tests/unit/test_router/test_router.py | 12 +- .../test_router/test_router_helper_utils.py | 28 + 9 files changed, 1416 insertions(+), 264 deletions(-) create mode 100644 tests/integration/routing/test_least_busy_concurrent_decisions.py diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 83a1251d3c0..0a02408a717 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -27,11 +27,12 @@ import inspect import sys import warnings from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload +from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet +from litellm.router_strategy.least_busy import picks_reserve if TYPE_CHECKING: from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged @@ -58,6 +59,7 @@ if TYPE_CHECKING: from litellm.types.utils import ModelInfo as ModelMapInfo _T: Final = TypeVar("_T") +_P: Final = ParamSpec("_P") DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: the live tally _note increments @@ -164,6 +166,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib return _DeprecatedAttribute(replacement, read) +def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]: + @functools.wraps(lookup) + async def looked_up( + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged + ) -> _T: + with picks_reserve(False): + return await lookup(*args, **kwargs) + + return looked_up + + # (id(router), member name) pairs whose subclass override is running in this context. _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar( "router_overrides_running", default=frozenset() @@ -333,7 +347,7 @@ class DeprecatedRouterMembers: deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map) async_get_available_deployment = _attr( "select(model, request_kwargs)", - lambda r: r._core_async_get_available_deployment, + lambda r: _without_reserving(r._core_async_get_available_deployment), ) # -- observers ------------------------------------------------------------------------------ diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index c0c5d74a88a..4b5ecc82581 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -282,7 +282,15 @@ from litellm.router_strategy.complexity_router.context_compaction import ( reject_recursive_compactor, surface_for_call, ) -from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import ( + CALL_IDENTITY, + SELECTION_RESERVATION_MAX_KEYS, + LeastBusyLoggingHandler, + call_id_of, + owning_picks, + pick_reserves, + picks_reserve, +) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler @@ -1462,13 +1470,14 @@ class PythonCore(DeprecatedRouterMembers): (``pass_through``). """ if kind == "pick": - return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + with picks_reserve(False): + return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + request_kwargs=request_kwargs if request_kwargs is not None else {}, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) if kind == "healthy": return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts model=model, @@ -1480,13 +1489,14 @@ class PythonCore(DeprecatedRouterMembers): health_check_probe=health_check_probe, ) if kind == "pass_through": - return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) + with picks_reserve(False): + return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + messages=messages, + input=input, + specific_deployment=specific_deployment, + request_kwargs=request_kwargs, + ) raise ValueError(f"unknown select kind {kind!r}") def apply(self, changes: ChangeSet) -> ApplyResult: @@ -1768,6 +1778,9 @@ class PythonCore(DeprecatedRouterMembers): self._core_cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._core_least_busy_reservations: InMemoryCache = InMemoryCache( + max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS + ) self._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), @@ -2271,7 +2284,9 @@ class PythonCore(DeprecatedRouterMembers): selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler(router_cache=self._core_cache) + selector = LeastBusyLoggingHandler( + router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations + ) case RoutingStrategy.USAGE_BASED_ROUTING.value: selector = LowestTPMLoggingHandler( router_cache=self._core_cache, @@ -2736,6 +2751,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Asks the strategy selector for a deployment. Caller handles @@ -2749,6 +2765,8 @@ class PythonCore(DeprecatedRouterMembers): return await selector.async_get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing": # `LowestTPMLoggingHandler` (v1) only exposes the sync @@ -2789,6 +2807,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Sync sibling of `_select_deployment_async`. Caller handles @@ -2805,6 +2824,8 @@ class PythonCore(DeprecatedRouterMembers): return selector.get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing" | "usage-based-routing-v2": return selector.get_available_deployments( @@ -3524,12 +3545,13 @@ class PythonCore(DeprecatedRouterMembers): ) # pick the one that is available (lowest TPM/RPM) - deployment = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._drop_unsupported_classifier_reasoning_effort( deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment model=model, @@ -6104,12 +6126,13 @@ class PythonCore(DeprecatedRouterMembers): deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + messages=[{"role": "user", "content": "prompt"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -6496,37 +6519,38 @@ class PythonCore(DeprecatedRouterMembers): **kwargs, ): messages: Final = [{"role": "user", "content": prompt}] - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + kwargs["model"] = model + kwargs["prompt"] = prompt + kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) + kwargs.setdefault("metadata", {}).update({"model_group": model}) - # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model): - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + # pick the one that is available (lowest TPM/RPM) + with endpoint_selection("text_completion", model=model), picks_reserve(True): + deployment: Final = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + + data: Final = deployment["litellm_params"].copy() + for k, v in self._core_default_litellm_params.items(): + if k not in kwargs: # prioritize model-specific params > default router params + kwargs[k] = v + elif k == "metadata": + kwargs[k].update(v) + pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) + if pinned_model_id is not None: + kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults + + # call via litellm.completion() + return self._host.invoke("text_completion")( + **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} ) - - data: Final = deployment["litellm_params"].copy() - for k, v in self._core_default_litellm_params.items(): - if k not in kwargs: # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) - pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults - - # call via litellm.completion() - return self._host.invoke("text_completion")( - **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} - ) - except Exception as e: - raise e + except Exception as e: + raise e async def atext_completion( self, @@ -7671,64 +7695,69 @@ class PythonCore(DeprecatedRouterMembers): handler_name: Final = original_function.__name__ metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) + self._update_kwargs_before_fallbacks( model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + kwargs=kwargs, + metadata_variable_name=metadata_variable_name, + ) + with ( + self._video_call_pin(handler_name, kwargs), + self._route_capability_pin(handler_name), + picks_reserve(True), + ): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="generic_api_call", + splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, ) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="generic_api_call", - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] - self.total_calls[model_name] += 1 + self.total_calls[model_name] += 1 - # For passthrough routes, use the actual model from deployment - # and swap model name in endpoint if present - if "endpoint" in kwargs and kwargs["endpoint"]: - kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) - kwargs["model"] = model_name + # For passthrough routes, use the actual model from deployment + # and swap model name in endpoint if present + if "endpoint" in kwargs and kwargs["endpoint"]: + kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) + kwargs["model"] = model_name - # Perform pre-call checks for routing strategy - self.routing_strategy_pre_call_checks(deployment=deployment) + # Perform pre-call checks for routing strategy + self.routing_strategy_pre_call_checks(deployment=deployment) - custom_llm_provider: Final = provider_for_generic_call(data) + custom_llm_provider: Final = provider_for_generic_call(data) - response: Final = original_function( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - **kwargs, - } - ) + response: Final = original_function( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + **kwargs, + } + ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) - return response - except Exception as e: - verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e + self.success_calls[model_name] += 1 + verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) + return response + except Exception as e: + verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) + if model is not None: + self.fail_calls[model] += 1 + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) + raise e def embedding( self, @@ -7752,7 +7781,7 @@ class PythonCore(DeprecatedRouterMembers): deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(): + with embedding_deployment_selection(), picks_reserve(True): deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, @@ -8061,68 +8090,71 @@ class PythonCore(DeprecatedRouterMembers): Returns: VectorStoreCreateResponse """ - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + # If model is None, use the factory function approach (direct SDK call) + if model is None: + from litellm.vector_stores.main import acreate - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) + # Use the factory function to handle the call + factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") + return await factory_fn(**kwargs) - from litellm.vector_stores import acreate as avector_store_create_sdk + from litellm.vector_stores import acreate as avector_store_create_sdk - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } + parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "vector-store-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="avector_store_create", ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + model_client: Final = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs, + ) + self.total_calls[model_name] += 1 - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e + # Get custom provider from deployment params + custom_llm_provider = data.get("custom_llm_provider") + _, inferred_custom_llm_provider, _, _ = get_llm_provider( + model=data["model"], + custom_llm_provider=custom_llm_provider, + ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider + + async with self._deployment_slot( + deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span + ): + response = await avector_store_create_sdk( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + "client": model_client, + **kwargs, + } + ) + + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + + return response + except Exception as e: + verbose_router_logger.exception( + "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e + ) + if model is not None: + self.fail_calls[model] += 1 + raise e def _override_vector_store_methods_for_router(self): """ @@ -8370,12 +8402,13 @@ class PythonCore(DeprecatedRouterMembers): try: verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "batch-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(False): + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "batch-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") data: Final = deployment["litellm_params"].copy() @@ -9520,38 +9553,39 @@ class PythonCore(DeprecatedRouterMembers): mock_timeout: Final = kwargs.pop("mock_timeout", None) - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + self._handle_mock_testing_fallbacks( + kwargs=kwargs, + model_group=model_group, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) + if mock_timeout is not None: + response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) + else: + response = await self.async_function_with_retries(*args, **kwargs) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug("Async Response: %s", response) + response = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=0, + ) + return response + except Exception as e: + return await self.async_function_with_fallbacks_common_utils( + e, + disable_fallbacks, + fallbacks, + context_window_fallbacks, + content_policy_fallbacks, + model_group, + args, + kwargs, + include_fallback_errors=include_fallback_errors, + ) def _handle_mock_testing_fallbacks( self, @@ -10856,6 +10890,7 @@ class PythonCore(DeprecatedRouterMembers): messages=None, input=None, request_kwargs=routing_kwargs, + reserve=False, ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( @@ -16214,6 +16249,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=pick_reserves(default=True), ) async def _async_get_available_deployment( @@ -16223,6 +16259,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, specific_deployment: bool | None, + reserve: bool, ): if ( self._core_routing_strategy != "usage-based-routing-v2" @@ -16319,6 +16356,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=reserve and request_kwargs not in (None, {}), ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -16488,6 +16526,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17365,6 +17404,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17545,6 +17585,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 8e0c3d6a7e7..22159079f98 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,6 +1,11 @@ +import contextlib import logging -from collections.abc import Mapping, Sequence -from typing import Final +import threading +import time +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextvars import ContextVar +from typing import Final, Protocol, TypeAlias from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -8,11 +13,31 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.types.utils import CallTypes IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 +SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 +SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 +_RESERVATIONS_LOCK: Final = threading.Lock() +_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves") +_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None) + + +@contextlib.contextmanager +def picks_reserve(reserve: bool) -> Generator[None, None, None]: + token: Final = _PICK_RESERVES.set(reserve) + try: + yield + finally: + _PICK_RESERVES.reset(token) + + +def pick_reserves(default: bool) -> bool: + return _PICK_RESERVES.get(default) class _ModelInfo(TypedDict, total=False): @@ -21,10 +46,12 @@ class _ModelInfo(TypedDict, total=False): class _Metadata(TypedDict, total=False): model_group: ReadOnly[str | None] + model_info: ReadOnly[_ModelInfo | None] class _LitellmParams(TypedDict, total=False): metadata: ReadOnly[_Metadata | None] + litellm_metadata: ReadOnly[_Metadata | None] model_info: ReadOnly[_ModelInfo | None] @@ -32,6 +59,10 @@ class _CallKwargs(TypedDict, total=False): litellm_params: ReadOnly[_LitellmParams | None] +class CallIdentity(TypedDict, total=False): + litellm_call_id: ReadOnly[object] + + class _DeploymentModelInfo(TypedDict): id: ReadOnly[str | int] @@ -41,38 +72,85 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +CALL_IDENTITY: Final = TypeAdapter(CallIdentity) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) -_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) +_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None) +_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) +_Pick: TypeAlias = tuple[float, str | None] +_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None) +_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: return f"{model_group}_request_count:{deployment_id}" +def _model_group(metadata: _Metadata | None) -> str | None: + return metadata.get("model_group") if metadata else None + + +def _deployment_id(model_info: _ModelInfo | None) -> str | int | None: + return model_info.get("id") if model_info else None + + +def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None: + return _deployment_id(metadata.get("model_info")) if metadata else None + + def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: - if is_batch_retrieve_call_type(kwargs.get("call_type")): + call_type: Final = kwargs.get("call_type") + if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value: return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) except ValidationError: return None - litellm_params: Final = call.get("litellm_params") - metadata: Final = litellm_params.get("metadata") if litellm_params else None - model_info: Final = litellm_params.get("model_info") if litellm_params else None - model_group: Final = metadata.get("model_group") if metadata else None - deployment_id: Final = model_info.get("id") if model_info else None + litellm_params: Final = call.get("litellm_params") or _LitellmParams() + routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata")) + model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None) + deployment_id: Final = next( + ( + deployment_id + for deployment_id in ( + _deployment_id(litellm_params.get("model_info")), + *map(_routed_deployment_id, routing_metadata), + ) + if deployment_id is not None + ), + None, + ) if model_group is None or deployment_id is None: return None return model_group, str(deployment_id) -def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: +def _reservation_key(deployment_id: str) -> str: + return f"least_busy_reservations:{deployment_id}" + + +def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str: + return f"least_busy_started:{call_id}:{_request_count_key(*ref)}" + + +def _picked_deployments_key(call_id: str) -> str: + return f"least_busy_picked:{call_id}" + + +def call_id_of(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( - _request_count_key(model_group, str(deployment["model_info"]["id"])) - for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) ) +def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]: + return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids) + + def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) @@ -84,12 +162,21 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) -def _least_busy( - healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] -) -> Mapping[str, object] | None: - if not healthy_deployments: - return None - return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] +def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]: + return tuple(pick for pick in picks if now - pick[0] < SELECTION_RESERVATION_TTL_SECONDS) + + +def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]: + tags: Final = tuple(tag for _, tag in picks) + owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None) + if owner not in tags: + return picks + index: Final = tags.index(owner) + return (*picks[:index], *picks[index + 1 :]) + + +def _least_busy_index(counts: tuple[int, ...]) -> int: + return min(range(len(counts)), key=lambda index: counts[index]) def _warn_unreadable(model_group: str, error: Exception) -> None: @@ -111,24 +198,156 @@ def _warn_unwritable(key: str, error: Exception) -> None: ) +class _ReservationStore(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]: + return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or () + + +def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]: + return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or () + + +def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]: + return _PICKED_DEPLOYMENTS.validate_python(store.get_cache(_picked_deployments_key(call_id))) or () + + +def _put(store: _ReservationStore, key: str, values: tuple[object, ...], ttl: float) -> None: + # InMemoryCache keeps a live key's first TTL, so drop the key to restart it + store.delete_cache(key) + if values: + store.set_cache(key, values, ttl=ttl) + + +def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) -> None: + kept: Final = tuple(pick for pick in _picks(store, deployment_id) if pick[1] != call_id) + _put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS) + + +def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None: + if call_id is None: + return + if not _picked_deployments(store, call_id): + return + with _RESERVATIONS_LOCK: + for deployment_id in frozenset(_picked_deployments(store, call_id)): + _release_call_on(store, deployment_id, call_id) + store.delete_cache(_picked_deployments_key(call_id)) + + +@contextlib.contextmanager +def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]: + owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex + token: Final = _PICK_OWNER.set(owner) + try: + yield + finally: + _PICK_OWNER.reset(token) + release_unstarted_picks(store, owner) + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache): + def __init__( + self, + router_cache: DualCache, + clock: Callable[[], float] = time.monotonic, + reservation_store: _ReservationStore | None = None, + ): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) + self._clock = clock + self._reservations: Final[_ReservationStore] = ( + reservation_store + if reservation_store is not None + else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock) + ) + + def _live_reservations(self, deployment_id: str, now: float) -> tuple[_Pick, ...]: + return _unexpired(_picks(self._reservations, deployment_id), now) + + def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None: + _put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS) + + def _take_reservation(self, deployment_id: str, call_id: str | None, now: float) -> None: + self._store_reservations(deployment_id, _without_one_pick(self._live_reservations(deployment_id, now), call_id)) + + def _record_pick(self, call_id: str, deployment_id: str) -> None: + picked: Final = (*_picked_deployments(self._reservations, call_id), deployment_id) + _put(self._reservations, _picked_deployments_key(call_id), picked, SELECTION_RESERVATION_TTL_SECONDS) + + def _select( + self, + healthy_deployments: Sequence[Mapping[str, object]], + deployment_ids: tuple[str, ...], + keys: tuple[str, ...], + shared: tuple[int, ...] | None, + reserve: bool, + call_id: str | None, + ) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + with _RESERVATIONS_LOCK: + now: Final = self._clock() + counts: Final = ( + self._this_workers_counts(keys) if shared is None else self._at_least_this_workers(keys, shared) + ) + live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids) + index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True))) + if reserve: + tag: Final = call_id if call_id is not None else _PICK_OWNER.get() + self._store_reservations(deployment_ids[index], (*live[index], (now, tag))) + if tag is not None: + self._record_pick(tag, deployment_ids[index]) + return healthy_deployments[index] + + def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: + with _RESERVATIONS_LOCK: + now: Final = self._clock() + self._take_reservation(ref[1], call_id, now) + self._count_locally(_request_count_key(*ref), 1) + if call_id is not None: + key: Final = _started_attempts_key(call_id, ref) + started: Final = (*_attempt_times(self._reservations, key), now) + _put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS) + + def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool: + if call_id is None: + with _RESERVATIONS_LOCK: + self._count_locally(_request_count_key(*ref), -1) + return True + key: Final = _started_attempts_key(call_id, ref) + with _RESERVATIONS_LOCK: + started: Final = _attempt_times(self._reservations, key) + if started: + _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) + self._count_locally(_request_count_key(*ref), -1) + else: + self._take_reservation(ref[1], call_id, self._clock()) + return bool(started) @with_service_target("router_usage") def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: - self._increment(kwargs, 1) + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + self._start_attempt(ref, call_id_of(kwargs)) + self._count_shared(_request_count_key(*ref), 1) @with_service_target("router_usage") def log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -136,7 +355,7 @@ class LeastBusyLoggingHandler(CustomLogger): def log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_failure += 1 @@ -144,7 +363,7 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -152,15 +371,20 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_failure += 1 @with_service_target("router_usage") def get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -168,15 +392,19 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id) @with_service_target("router_usage") async def async_get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -184,42 +412,51 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id) - def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + def _settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None: + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): return - key: Final = _request_count_key(*ref) - redis_cache: Final = self.router_cache.redis_cache + self._count_shared(_request_count_key(*ref), -1) + + async def _async_settle(self, kwargs: Mapping[str, object]) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): + return + await self._async_count_shared(_request_count_key(*ref), -1) + + def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: + return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + + def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]: + return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True)) + + def _count_locally(self, key: str, delta: int) -> None: try: local: Final = self.router_cache.increment_cache( key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS ) if local < 0: self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return + except TypeError as e: + _warn_unwritable(key, e) + + def _count_shared(self, key: str, delta: int) -> None: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return + try: redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) - async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) - if ref is None: - return - key: Final = _request_count_key(*ref) + async def _async_count_shared(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return try: - local: Final = await self.router_cache.async_increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - if local is not None and local < 0: - await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/integration/routing/test_least_busy_concurrent_decisions.py b/tests/integration/routing/test_least_busy_concurrent_decisions.py new file mode 100644 index 00000000000..bb923534db4 --- /dev/null +++ b/tests/integration/routing/test_least_busy_concurrent_decisions.py @@ -0,0 +1,90 @@ +import json +import threading +from collections import Counter +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark = pytest.mark.timeout(240) + +_GROUP: Final = "least-busy-decider" +_DEPLOYMENTS: Final = ("least-busy-decider-a", "least-busy-decider-b") +_BURST: Final = 4 +_REPLY: Final = json.dumps( + { + "model": "Qwen/Qwen3-0.6B", + "answers": {"billing": {"type": "noul", "noul": 0.91}}, + "usage": {"input_tokens": 367, "output_tokens": 3}, + } +).encode() +_DECISION: Final = { + "model": _GROUP, + "input": "I was charged twice for order 4417.", + "questions": [{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}], +} + + +def _held_until_the_burst_arrives(arrived: threading.Barrier) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + arrived.wait() + return Reply(body=_REPLY) + + return respond + + +def _write_config(path: Path, provider_url: str) -> Path: + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _GROUP, + "litellm_params": {"model": "hosted_vllm/Qwen/Qwen3-0.6B", "api_base": provider_url}, + "model_info": {"id": deployment_id}, + } + for deployment_id in _DEPLOYMENTS + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": {"routing_strategy": "least-busy", "num_retries": 0, "disable_cooldowns": True}, + } + ) + ) + return path + + +def _decide(proxy: Gateway) -> httpx.Response: + return proxy.client.post( + "/v1/decisions", json=_DECISION, headers={"Authorization": f"Bearer {proxy.key}"}, timeout=120 + ) + + +def test_least_busy_spreads_a_burst_of_decisions_calls_across_the_group(tmp_path: Path) -> None: + arrived: Final = threading.Barrier(_BURST, timeout=60) + with ( + wire_server(_held_until_the_burst_arrives(arrived)) as upstream, + gateway_from_environment() as rig, + owned_proxy( + rig, + tmp_path, + {"STORE_MODEL_IN_DB": "False"}, + config=_write_config(tmp_path / "least-busy.yaml", upstream.url), + ) as proxy, + ThreadPoolExecutor(max_workers=_BURST) as pool, + ): + responses: Final = tuple(pool.map(lambda _: _decide(proxy), range(_BURST))) + assert [response.status_code for response in responses] == [200] * _BURST, [r.text for r in responses] + assert Counter(response.headers["x-litellm-model-id"] for response in responses) == Counter( + {deployment_id: _BURST // len(_DEPLOYMENTS) for deployment_id in _DEPLOYMENTS} + ) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 1e42e74806d..0e2813e0a0d 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,8 +6,10 @@ import copy import json import re import threading +import uuid from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import wait +from datetime import datetime from functools import partial from types import MappingProxyType from typing import Final, TypeAlias @@ -300,6 +302,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base( assert response.answers["sentiment"].choice == "positive" +@pytest.mark.asyncio +async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "jev", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("jev-a", "jev-b") + ], + routing_strategy="least-busy", + ) + + async def call_through_proxy_logging() -> object: + logging_obj, _ = litellm.utils.function_setup( + original_function="adecisions", + rules_obj=litellm.utils.Rules(), + start_time=datetime(2026, 1, 1), + model="jev", + litellm_call_id=str(uuid.uuid4()), + ) + response: Final = await router.adecisions( + model="jev", + input="Is this a defect?", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + litellm_logging_obj=logging_obj, + ) + await _drain_logging_worker() + return response._hidden_params["model_id"] + + assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a") + + def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 5b3af388ae5..f156a5ee405 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,9 +1,16 @@ import asyncio +import contextlib +import contextvars +import datetime import json import logging +import threading +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx +import openai import pytest import respx @@ -11,11 +18,19 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler -from litellm.types.utils import ModelResponse +from litellm.router_strategy.least_busy import ( + IN_FLIGHT_COUNT_TTL_SECONDS, + SELECTION_RESERVATION_TTL_SECONDS, + LeastBusyLoggingHandler, + owning_picks, + picks_reserve, + release_unstarted_picks, +) +from litellm.types.utils import CallTypes, ModelResponse GROUP: Final = "least-busy-group" DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} @@ -23,8 +38,11 @@ DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] -def _call_kwargs(deployment_id: str) -> dict[str, object]: - return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} +def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]: + return { + "litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}, + "litellm_call_id": call_id, + } class SharedRedisCounters: @@ -58,6 +76,20 @@ class SharedRedisCounters: return self.batch_get_counts(key_list) +class SlowRedisRead(SharedRedisCounters): + """Answers a count read with the counts from when the read started, after letting other work run.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + counts_when_asked: Final = self.batch_get_counts(key_list) + await asyncio.sleep(0) + self.while_reading() + return counts_when_asked + + def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double return LeastBusyLoggingHandler(router_cache=cache) @@ -81,6 +113,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None: assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +@pytest.mark.asyncio +async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None: + redis: Final = SlowRedisRead() + worker: Final = _worker(redis) + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + redis.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + def test_sync_pick_reads_the_shared_counts() -> None: shared: Final = SharedRedisCounters() streaming_worker: Final = _worker(shared) @@ -98,6 +140,44 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: + return { + "litellm_params": { + "metadata": {"tags": ["caller-tag"]}, + "litellm_metadata": {"model_group": GROUP}, + "model_info": {"id": deployment_id}, + }, + "litellm_call_id": "call-1", + } + + +def test_a_request_routed_through_litellm_metadata_counts_while_in_flight() -> None: + shared: Final = SharedRedisCounters() + calling_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + calling_worker.log_pre_api_call(model="m", messages=[], kwargs=_routed_through_litellm_metadata("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + calling_worker.log_success_event(_routed_through_litellm_metadata("dep-a"), None, None, None) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_closed_realtime_session_leaves_its_deployment_idle() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + session_event: Final = {**_routed_through_litellm_metadata("dep-a"), "call_type": CallTypes.arealtime.value} + + for _ in range(5): + worker.log_pre_api_call(model="m", messages=[], kwargs=session_event) + worker.log_success_event(session_event, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that stuck count an hour after the key was created, which only works while nothing writes the TTL @@ -259,12 +339,16 @@ ROUTER_GROUP: Final = "least-busy-router-group" ROUTER_DEPLOYMENT_IDS: Final = ("1", "2", "3") -def _least_busy_router() -> Router: +def _least_busy_router(group: str = ROUTER_GROUP, *, use_in_pass_through: bool = False) -> Router: return Router( model_list=[ { - "model_name": ROUTER_GROUP, - "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": f"key-{deployment_id}"}, + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": f"key-{deployment_id}", + "use_in_pass_through": use_in_pass_through, + }, "model_info": {"id": deployment_id}, } for deployment_id in ROUTER_DEPLOYMENT_IDS @@ -391,3 +475,575 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts assert response._hidden_params["model_id"] == "1" assert route.calls.last.request.headers["authorization"] == "Bearer key-1" assert _router_counts(router) == seeded + + +@pytest.mark.asyncio +async def test_router_spreads_requests_that_start_together( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + burst: Final = 2 * len(ROUTER_DEPLOYMENT_IDS) + all_started: Final = asyncio.Barrier(burst) + + async def completion_after_all_started(_: httpx.Request) -> httpx.Response: + await all_started.wait() + return httpx.Response( + 200, + json={ + "id": "chatcmpl-least-busy", + "object": "chat.completion", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=completion_after_all_started + ) + router: Final = _least_busy_router() + + await asyncio.gather( + *(router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(burst)) + ) + + one_per_deployment: Final = tuple(f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS) + assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted(one_per_deployment * 2) + + +class _ManualClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +def test_a_reservation_that_never_reaches_the_api_call_expires() -> None: + clock: Final = _ManualClock() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=clock) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = 1.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + clock.now = 2.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.5 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_a_fresh_reservation_outlives_the_ttl_of_the_first_one_on_its_deployment() -> None: + clock: Final = _ManualClock() + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + worker: Final = LeastBusyLoggingHandler(router_cache=cache, clock=clock) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(kwargs=_call_kwargs("dep-a"), response_obj=None, start_time=None, end_time=None) + clock.now = SELECTION_RESERVATION_TTL_SECONDS - 0.1 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.1 + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None: + worker: Final = _worker(None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None: + cache: Final = DualCache() + reservations: Final = InMemoryCache() + picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +class _StoreWhereThreadsReadBeforeWriting: + def __init__(self) -> None: + self._store: Final = InMemoryCache() + self._first_reads: Final = threading.Barrier(2, timeout=0.5) + self._thread_state: Final = threading.local() + self._meeting: Final = threading.Event() + + @contextlib.contextmanager + def first_reads_meet(self) -> Generator[None, None, None]: + self._meeting.set() + yield + self._meeting.clear() + + def get_cache(self, key: str) -> object: + value: Final = self._store.get_cache(key) + if self._meeting.is_set() and not getattr(self._thread_state, "has_read", False): + self._thread_state.has_read = True + with contextlib.suppress(threading.BrokenBarrierError): + self._first_reads.wait() + return value + + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: + self._store.set_cache(key, value, ttl=ttl) + + def delete_cache(self, key: str) -> None: + self._store.delete_cache(key) + + +def test_two_threads_picking_at_once_reserve_different_deployments() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2))) + + assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"] + + +def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + reserved: Final = pool.submit(worker.get_available_deployments, GROUP, HEALTHY, True) + started: Final = pool.submit(worker.log_pre_api_call, "m", [], _call_kwargs("dep-a")) + assert reserved.result() is DEPLOYMENT_A + started.result() + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + + +class _SlowLocalRead(InMemoryCache): + """Answers a count read with the counts from when the read started, once another thread has begun its request.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + self.other_thread: threading.Thread | None = None + + def batch_get_cache(self, keys: list[str], **kwargs: object) -> list[object]: + counts_when_asked: Final = super().batch_get_cache(keys, **kwargs) + began: Final = threading.Event() + + def begin_then_run() -> None: + began.set() + self.while_reading() + + self.other_thread = threading.Thread(target=begin_then_run) + self.other_thread.start() + began.wait() + self.other_thread.join(timeout=0.2) + return counts_when_asked + + +def test_a_request_that_starts_on_another_thread_while_counts_are_read_still_counts_on_its_deployment() -> None: + local: Final = _SlowLocalRead() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(in_memory_cache=local)) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + local.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + picked_while_the_other_thread_started: Final = worker.get_available_deployments(GROUP, HEALTHY, reserve=True) + if local.other_thread is not None: + local.other_thread.join() + + assert picked_while_the_other_thread_started is DEPLOYMENT_B + + +def _router_pick(router: Router, group: str = ROUTER_GROUP) -> str: + deployment: Final = router.get_available_deployment(group, request_kwargs={"metadata": {}}) + return deployment["model_info"]["id"] + + +def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None: + router: Final = _least_busy_router() + + with picks_reserve(True): + first: Final = _router_pick(router) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) + + assert _router_pick(router) != first + + +@pytest.mark.asyncio +async def test_looking_up_a_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router() + + sync_picks: Final = [_router_pick(router) for _ in range(3)] + async_picks: Final = [ + (await router.async_get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}))["model_info"]["id"] + for _ in range(3) + ] + + assert sync_picks + async_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6 + + +@pytest.mark.asyncio +async def test_looking_up_a_pass_through_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router(use_in_pass_through=True) + + selected: Final = [ + (await router.select(ROUTER_GROUP, {"metadata": {}}, kind="pass_through"))["model_info"]["id"] for _ in range(2) + ] + async_picks: Final = [ + (await router.async_get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + sync_picks: Final = [ + router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"]["id"] + for _ in range(2) + ] + + assert selected + async_picks + sync_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6 + + +@pytest.mark.asyncio +async def test_a_cache_hit_releases_the_pick_it_reserved() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + await worker.async_log_success_event({**_call_kwargs("dep-b"), "cache_hit": True}, None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + worker.log_success_event({**_call_kwargs("dep-a", "cached-call"), "cache_hit": True}, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event]) +def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick( + outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None], +) -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + outcome(worker, _call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None: + worker: Final = _worker(None) + rejected: Final = {**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} + + worker.log_pre_api_call(model="m", messages=[], kwargs=rejected) + worker.log_failure_event(rejected, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments("tier-group", HEALTHY) is DEPLOYMENT_A + + +def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pick() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "busy-call")) + first_attempt_started_at: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + first_attempt: Final = {**_call_kwargs("dep-a", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_pre_api_call(model="m", messages=[], kwargs=first_attempt) + worker.log_failure_event(first_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + second_attempt: Final = {**_call_kwargs("dep-b", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_failure_event(second_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_ending_one_request_frees_only_the_picks_it_never_started() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="ended-call") is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="running-call") is DEPLOYMENT_A + + release_unstarted_picks(reservations, "ended-call") + + assert [worker.get_available_deployments(GROUP, HEALTHY, reserve=True) for _ in range(2)] == [ + DEPLOYMENT_A, + DEPLOYMENT_B, + ] + + +def test_a_call_without_a_call_id_turns_its_own_pick_into_its_in_flight_count() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "attempt-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_another_call_starting_leaves_the_pick_of_a_call_without_a_call_id_reserved() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + other_call: Final = contextvars.copy_context() + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + other_call.run(worker.log_pre_api_call, model="m", messages=[], kwargs=_call_kwargs("dep-a", "other-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "looked-up-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def _fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "unreachable-group", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-unreachable"}, + "model_info": {"id": "unreachable"}, + }, + *( + { + "model_name": ROUTER_GROUP, + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": f"key-{deployment_id}"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ROUTER_DEPLOYMENT_IDS + ), + ], + routing_strategy="least-busy", + num_retries=0, + fallbacks=[{"unreachable-group": [ROUTER_GROUP]}], + ) + + +@pytest.mark.asyncio +async def test_a_fallback_hop_rejected_before_its_provider_frees_its_pick_when_every_hop_shares_one_logging_object( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(500, json={"error": {"message": "upstream down"}}) + ) + router: Final = _fallback_router() + logging_obj, request = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime(2026, 1, 1), + model="unreachable-group", + messages=[{"role": "user", "content": "hi"}], + logit_bias={"1": 1}, + litellm_call_id="fallback-call", + ) + + with pytest.raises(litellm.InternalServerError): + await router.acompletion(**request, litellm_logging_obj=logging_obj) + await _drain_logging() + + assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-unreachable"} + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + + +def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None: + worker: Final = _worker(None) + attempt: Final = _call_kwargs("dep-a", "retried-call") + + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_failure_event(attempt, None, None, None) + worker.log_success_event(attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +async def _drain_logging() -> None: + for _ in range(10): + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + + +def _called_deployment(route: respx.Route) -> str: + return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-") + + +_BATCH: Final = { + "id": "batch_least_busy", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-least-busy", + "completion_window": "24h", + "status": "validating", + "created_at": 1, +} +_REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "invalid_request_error"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)]) +async def test_a_batch_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.acreate_batch( + model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy" + ) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +_VECTOR_STORE: Final = { + "id": "vs_least_busy", + "object": "vector_store", + "created_at": 1, + "name": "least-busy", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", +} +_REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}} +_CALL_IDENTITIES: Final = ({"litellm_call_id": "least-busy-call"}, {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +@pytest.mark.parametrize( + "provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)] +) +async def test_a_vector_store_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, + call_identity: dict[str, str], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.get(url__startswith="https://api.openai.com/v1/fine_tuning/jobs").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False}) + ) + router: Final = _least_busy_router() + + await router.alist_fine_tuning_jobs(model=ROUTER_GROUP, **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +TEXT_GROUP: Final = "openai/least-busy-text-group" + + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_text_completion_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "cmpl-least-busy", + "object": "text_completion", + "created": 1, + "model": "least-busy-text-group", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + ) + router: Final = _least_busy_router(TEXT_GROUP) + + router.text_completion(model=TEXT_GROUP, prompt="Say ok", **call_identity) + + assert _router_pick(router, TEXT_GROUP) == _called_deployment(route) + + + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick( + call_identity: dict[str, str], +) -> None: + router: Final = _least_busy_router() + + with pytest.raises(litellm.InternalServerError): + router.responses(model=ROUTER_GROUP, input="hi", api_base="not-a-url", **call_identity) + + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + +@pytest.mark.asyncio +async def test_a_batch_cancel_does_not_keep_its_deployment_busy( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches/batch_least_busy/cancel").mock( + return_value=httpx.Response(200, json={**_BATCH, "status": "cancelling"}) + ) + router: Final = _least_busy_router() + + await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy") + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() -> None: + router: Final = _least_busy_router() + + picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}}) + + assert _router_pick(router) == picked["model_info"]["id"] diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index afcd1e2a7ff..951182051c4 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -21,6 +21,7 @@ from litellm import Router from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext +from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -482,6 +483,38 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 +def _filtered_model_pick(router: Router) -> str: + deployment: Final = router.get_available_deployment( + "filtered-model", request_kwargs=_proxy_shaped_request(model="filtered-model") + ) + return deployment["model_info"]["id"] + + +def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = _build_router( + routing_strategy="least-busy", + routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}], + ) + + request: Final = _proxy_shaped_request(model="filtered-model") + with picks_reserve(True): + first: Final = router.get_available_deployment("filtered-model", request_kwargs=request)["model_info"]["id"] + kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}, + "litellm_call_id": request["litellm_call_id"], + } + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert _filtered_model_pick(router) == first + + def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): """ Least-busy hangs its counting off litellm's global callback lists, and those lists keep one @@ -614,6 +647,7 @@ async def test_select_deployment_async_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -628,6 +662,7 @@ def test_select_deployment_sync_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -662,6 +697,7 @@ def test_select_deployment_sync_does_not_dispatch_cost_based_routing(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -708,6 +744,7 @@ async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): messages=[{"role": "user", "content": "hi"}], input=None, request_kwargs={}, + reserve=False, ) assert v1_spy.called, "async dispatch must route v1 strategy through sync method" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index c16c1c62202..f5962cb4933 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -30102,11 +30102,21 @@ def test_reset_custom_routing_strategy(): assert "get_available_deployment" not in router._core.__dict__ assert "_core_async_get_available_deployment" not in router._core.__dict__ - assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment + assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment router._reset_custom_routing_strategy() +@pytest.mark.asyncio +async def test_select_picks_through_a_custom_routing_strategy() -> None: + router: Final = _create_custom_routing_router() + router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router)) + + picked: Final = await router.select("azure-model", {"metadata": {}}) + + assert picked["model_info"]["id"] == "very-special-endpoint" + + @pytest.mark.asyncio async def test_azure_router_acompletion_passes_deployment_credentials_for_plain_and_stream( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5ad6eb819ef..3bbb69413d9 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -844,6 +844,34 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b await admit_turn() +@pytest.mark.asyncio +async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): + router: Final = Router( + model_list=[ + { + "model_name": "ws-least-busy", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("ws-least-busy-a", "ws-least-busy-b") + ], + routing_strategy="least-busy", + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + for _ in range(3): + await admit_turn() + + picks: Final = [ + (await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + + assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"] + + @pytest.mark.asyncio async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): router = Router(