diff --git a/litellm/constants.py b/litellm/constants.py index e07d727973d..cfee777f36a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from typing import Final, Literal from litellm.litellm_core_utils.env_utils import ( get_env_float, + get_env_float_in_range, get_env_float_or_none, get_env_int, get_env_int_in_range, @@ -149,6 +150,9 @@ DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float( "DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5 ) # default cooldown a deployment if 50% of requests fail in a given minute DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096) +NO_MAX_TOKENS_TPM_RESERVATION_FRACTION: Final = get_env_float_in_range( + "NO_MAX_TOKENS_TPM_RESERVATION_FRACTION", 0.25, 0.01, 1.0 +) DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3) DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1) DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 28211b3dab2..12425c45b01 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -9,6 +9,7 @@ ## LiteLLM versions of the OpenAI Exception Types +import datetime import enum from collections.abc import Mapping, Sequence from typing import Any, Final @@ -115,6 +116,54 @@ def validate_rate_limit_type(value: object) -> str | None: return None +def _upstream_body(response: object) -> bytes: + if not isinstance(response, httpx.Response): + return b"" + try: + return response.content + except httpx.ResponseNotRead: + return b"" + + +def _request_without_credentials(upstream: object, fallback_url: str) -> httpx.Request: + if not isinstance(upstream, httpx.Response): + return httpx.Request(method="POST", url=fallback_url) + try: + upstream_request: Final = upstream.request + except RuntimeError: + return httpx.Request(method="POST", url=fallback_url) + return httpx.Request( + method=upstream_request.method, + url=upstream_request.url.copy_with(userinfo=b"", query=None, fragment=None), + ) + + +def _upstream_elapsed(upstream: object) -> datetime.timedelta | None: + if not isinstance(upstream, httpx.Response): + return None + try: + return upstream.elapsed + except RuntimeError: + return None + + +def error_response_without_credentials(status_code: int, upstream: object, fallback_url: str) -> httpx.Response: + response: Final = httpx.Response( + status_code=status_code, + stream=httpx.ByteStream(_upstream_body(upstream)), + request=_request_without_credentials(upstream, fallback_url), + ) + _ = response.read() + response.headers = httpx.Headers(upstream.headers if isinstance(upstream, httpx.Response) else None) + elapsed: Final = _upstream_elapsed(upstream) + if elapsed is not None: + response.elapsed = elapsed + return response + + +_NO_UPSTREAM_REQUEST_URL: Final = "https://cloud.google.com/vertex-ai/" + + _MINIMAL_ERROR_RESPONSE: httpx.Response | None = None @@ -519,20 +568,11 @@ class RateLimitError(openai.RateLimitError): # headers stay reachable on `e.response.headers` for callers that # explicitly want them; only the proxy-supplied `headers=` kwarg # makes it onto `self.headers`. - _response_headers: Final = getattr(response, "headers", None) if response is not None else None self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Mirrors FastAPI HTTPException.detail so the same instance can be # serialized through both the ProxyException and HTTPException paths. self.detail = detail if detail is not None else self.message - self.response = httpx.Response( - status_code=429, - headers=_response_headers, - content=response.content if response is not None else None, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(429, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -757,15 +797,7 @@ class ServiceUnavailableError(openai.APIStatusError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, body=None ) # Call the base class constructor with the parameters it needs @@ -808,15 +840,7 @@ class BadGatewayError(openai.APIStatusError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -862,15 +886,7 @@ class InternalServerError(openai.InternalServerError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -908,6 +924,7 @@ class APIError(openai.APIError): max_retries: int | None = None, num_retries: int | None = None, body: object | None = None, + response: httpx.Response | None = None, ): self.status_code = status_code self.message = f"litellm.APIError: {message}" @@ -916,6 +933,8 @@ class APIError(openai.APIError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + if response is not None: + self.response = response if request is None: request = httpx.Request(method="POST", url="https://api.openai.com/v1") super().__init__(self.message, request=request, body=body) diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py index 6883af98bbd..462b0b3df47 100644 --- a/litellm/litellm_core_utils/env_utils.py +++ b/litellm/litellm_core_utils/env_utils.py @@ -71,7 +71,15 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) Values outside the range fall back to the default and warn, so a misconfigured knob can neither crash the caller nor silently change the meaning of what it computes. """ - value: Final = get_env_int(env_var, default) + return _within_range(env_var, get_env_int(env_var, default), default, minimum, maximum) + + +def get_env_float_in_range(env_var: str, default: float, minimum: float, maximum: float) -> float: + """Float counterpart of `get_env_int_in_range`.""" + return _within_range(env_var, get_env_float(env_var, default), default, minimum, maximum) + + +def _within_range(env_var: str, value: _NumberT, default: _NumberT, minimum: _NumberT, maximum: _NumberT) -> _NumberT: if minimum <= value <= maximum: return value logging.getLogger("LiteLLM").warning( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 64a95eb5f28..5a27f5446ad 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -37,6 +37,7 @@ from ..exceptions import ( ServiceUnavailableError, Timeout, UnprocessableEntityError, + error_response_without_credentials, ) @@ -305,6 +306,16 @@ def _litellm_proxy_response( ) +def _without_credentials(response: httpx.Response | None) -> httpx.Response | None: + if not isinstance(response, httpx.Response): + return response + return error_response_without_credentials(response.status_code, response, "https://docs.litellm.ai/docs") + + +def _upstream_reply(response: httpx.Response | None) -> httpx.Response | None: + return response if response is not None and (response.headers or response.content) else None + + def _map_openai_exception( *, model: str, @@ -315,7 +326,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: - response: Final = _litellm_proxy_response(original_exception, custom_llm_provider) + response: Final = _without_credentials(_litellm_proxy_response(original_exception, custom_llm_provider)) if (cyber_policy_body := _openai_cyber_policy_body(original_exception)) is not None: raise ContentPolicyViolationError( message=f"ContentPolicyViolationError: {exception_provider} - {original_exception}", @@ -378,6 +389,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif ( ("invalid_request_error" in error_str and "content_policy_violation" in error_str) @@ -449,6 +461,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, body=getattr(original_exception, "body", None), + response=_upstream_reply(response), ) elif "Request too large" in error_str: raise RateLimitError( @@ -479,6 +492,7 @@ def _map_openai_exception( model=model, request=_request, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif hasattr(original_exception, "status_code"): if original_exception.status_code == 400: @@ -512,6 +526,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -563,6 +578,7 @@ def _map_openai_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, + response=_upstream_reply(response), ) else: raise APIError( @@ -572,6 +588,7 @@ def _map_openai_exception( model=model, request=getattr(original_exception, "request", None), litellm_debug_info=extra_information, + response=_upstream_reply(response), ) else: # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors @@ -2019,6 +2036,7 @@ def _map_azure_exception( exception_provider: str, extra_information: str, ) -> None: + response: Final = _without_credentials(getattr(original_exception, "response", None)) message = get_error_message(error_obj=original_exception) if message is None: if hasattr(original_exception, "message"): @@ -2054,7 +2072,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif "DeploymentNotFound" in error_str: raise NotFoundError( @@ -2062,7 +2080,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif azure_error_code == "content_policy_violation" or ExceptionCheckers.is_azure_content_policy_violation_error( error_str @@ -2076,6 +2094,7 @@ def _map_azure_exception( model=model, extra_information=extra_information, original_exception=original_exception, + response=response, ) elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str: helpful_message: Final = ( @@ -2094,7 +2113,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif getattr(original_exception, "status_code", None) in (None, 400) and ( @@ -2106,7 +2125,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400): @@ -2115,7 +2134,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif "The api_key client option must be set either by passing api_key to the client or by setting" in error_str: @@ -2124,7 +2143,7 @@ def _map_azure_exception( llm_provider=custom_llm_provider, model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif "Connection error" in error_str: raise APIConnectionError( @@ -2140,7 +2159,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif original_exception.status_code == 401: @@ -2149,7 +2168,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 404: raise NotFoundError( @@ -2157,7 +2176,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 408: raise Timeout( @@ -2165,6 +2184,7 @@ def _map_azure_exception( model=model, litellm_debug_info=extra_information, llm_provider="azure", + response=_upstream_reply(response), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -2172,7 +2192,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 429: raise RateLimitError( @@ -2180,7 +2200,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 502: raise BadGatewayError( @@ -2188,7 +2208,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 503: raise ServiceUnavailableError( @@ -2196,7 +2216,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 504: # gateway timeout error raise Timeout( @@ -2205,6 +2225,7 @@ def _map_azure_exception( litellm_debug_info=extra_information, llm_provider="azure", exception_status_code=original_exception.status_code, + response=_upstream_reply(response), ) else: raise APIError( @@ -2214,6 +2235,7 @@ def _map_azure_exception( litellm_debug_info=extra_information, model=model, request=httpx.Request(method="POST", url="https://openai.com/"), + response=_upstream_reply(response), ) else: # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors @@ -2589,17 +2611,20 @@ def exception_type( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply( + _without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)) + ), ) if ( custom_llm_provider == "litellm_proxy" ): # handle special case where calling litellm proxy + exception str contains error message extract_and_raise_litellm_exception( - response=_litellm_proxy_response(mappable_exception, custom_llm_provider), + response=_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)), error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, - body=getattr(original_exception, "body", None), + body=getattr(mappable_exception, "body", None), ) if ( custom_llm_provider == "openai" diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index a1ba8de6d7f..69ffac5d6f9 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -42,6 +42,7 @@ from ..openai.common_utils import ( build_output_token_limit_response, is_output_token_limit_error, model_dump_provider_response, + upstream_error_response, ) from .common_utils import ( AzureOpenAIError, @@ -399,6 +400,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): raise AzureOpenAIError( status_code=status_code, message=str(e), + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -662,6 +664,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): raise AzureOpenAIError( status_code=status_code, message=message, + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -855,7 +858,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) error_text = error_response.text - raise AzureOpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def make_async_azure_httpx_request( self, @@ -1190,7 +1198,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: error_code: Final[object] = getattr(e, "status_code", None) if error_code is not None: - raise AzureOpenAIError(status_code=error_code, message=str(e)) + raise AzureOpenAIError(status_code=error_code, message=str(e), response=upstream_error_response(e)) else: raise AzureOpenAIError(status_code=500, message=str(e)) diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 34e51eaf1ec..c2744e7f750 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResponse +from ...openai.common_utils import upstream_error_response from ...openai.completion.transformation import OpenAITextCompletionConfig from ..common_utils import AzureOpenAIError, BaseAzureLLM @@ -183,7 +184,12 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) async def acompletion( self, @@ -244,7 +250,12 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) def streaming( self, @@ -355,4 +366,9 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index f7146a36384..6eb5161f7e4 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -1,5 +1,7 @@ from typing import Final +import httpx + from litellm.exceptions import ContentPolicyViolationError @@ -14,6 +16,7 @@ class AzureOpenAIExceptionMapping: model: str, extra_information: str, original_exception: Exception, + response: httpx.Response | None, ) -> ContentPolicyViolationError: """ Create a content policy violation error @@ -39,7 +42,7 @@ class AzureOpenAIExceptionMapping: llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, provider_specific_fields={ # Preserve legacy key for backward compatibility. "innererror": inner_error, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 6c416124d1a..aa14a985afe 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from aiohttp import ClientSession import litellm +from litellm.exceptions import error_response_without_credentials from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model @@ -148,6 +149,14 @@ class OpenAIError(BaseLLMException): ) +def upstream_error_response(error: Exception) -> httpx.Response | None: + if not isinstance(error, (openai.APIStatusError, httpx.HTTPStatusError)): + return None + return error_response_without_credentials( + error.response.status_code, error.response, "https://docs.litellm.ai/docs" + ) + + _DECODED_JSON: Final = TypeAdapter(object) diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 278427e8d77..8bfdd536f90 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -10,7 +10,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse from litellm.utils import ProviderConfigManager -from ..common_utils import BaseOpenAILLM, OpenAIError +from ..common_utils import BaseOpenAILLM, OpenAIError, upstream_error_response from ..workload_identity import build_async_openai_client, build_openai_client from .transformation import OpenAITextCompletionConfig @@ -164,7 +164,12 @@ class OpenAITextCompletion(BaseLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def acompletion( self, @@ -220,7 +225,12 @@ class OpenAITextCompletion(BaseLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def streaming( self, @@ -261,7 +271,12 @@ class OpenAITextCompletion(BaseLLM): error_response: object = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) streamwrapper: Final = CustomStreamWrapper( completion_stream=response, model=model, diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 01f01d3de1d..504dfb50015 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj -from ..common_utils import OpenAIError +from ..common_utils import OpenAIError, upstream_error_response from ..workload_identity import build_async_openai_client, build_openai_client @@ -129,7 +129,12 @@ class OpenAIImageVariationsHandler: error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def image_variations( self, @@ -245,4 +250,9 @@ class OpenAIImageVariationsHandler: error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index af2e85c5f20..e0a23d6c99f 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -64,6 +64,7 @@ from .common_utils import ( is_openai_backed_api_base, is_output_token_limit_error, model_dump_provider_response, + upstream_error_response, ) from .workload_identity import ( build_async_openai_client, @@ -983,6 +984,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=status_code, message=error_text, + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -1110,6 +1112,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=status_code, message=message, + response=upstream_error_response(e), headers=error_headers, body=exception_body, ) @@ -1282,6 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=getattr(e, "status_code", 500), message=str(e), + response=upstream_error_response(e), headers=error_headers, body=exception_body, ) @@ -1409,7 +1413,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def embedding( self, @@ -1499,7 +1508,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def aimage_generation( self, @@ -1666,7 +1680,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): original_response=str(e), ) if hasattr(e, "status_code"): - raise OpenAIError(status_code=getattr(e, "status_code", 500), message=str(e)) + raise OpenAIError( + status_code=getattr(e, "status_code", 500), + message=str(e), + response=upstream_error_response(e), + ) else: raise OpenAIError(status_code=500, message=str(e)) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 96a04c96f44..0799a298942 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6014,11 +6014,14 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6313,16 +6316,19 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-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/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -11303,16 +11309,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "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", "/v1/batch", @@ -11411,11 +11420,14 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -11513,16 +11525,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "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", "/v1/batch", diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 404aff84994..14b08a09116 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -42,7 +42,11 @@ from litellm.caching.redis_batch import ( active_request_redis_batch, ) from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import ( + DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + NO_MAX_TOKENS_TPM_RESERVATION_FRACTION, +) from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens @@ -508,11 +512,6 @@ REDIS_NODE_HASHTAG_NAME: Final = "all_keys" # *some* output budget; these define that fallback estimate. DEFAULT_MAX_TOKENS_ESTIMATE: Final = 4096 DEFAULT_CHARS_PER_TOKEN: Final = 4 -# Fraction of the available output budget reserved as the upfront floor when -# the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE -# (baseline floor) and to the smallest configured TPM limit (capped floor for -# small per-tenant TPM caps). -_TPM_FLOOR_FRACTION: Final = 4 # Both embeddings and the Responses API put their prompt in data["input"], # but only embeddings have no output tokens. Every "is this an embedding" # check on data["input"] must exclude these call types, or a Responses call @@ -1003,10 +1002,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): per-tenant cap can't be tripped by the floor alone. Returns the baseline floor when no limit is provided. """ - baseline: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + baseline: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) if min_configured_tpm_limit is None: return baseline - return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION)) + return min(baseline, max(1, int(min_configured_tpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION))) @staticmethod def _is_embedding_request(data: object, call_type: str | None) -> bool: @@ -1138,7 +1137,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if call_type in RESPONSES_API_CALL_TYPES else base_capped_floor ) - baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + baseline_floor: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) is_embedding: Final = PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) if ( capped_floor >= baseline_floor diff --git a/litellm/router.py b/litellm/router.py index 3a423b79e49..20af2e5649e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -377,6 +377,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor EncryptedContentAffinityCheck, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_io_token_refund_failure, build_io_token_rate_limit_headers, deployment_has_io_token_limits, refund_stale_reservation_before_retry, @@ -9500,7 +9501,7 @@ class Router(DeprecatedRouterMembers): raise original_exception - async def make_call(self, original_function: Any, *args, **kwargs): + async def make_call(self, original_function: Any, *args: object, **kwargs: object): """ Handler for making a call to the .completion()/.embeddings()/etc. functions. """ @@ -9509,17 +9510,26 @@ class Router(DeprecatedRouterMembers): endpoint_function: Final = ( generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) ) - with endpoint_selection( - getattr(endpoint_function, "__name__", None), - getattr(endpoint_function, "__module__", None), - model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), - ): - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response + try: + with endpoint_selection( + getattr(endpoint_function, "__name__", None), + getattr(endpoint_function, "__module__", None), + model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), + ): + response = original_function(*args, **kwargs) + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + response = await response + except Exception: + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self.cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) + ) + raise await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS - response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) + response = await self.set_response_headers( + response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs + ) return response diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 0b3d48c9e1d..7580c7dadd3 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -1,7 +1,8 @@ """ Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits. -- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM +- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM; a request + without max_tokens reserves the model's max output, capped at a fraction of the OTPM limit - Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens - Cached prompt-read tokens are excluded from ITPM post-call accounting @@ -22,6 +23,7 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET @@ -127,10 +129,15 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | return itpm_key, otpm_key -def _tokenizer_model(deployment: object) -> str: +def _served_model(deployment: object) -> str: params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None served: Final = params.get("model") if is_str_keyed_mapping(params) else None - return deployment_openai_model(deployment) or (served if isinstance(served, str) else "") + return served if isinstance(served, str) else "" + + +def _tokenizer_and_served_models(deployment: object) -> tuple[str, str]: + served: Final = _served_model(deployment) + return deployment_openai_model(deployment) or served, served def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int: @@ -160,24 +167,61 @@ def _model_max_output_tokens(model_name: str) -> int | None: return None -def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict) -> int: - if request_kwargs: - # An explicit max_tokens=0 must be honored, not treated as absent and - # replaced by the model default. - explicit = request_kwargs.get("max_tokens") - if explicit is None: - explicit = request_kwargs.get("max_completion_tokens") - if explicit is None: - explicit = request_kwargs.get("max_output_tokens") - if explicit is not None: - return max(0, int(explicit)) +_CHAT_OUTPUT_CAP_KEYS: Final = ("max_tokens", "max_completion_tokens") +_RESPONSES_OUTPUT_CAP_KEYS: Final = ("max_output_tokens",) +_OUTPUT_CAP_KEYS: Final = _CHAT_OUTPUT_CAP_KEYS + _RESPONSES_OUTPUT_CAP_KEYS - model_name: Final = (deployment.get("litellm_params") or {}).get("model") - if model_name: - model_max: Final = _model_max_output_tokens(model_name) - if model_max is not None: - return model_max - return 4096 + +def _explicit_max_tokens( + request_kwargs: Mapping[str, object] | None, cap_keys: tuple[str, ...] = _OUTPUT_CAP_KEYS +) -> int | None: + if not request_kwargs: + return None + explicit: Final = next( + (request_kwargs[key] for key in cap_keys if request_kwargs.get(key) is not None), + None, + ) + match explicit: + case int() | float() | str(): + return max(0, int(explicit)) + case _: + return None + + +def _sent_deployment_cap_keys(request_kwargs: Mapping[str, object] | None) -> tuple[str, ...]: + is_responses_request: Final = ( + request_kwargs is not None and "input" in request_kwargs and "messages" not in request_kwargs + ) + return _RESPONSES_OUTPUT_CAP_KEYS if is_responses_request else _CHAT_OUTPUT_CAP_KEYS + + +def _deployment_output_cap(request_kwargs: Mapping[str, object] | None, deployment: Mapping[str, object]) -> int | None: + params: Final = deployment.get("litellm_params") + return ( + _explicit_max_tokens(params, _sent_deployment_cap_keys(request_kwargs)) + if is_str_keyed_mapping(params) + else None + ) + + +def _default_otpm_reservation(model_name: str, otpm_limit: int) -> int: + model_max: Final = _model_max_output_tokens(model_name) if model_name else None + share_of_limit: Final = max(1, int(otpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)) + return min(4096 if model_max is None else model_max, share_of_limit) + + +def _otpm_reservation( + request_kwargs: Mapping[str, object] | None, + deployment: Mapping[str, object], + model_name: str, + otpm_limit: int, +) -> int: + from_request: Final = _explicit_max_tokens(request_kwargs) + if from_request is not None: + return from_request + default: Final = _default_otpm_reservation(model_name, otpm_limit) + from_deployment: Final = _deployment_output_cap(request_kwargs, deployment) + return default if from_deployment is None else min(from_deployment, default) def _get_usage_tokens(usage: object) -> tuple[int, int, int]: @@ -210,6 +254,12 @@ def _extract_response_usage(response_obj: object) -> object: return getattr(response_obj, "usage", None) +def _openai_shaped_usage(usage: object) -> object: + if isinstance(usage, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage): + return litellm.AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) + return usage + + def _usage_is_present(usage: object) -> bool: """ True only if usage carries an actual input/output breakdown. @@ -238,7 +288,7 @@ def _resolve_reconcile_usage_tokens( fields. When usage cannot be resolved, return ``usage_resolved=False`` so callers keep the pre-call reservation instead of refunding it as zero usage. """ - usage: Final = _extract_response_usage(response_obj) + usage: Final = _openai_shaped_usage(_extract_response_usage(response_obj)) if _usage_is_present(usage): prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage) return max(0, prompt_tokens - cached_tokens), completion_tokens, True @@ -261,6 +311,21 @@ def _resolve_reconcile_usage_tokens( return 0, 0, False +def _key_for_this_minute(reserved_key: str) -> str: + counter, _, _ = reserved_key.rpartition(":") + return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" + + +def _reconcile_increments( + reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int +) -> tuple[tuple[str, int], ...]: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation + deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved)) + return tuple( + (key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta + ) + + def _stash_reservation_in_metadata( request_kwargs: dict[str, object] | None, *, @@ -311,7 +376,9 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, return () top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata")) litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params")) - from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),) + from_params: Final = ( + () if litellm_params is None else (litellm_params.get("metadata"), litellm_params.get("litellm_metadata")) + ) standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object")) from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),) return top_level + from_params + from_logging_object @@ -433,9 +500,11 @@ def io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = _tokenizer_model(deployment) + _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) - max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) + otpm_to_reserve: Final = ( + _otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0 + ) dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") @@ -458,7 +527,7 @@ def io_token_pre_call_check( ) if otpm_limit is not None: - otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + otpm_reserved = otpm_to_reserve try: _sync_increment_with_rollback( dual_cache, @@ -497,9 +566,11 @@ async def async_io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = _tokenizer_model(deployment) + _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) - max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) + otpm_to_reserve: Final = ( + _otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0 + ) dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") @@ -523,7 +594,7 @@ async def async_io_token_pre_call_check( ) if otpm_limit is not None: - otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + otpm_reserved = otpm_to_reserve try: await _increment_with_rollback( dual_cache, @@ -564,7 +635,8 @@ def io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _read_reservation_from_kwargs(request_kwargs) + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return @@ -572,23 +644,8 @@ def io_token_reconcile_success( try: if usage_resolved: - if itpm_key is not None: - itpm_delta: Final = billable_input - itpm_reserved - if itpm_delta != 0: - dual_cache.increment_cache( - key=itpm_key, - value=itpm_delta, - ttl=RoutingArgsTTL, - ) - - if otpm_key is not None: - otpm_delta: Final = completion_tokens - otpm_reserved - if otpm_delta != 0: - dual_cache.increment_cache( - key=otpm_key, - value=otpm_delta, - ttl=RoutingArgsTTL, - ) + for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", @@ -618,39 +675,23 @@ async def async_io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _read_reservation_from_kwargs(request_kwargs) + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) - # Reconcile against the exact key that held the reservation (which encodes - # the reservation's minute), not a key recomputed at response time. This - # tracks actual usage even when the pre-call estimate was 0, and avoids a - # minute-boundary mismatch for calls that span into the next minute. Always - # clear the stash afterwards (even if an increment throws) so a retry or a - # duplicate success event can't re-process it. + # A refund goes back to the minute that held the reservation, while usage above the reservation is charged + # to the minute the response finished in, so a call that runs past its minute still counts against a live + # window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success + # event can't re-process it. try: if usage_resolved: - if itpm_key is not None: - itpm_delta: Final = billable_input - itpm_reserved - if itpm_delta != 0: - await dual_cache.async_increment_cache( - key=itpm_key, - value=itpm_delta, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) - - if otpm_key is not None: - otpm_delta: Final = completion_tokens - otpm_reserved - if otpm_delta != 0: - await dual_cache.async_increment_cache( - key=otpm_key, - value=otpm_delta, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + await dual_cache.async_increment_cache( + key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span + ) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index f4b41211c6e..f41d19cd62c 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -26,7 +26,6 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_RESERVED_KEY, async_io_token_pre_call_check, async_io_token_reconcile_success, async_io_token_refund_failure, @@ -485,24 +484,12 @@ class ModelRateLimitingCheck(CustomLogger): try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - - # IO token reconciliation works purely from the cache keys stashed in - # kwargs/metadata, so it must run before the model_id guard below - # (which only the TPM-tracking path needs). Otherwise a request whose - # standard_logging_object lacks model_id would never return its - # reservation, leaving the counter elevated until the TTL expires. - slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} - kwargs_metadata: Final = kwargs.get("metadata") or {} - if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: - await async_io_token_reconcile_success( - self.dual_cache, - kwargs, - response_obj, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - # Fall through: a deployment can also configure tpm/rpm alongside - # itpm/otpm, and that path's pre-call check reads the tpm_key - # counter tracked below, so it must still be incremented here. + await async_io_token_reconcile_success( + self.dual_cache, + kwargs, + response_obj, + parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), + ) if standard_logging_object is None: return @@ -557,17 +544,7 @@ class ModelRateLimitingCheck(CustomLogger): """ try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} - kwargs_metadata: Final = kwargs.get("metadata") or {} - if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: - io_token_reconcile_success( - self.dual_cache, - kwargs, - response_obj, - ) - # Fall through: a deployment can also configure tpm/rpm alongside - # itpm/otpm, and that path's pre-call check reads the tpm_key - # counter tracked below, so it must still be incremented here. + io_token_reconcile_success(self.dual_cache, kwargs, response_obj) if standard_logging_object is None: return diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 96a04c96f44..0799a298942 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6014,11 +6014,14 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6313,16 +6316,19 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-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/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -11303,16 +11309,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "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", "/v1/batch", @@ -11411,11 +11420,14 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -11513,16 +11525,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "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", "/v1/batch", diff --git a/tests/integration/routing/test_otpm_default_reservation_fraction.py b/tests/integration/routing/test_otpm_default_reservation_fraction.py new file mode 100644 index 00000000000..589a2535a03 --- /dev/null +++ b/tests/integration/routing/test_otpm_default_reservation_fraction.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import dataclasses +import threading +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GROUPS: Final = ("otpm-fraction-first", "otpm-fraction-retry") +OTPM_LIMIT: Final = 1000 +FRACTION: Final = 0.5 +DEFAULT_RESERVATION: Final = int(OTPM_LIMIT * FRACTION) + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + upstream: Wire + gates: SimpleQueue[threading.Event] + + +@dataclass(frozen=True, slots=True) +class Admissions: + held_forwarded_max_tokens: bool + statuses: tuple[int, int, int] + one_minute: bool + + +def _reply_for(gates: SimpleQueue[threading.Event]) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + return dataclasses.replace( + chat_reply("chatcmpl-held", OPENAI_MODEL, "held", stream=True), gate_before_first=gates.get(timeout=60) + ) + return chat_reply("chatcmpl-capped", OPENAI_MODEL, "capped", stream=False) + + return reply + + +def _config(directory: Path, upstream: Wire) -> Path: + return owned_config( + directory, + [ + { + "model_name": group, + "litellm_params": { + "model": OPENAI_MODEL, + "api_base": upstream.url + "/v1", + "api_key": "synthetic-openai-key", + "otpm": OTPM_LIMIT, + }, + "model_info": {"id": group}, + } + for group in GROUPS + ], + litellm_settings={"cache": False}, + router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 0}, + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("otpm-fraction") + gates: Final[SimpleQueue[threading.Event]] = SimpleQueue() + with ( + gateway_from_environment() as gateway, + wire_server(_reply_for(gates)) as upstream, + owned_proxy_process( + gateway, + directory, + {"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION": str(FRACTION)}, + config=_config(directory, upstream), + ) as owned, + ): + yield Rig(owned, upstream, gates) + + +def _chat(rig: Rig, group: str, **params: int | bool) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": group, "messages": [{"role": "user", "content": f"fraction {uuid.uuid4().hex}"}], **params}, + ) + + +def _await_forwarded_chat(upstream: Wire) -> Request: + request: Final = upstream.received.get(timeout=60) + return _await_forwarded_chat(upstream) if is_model_info_probe(request) else request + + +def _utc_minute() -> str: + return datetime.now(UTC).strftime("%Y-%m-%d %H:%M") + + +def _admissions_beside_a_held_stream(rig: Rig, group: str) -> Admissions: + release: Final = threading.Event() + rig.gates.put(release) + started: Final = _utc_minute() + with ThreadPoolExecutor(max_workers=1) as pool: + held: Final = pool.submit(_chat, rig, group, stream=True) + try: + forwarded: Final = JSON_OBJECT.validate_json(_await_forwarded_chat(rig.upstream).body) + over: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION + 1) + fits: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION) + ended: Final = _utc_minute() + finally: + release.set() + return Admissions( + held_forwarded_max_tokens="max_tokens" in forwarded, + statuses=(over.status_code, fits.status_code, held.result(timeout=60).status_code), + one_minute=started == ended, + ) + + +def test_reservation_fraction_env_var_sets_what_a_request_without_max_tokens_holds(rig: Rig) -> None: + first: Final = _admissions_beside_a_held_stream(rig, GROUPS[0]) + admissions: Final = first if first.one_minute else _admissions_beside_a_held_stream(rig, GROUPS[1]) + + assert admissions == Admissions(held_forwarded_max_tokens=False, statuses=(429, 200, 200), one_minute=True) diff --git a/tests/unit/litellm_core_utils/test_env_utils.py b/tests/unit/litellm_core_utils/test_env_utils.py index 1ddd172e274..89b091e47c2 100644 --- a/tests/unit/litellm_core_utils/test_env_utils.py +++ b/tests/unit/litellm_core_utils/test_env_utils.py @@ -6,6 +6,7 @@ import pytest from litellm._logging import verbose_logger from litellm.litellm_core_utils.env_utils import ( get_env_float, + get_env_float_in_range, get_env_float_or_none, get_env_int, get_env_int_or_none, @@ -94,3 +95,12 @@ def test_unparsable_stdout_log_limit_does_not_break_logging( assert "first line" in messages assert "second line" in messages assert len([message for message in messages if raw in message]) == 1 + + +@pytest.mark.parametrize(("raw", "expected"), [("0.5", 0.5), ("1", 1.0), ("4", 0.25), ("0", 0.25), ("half", 0.25)]) +def test_float_in_range_keeps_values_inside_the_range_and_falls_back_otherwise( + monkeypatch: pytest.MonkeyPatch, raw: str, expected: float +) -> None: + monkeypatch.setenv("ENV_UTILS_TEST_FRACTION", raw) + + assert get_env_float_in_range("ENV_UTILS_TEST_FRACTION", 0.25, 0.01, 1.0) == expected diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 19e5387da9a..88501a23235 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -24,6 +24,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.proxy.common_utils.openai_error_payload import openai_error_type from litellm.types.utils import LlmProviders @@ -54,8 +55,8 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.code == code assert caught.value.body == body assert caught.value.status_code == 400 - assert caught.value.response is response - assert caught.value.request is request + assert caught.value.response.status_code == 400 + assert caught.value.request.url == request.url assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" @@ -2940,3 +2941,80 @@ def test_transport_timeout_maps_to_timeout( assert exc_info.value.status_code == 408 assert f"Request timed out. Error_str: {reason}" in str(exc_info.value) + + +def test_bedrock_server_error_keeps_the_aws_body_without_the_signed_request() -> None: + body: Final = '{"message":"Bedrock is unable to process your request."}' + signed_request: Final = httpx.Request( + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5-5/converse", + headers={"authorization": "AWS4-HMAC-SHA256 Credential=AKIDSIGNED/20261029/us-east-1/bedrock/aws4_request"}, + content=b'{"messages":[{"role":"user","content":[{"text":"hi"}]}]}', + ) + upstream: Final = httpx.Response( + 500, content=body, headers={"x-amzn-requestid": "req-aws"}, request=signed_request + ) + + with pytest.raises(litellm.ServiceUnavailableError) as raised: + exception_type( + model="bedrock/anthropic.claude-opus-5-5", + original_exception=BedrockError(status_code=500, message=body, response=upstream), + custom_llm_provider="bedrock", + ) + + assert raised.value.response.text == body + assert raised.value.response.headers["x-amzn-requestid"] == "req-aws" + assert "AKIDSIGNED" not in str(raised.value.request.headers) + assert raised.value.request.content == b"" + + +@pytest.mark.parametrize("provider", ("gemini", "vertex_ai")) +def test_gemini_server_error_response_carries_the_upstream_error_body(provider: str) -> None: + body: Final = '{"error":{"code":500,"message":"Internal error encountered.","status":"INTERNAL"}}' + + with pytest.raises(litellm.InternalServerError) as raised: + exception_type( + model="gemini-3.1-pro", + original_exception=VertexAIError(status_code=500, message=body, headers=None), + custom_llm_provider=provider, + ) + + assert raised.value.response.status_code == 500 + assert raised.value.response.json() == json.loads(body) + + +@pytest.mark.parametrize( + ("status_code", "message", "mapped_error"), + ( + (500, "The server had an error processing your request. Sorry about that!", litellm.InternalServerError), + (520, "Web server is returning an unknown error", litellm.InternalServerError), + (504, "A timeout occurred", litellm.Timeout), + (500, "Mistral API raised a streaming error", litellm.APIError), + ), +) +def test_openai_error_matched_by_its_message_keeps_the_upstream_body( + status_code: int, message: str, mapped_error: type[Exception] +) -> None: + body: Final = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": None}}) + sent_request: Final = httpx.Request( + "POST", + "https://api.openai.com/v1/chat/completions", + headers={"authorization": "Bearer sk-openai-secret"}, + content=b'{"messages":[{"role":"user","content":"hi"}]}', + ) + upstream: Final = httpx.Response( + status_code, content=body, headers={"x-request-id": "req-openai"}, request=sent_request + ) + + with pytest.raises(mapped_error) as raised: + exception_type( + model="gpt-5.6", + original_exception=openai.APIStatusError(message, response=upstream, body=json.loads(body)), + custom_llm_provider="openai", + ) + + error_response: Final = getattr(raised.value, "response") + assert error_response.text == body + assert error_response.headers["x-request-id"] == "req-openai" + assert "sk-openai-secret" not in str(error_response.request.headers) + assert error_response.request.content == b"" diff --git a/tests/unit/llms/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py index cbbf70c691b..018b499f401 100644 --- a/tests/unit/llms/azure/test_azure_exception_mapping.py +++ b/tests/unit/llms/azure/test_azure_exception_mapping.py @@ -1,3 +1,4 @@ +import datetime import json from collections.abc import Awaitable, Callable, Mapping from typing import Final @@ -12,6 +13,7 @@ import respx import litellm from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.llms.azure.common_utils import AzureOpenAIError @respx.mock @@ -632,3 +634,202 @@ def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest. error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error)) assert type(error) is ContentPolicyViolationError + + +_AZURE_REFUSAL_BODY: Final = '{"error":{"message":"azure refused","type":"upstream_refusal","code":"refused"}}' +_AZURE_DEPLOYMENT_NOT_FOUND_BODY: Final = '{"error":{"code":"DeploymentNotFound","message":"no such deployment"}}' +_AZURE_INTERNAL_ERROR_BODY: Final = '{"error":{"code":"InternalServerError","message":"Internal server error"}}' +_AZURE_CONTENT_FILTER_BODY: Final = json.dumps( + { + "error": { + "message": "The response was filtered due to the prompt triggering Azure OpenAI's content management policy.", + "type": None, + "param": "prompt", + "code": "content_filter", + "status": 400, + "innererror": { + "code": "ResponsibleAIPolicyViolation", + "content_filter_result": {"violence": {"filtered": True, "severity": "high"}}, + }, + } + } +) +_AZURE_KEY: Final = "azure-provider-secret" +_AZURE_BASE: Final = "https://azure-upstream.openai.azure.com" +_AZURE_ARGS: Final = {"api_base": _AZURE_BASE, "api_version": "2024-10-21", "api_key": _AZURE_KEY, "max_retries": 0} + + +async def _azure_completion() -> None: + litellm.completion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS) + + +async def _azure_completion_stream() -> None: + list( + litellm.completion( + model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS + ) + ) + + +async def _azure_acompletion() -> None: + await litellm.acompletion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS) + + +async def _azure_acompletion_stream() -> None: + stream: Final = await litellm.acompletion( + model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS + ) + async for _ in stream: + pass + + +async def _azure_embedding() -> None: + litellm.embedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS) + + +async def _azure_aembedding() -> None: + await litellm.aembedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS) + + +async def _azure_image_generation() -> None: + litellm.image_generation(model="azure/gpt-image-2", prompt="a cat", **_AZURE_ARGS) + + +async def _azure_text_completion() -> None: + litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS) + + +async def _azure_text_completion_stream() -> None: + list(litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS)) + + +async def _azure_atext_completion() -> None: + await litellm.atext_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS) + + +async def _azure_atext_completion_stream() -> None: + stream: Final = await litellm.atext_completion( + model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS + ) + async for _ in stream: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + _azure_completion, + _azure_completion_stream, + _azure_acompletion, + _azure_acompletion_stream, + _azure_embedding, + _azure_aembedding, + _azure_image_generation, + _azure_text_completion, + _azure_text_completion_stream, + _azure_atext_completion, + _azure_atext_completion_stream, + ), + ids=lambda route: route.__name__.removeprefix("_azure_"), +) +@pytest.mark.parametrize( + ("status_code", "body", "mapped_error"), + ( + (400, _AZURE_REFUSAL_BODY, BadRequestError), + (400, _AZURE_CONTENT_FILTER_BODY, BadRequestError), + (403, _AZURE_REFUSAL_BODY, litellm.APIError), + (404, _AZURE_DEPLOYMENT_NOT_FOUND_BODY, litellm.NotFoundError), + (500, _AZURE_INTERNAL_ERROR_BODY, openai.APIError), + (502, _AZURE_REFUSAL_BODY, litellm.BadGatewayError), + (504, _AZURE_REFUSAL_BODY, litellm.Timeout), + ), +) +@respx.mock +async def test_azure_error_keeps_the_upstream_body_without_the_api_key( + route: Callable[[], Awaitable[None]], + status_code: int, + body: str, + mapped_error: type[Exception], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock( + return_value=httpx.Response(status_code, content=body, headers={"content-type": "application/json"}) + ) + + with pytest.raises(mapped_error) as raised: + await route() + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == body + assert error_response.json() == json.loads(body) + for exposed in (error_response.request, getattr(raised.value, "request")): + assert _AZURE_KEY not in str(exposed.headers) + assert exposed.content == b"" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + _azure_completion, + _azure_completion_stream, + _azure_acompletion, + _azure_acompletion_stream, + _azure_embedding, + _azure_aembedding, + _azure_text_completion, + _azure_atext_completion, + _azure_atext_completion_stream, + ), + ids=lambda route: route.__name__.removeprefix("_azure_"), +) +@respx.mock +async def test_azure_error_response_keeps_the_upstream_elapsed_time( + route: Callable[[], Awaitable[None]], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock( + return_value=httpx.Response(400, content=_AZURE_REFUSAL_BODY, headers={"content-type": "application/json"}) + ) + + with pytest.raises(BadRequestError) as raised: + await route() + + assert raised.value.response.elapsed > datetime.timedelta(0) + + +def test_azure_client_side_missing_key_error_keeps_the_empty_placeholder_response() -> None: + missing_key: Final = openai.OpenAIError( + "The api_key client option must be set either by passing api_key to the client or by setting the " + "OPENAI_API_KEY environment variable" + ) + + with pytest.raises(litellm.AuthenticationError) as raised: + exception_type(model="azure/gpt-4.1-mini", original_exception=missing_key, custom_llm_provider="azure") + + assert raised.value.response.status_code == 401 + assert raised.value.response.text == "" + + +@pytest.mark.parametrize( + ("status_code", "message", "mapped_error"), + ( + (500, "azure_client is not an instance of AzureOpenAI or OpenAI", litellm.APIError), + (408, "the deployment did not answer in time", litellm.Timeout), + (504, "the deployment did not answer in time", litellm.Timeout), + ), +) +def test_azure_error_litellm_raises_itself_has_no_upstream_response( + status_code: int, message: str, mapped_error: type[Exception] +) -> None: + with pytest.raises(mapped_error) as raised: + exception_type( + model="azure/gpt-4.1-mini", + original_exception=AzureOpenAIError(status_code=status_code, message=message), + custom_llm_provider="azure", + ) + + assert getattr(raised.value, "response", None) is None diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index 9e2754eef68..628c9bbd783 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -93,24 +93,34 @@ async def test_acompletion_forwards_client_headers_to_provider( assert request_headers["x-mycorp-llmcall-id"] == "abc-123" +_REJECTION_BODY: Final = '{"error":{"message":"bad prompt","type":"invalid_request_error"}}' + + @pytest.fixture def rejected_completions_endpoint(): return respx.post("https://api.openai.com/v1/completions").mock( return_value=Response( 400, - json={"error": {"message": "bad prompt", "type": "invalid_request_error"}}, - headers={"x-request-id": "req-1"}, + content=_REJECTION_BODY, + headers={"x-request-id": "req-1", "content-type": "application/json"}, ) ) +def _assert_carries_the_upstream_rejection(rejected: litellm.BadRequestError) -> None: + assert rejected.status_code == 400 + assert rejected.litellm_response_headers["x-request-id"] == "req-1" + assert rejected.response.text == _REJECTION_BODY + assert rejected.response.json() == json.loads(_REJECTION_BODY) + assert rejected.response.headers["x-request-id"] == "req-1" + + @respx.mock def test_completion_surfaces_the_status_and_headers_of_a_rejected_request(rejected_completions_endpoint): with pytest.raises(litellm.BadRequestError) as rejected: text_completion(model="gpt-3.5-turbo-instruct", prompt="hello") - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock @@ -118,8 +128,7 @@ def test_streaming_completion_surfaces_the_status_and_headers_of_a_rejected_requ with pytest.raises(litellm.BadRequestError) as rejected: list(text_completion(model="gpt-3.5-turbo-instruct", prompt="hello", stream=True)) - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock @@ -130,8 +139,7 @@ async def test_acompletion_surfaces_the_status_and_headers_of_a_rejected_request with pytest.raises(litellm.BadRequestError) as rejected: await atext_completion(model="gpt-3.5-turbo-instruct", prompt="hello") - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock diff --git a/tests/unit/llms/openai/image_variations/__init__.py b/tests/unit/llms/openai/image_variations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py b/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py new file mode 100644 index 00000000000..83ddbec6038 --- /dev/null +++ b/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py @@ -0,0 +1,51 @@ +import json +from io import BytesIO +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.openai.common_utils import OpenAIError + +_UPSTREAM_ERROR_BODY: Final = '{"error":{"message":"image is not a valid PNG","type":"invalid_request_error"}}' + + +def _upstream_rejection(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + content=_UPSTREAM_ERROR_BODY, + headers={"x-request-id": "req_variation_400", "content-type": "application/json"}, + ) + + +_VARIATION_KWARGS: Final = { + "custom_llm_provider": "openai", + "api_key": "transport-only", + "api_base": "https://upstream.example/v1", + "max_retries": 0, +} + + +async def _request_variation(is_async: bool) -> None: + if is_async: + await litellm.aimage_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS) + return + litellm.image_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_rejected_image_variation_carries_the_upstream_error_response( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + transport: Final = httpx.MockTransport(_upstream_rejection) + monkeypatch.setattr(litellm, "client_session", httpx.Client(transport=transport)) + monkeypatch.setattr(litellm, "aclient_session", httpx.AsyncClient(transport=transport)) + + with pytest.raises((litellm.BadRequestError, OpenAIError)) as raised: + await _request_variation(is_async) + + assert raised.value.response.text == _UPSTREAM_ERROR_BODY + assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY) + assert raised.value.response.headers["x-request-id"] == "req_variation_400" diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 5b27a446fb3..3e1cbad9308 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -1,8 +1,9 @@ import asyncio +import gzip import importlib import json import os -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from types import ModuleType from typing import Final, Literal, Protocol from unittest.mock import MagicMock, Mock, patch @@ -375,6 +376,277 @@ async def test_acompletion_preserves_public_errors_for_both_http_clients( assert not http_client.is_closed +_UPSTREAM_ERROR_BODY: Final = ( + '{"error":{"message":"upstream rejected the request","type":"invalid_request_error",' + '"param":"response_format.type","code":"invalid_value"}}' +) +_UPSTREAM_BASE: Final = "https://upstream.example/v1" +_UPSTREAM_TRANSPORT: Final = httpx.MockTransport( + lambda _request: httpx.Response( + 400, + content=_UPSTREAM_ERROR_BODY, + headers={"x-request-id": "req_upstream_400", "x-upstream-trace": "trace-7", "content-type": "application/json"}, + ) +) +_MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +def _rejecting_client() -> OpenAI: + return OpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.Client(transport=_UPSTREAM_TRANSPORT), + ) + + +def _rejecting_async_client() -> AsyncOpenAI: + return AsyncOpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.AsyncClient(transport=_UPSTREAM_TRANSPORT), + ) + + +def _upstream(provider: str, model: str) -> dict[str, object]: + return {"model": f"{provider}/{model}", "api_base": _UPSTREAM_BASE, "api_key": "transport-only", "num_retries": 0} + + +async def _completion(provider: str) -> None: + litellm.completion(messages=_MESSAGES, client=_rejecting_client(), **_upstream(provider, "gpt-5.6")) + + +async def _completion_stream(provider: str) -> None: + stream: Final = litellm.completion( + messages=_MESSAGES, client=_rejecting_client(), stream=True, **_upstream(provider, "gpt-5.6") + ) + list(stream) + + +async def _acompletion(provider: str) -> None: + await litellm.acompletion(messages=_MESSAGES, client=_rejecting_async_client(), **_upstream(provider, "gpt-5.6")) + + +async def _acompletion_stream(provider: str) -> None: + stream: Final = await litellm.acompletion( + messages=_MESSAGES, client=_rejecting_async_client(), stream=True, **_upstream(provider, "gpt-5.6") + ) + async for _ in stream: + pass + + +async def _embedding(provider: str) -> None: + litellm.embedding(input=["hi"], client=_rejecting_client(), **_upstream(provider, "text-embedding-3-large")) + + +async def _aembedding(provider: str) -> None: + await litellm.aembedding( + input=["hi"], client=_rejecting_async_client(), **_upstream(provider, "text-embedding-3-large") + ) + + +async def _image_generation(provider: str) -> None: + litellm.image_generation(prompt="a cat", client=_rejecting_client(), **_upstream(provider, "gpt-image-2")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize( + "route", + (_completion, _completion_stream, _acompletion, _acompletion_stream, _embedding, _aembedding, _image_generation), + ids=lambda route: route.__name__.lstrip("_"), +) +async def test_mapped_error_carries_the_upstream_error_response( + provider: str, route: Callable[[str], Awaitable[None]] +) -> None: + with pytest.raises(litellm.BadRequestError) as raised: + await route(provider) + + assert raised.value.response.text == _UPSTREAM_ERROR_BODY + assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY) + assert raised.value.response.headers["x-request-id"] == "req_upstream_400" + assert raised.value.response.headers["x-upstream-trace"] == "trace-7" + + +_UPSTREAM_SERVER_ERROR_BODY: Final = ( + '{"error":{"message":"The server had an error while processing your request.","type":"server_error",' + '"param":null,"code":null}}' +) +_CHUNKED_UPSTREAM_HEADERS: Final = { + "content-type": "application/json", + "transfer-encoding": "chunked", + "x-request-id": "req_upstream_5xx", +} + + +def _chunked_upstream_client(status_code: int, content_encoding: Literal["gzip"] | None) -> httpx.AsyncClient: + body: Final = _UPSTREAM_SERVER_ERROR_BODY.encode() + encoding_headers: Final = {"content-encoding": content_encoding} if content_encoding else {} + return httpx.AsyncClient( + transport=httpx.MockTransport( + lambda _request: httpx.Response( + status_code, + headers={**_CHUNKED_UPSTREAM_HEADERS, **encoding_headers}, + stream=httpx.ByteStream(gzip.compress(body) if content_encoding else body), + ) + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize( + ("status_code", "mapped_error"), + ( + (429, litellm.RateLimitError), + (500, litellm.InternalServerError), + (502, litellm.BadGatewayError), + (503, litellm.ServiceUnavailableError), + ), +) +@pytest.mark.parametrize("content_encoding", (None, "gzip")) +async def test_retryable_error_keeps_the_decoded_upstream_body_and_its_headers( + provider: str, + stream: bool, + status_code: int, + mapped_error: type[Exception], + content_encoding: Literal["gzip"] | None, +) -> None: + client: Final = AsyncOpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=_chunked_upstream_client(status_code, content_encoding), + ) + + with pytest.raises(mapped_error) as raised: + await litellm.acompletion(messages=_MESSAGES, client=client, stream=stream, **_upstream(provider, "gpt-5.6")) + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == _UPSTREAM_SERVER_ERROR_BODY + assert error_response.json() == json.loads(_UPSTREAM_SERVER_ERROR_BODY) + assert dict(error_response.headers) == { + **_CHUNKED_UPSTREAM_HEADERS, + **({"content-encoding": content_encoding} if content_encoding else {}), + } + + +_REFUSAL_BODY: Final = '{"error":{"message":"upstream refused","type":"upstream_refusal","param":null,"code":null}}' +_GATEWAY_TIMEOUT_BODY: Final = ( + '{"error":{"message":"litellm.Timeout: APITimeoutError - Request timed out.","type":"invalid_request_error",' + '"param":null,"code":"408"}}' +) +_PROVIDER_KEY: Final = "sk-provider-secret" +_QUERY_SECRET: Final = "query-secret" + + +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize( + ("status_code", "mapped_error", "body"), + ( + (400, litellm.BadRequestError, _REFUSAL_BODY), + (403, litellm.APIError, _REFUSAL_BODY), + (408, litellm.Timeout, _REFUSAL_BODY), + (408, litellm.Timeout, _GATEWAY_TIMEOUT_BODY), + (409, litellm.APIError, _REFUSAL_BODY), + (413, litellm.APIError, _REFUSAL_BODY), + (429, litellm.RateLimitError, _REFUSAL_BODY), + (500, litellm.InternalServerError, _REFUSAL_BODY), + (504, litellm.Timeout, _REFUSAL_BODY), + ), +) +def test_mapped_error_keeps_the_upstream_body_without_the_provider_credentials( + provider: str, status_code: int, mapped_error: type[Exception], body: str +) -> None: + client: Final = OpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + default_query={"key": _QUERY_SECRET}, + max_retries=0, + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda _request: httpx.Response( + status_code, content=body, headers={"content-type": "application/json"} + ) + ) + ), + ) + + with pytest.raises(mapped_error) as raised: + litellm.completion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + client=client, + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == body + assert error_response.json() == json.loads(body) + assert str(error_response.request.url) == f"{_UPSTREAM_BASE}/chat/completions" + for exposed in (error_response.request, getattr(raised.value, "request")): + assert _QUERY_SECRET not in str(exposed.url) + assert _PROVIDER_KEY not in str(exposed.headers) + assert exposed.content == b"" + + +def _time_out(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout("no reply", request=request) + + +async def _sync_completion_timing_out(provider: str) -> None: + litellm.completion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + client=OpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(_time_out)), + ), + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + + +async def _async_stream_timing_out(provider: str) -> None: + stream: Final = await litellm.acompletion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + stream=True, + client=AsyncOpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_time_out)), + ), + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + async for _ in stream: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize("call", (_sync_completion_timing_out, _async_stream_timing_out)) +async def test_timeout_before_any_reply_has_no_upstream_response( + provider: str, call: Callable[[str], Awaitable[None]] +) -> None: + with pytest.raises(litellm.Timeout) as raised: + await call(provider) + + assert getattr(raised.value, "response", None) is None + + @pytest.mark.asyncio async def test_acompletion_streams_text_deltas_over_injected_transport( http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index aeab7a3d348..c867db2fd02 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,6 +2,7 @@ import asyncio import contextvars import copy import datetime +import gzip import json from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace @@ -13,6 +14,7 @@ from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stub import anthropic import httpx +import openai import pytest import respx from fastapi import HTTPException, Request, Response, status @@ -11675,6 +11677,67 @@ async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_r assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400" +_UPSTREAM_REJECTION_BODY: Final = b'{"error":{"message":"upstream said no","type":"upstream_rejection"}}' + + +def _upstream_rejection_wire(content_encoding: Literal["gzip"] | None) -> tuple[bytes, dict[str, str]]: + wire_body: Final = gzip.compress(_UPSTREAM_REJECTION_BODY) if content_encoding else _UPSTREAM_REJECTION_BODY + return wire_body, { + "x-request-id": "req_upstream", + "x-upstream-trace": "trace-7", + "content-type": "application/json", + "content-length": str(len(wire_body)), + **({"content-encoding": content_encoding} if content_encoding else {}), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "mapped_error"), + ( + (400, litellm.BadRequestError), + (429, litellm.RateLimitError), + (503, litellm.ServiceUnavailableError), + ), +) +@pytest.mark.parametrize("content_encoding", (None, "gzip")) +async def test_handle_llm_api_exception_forwards_exactly_the_upstream_headers_of_a_mapped_openai_error( + status_code: int, mapped_error: type[Exception], content_encoding: Literal["gzip"] | None +) -> None: + wire_body, upstream_headers = _upstream_rejection_wire(content_encoding) + transport: Final = httpx.MockTransport( + lambda _request: httpx.Response(status_code, content=wire_body, headers=upstream_headers) + ) + with pytest.raises(mapped_error) as mapped: + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + num_retries=0, + client=openai.AsyncOpenAI( + api_key="transport-only", + base_url="https://upstream.example/v1", + max_retries=0, + http_client=httpx.AsyncClient(transport=transport), + ), + ) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(ProxyException) as forwarded: + await ProxyBaseLLMRequestProcessing(data={}).handle_llm_api_exception( + e=mapped.value, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert forwarded.value.code == str(status_code) + assert "upstream said no" in forwarded.value.message + assert {name: value for name, value in forwarded.value.headers.items() if name.startswith("llm_provider-")} == { + f"llm_provider-{name}": value for name, value in upstream_headers.items() + } + + @pytest.mark.asyncio async def test_handle_llm_api_exception_logs_bug_report_for_unmapped_error( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 99b049f1803..d9c34e95364 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -19,7 +19,14 @@ from unittest.mock import patch _ENV_HELPERS = frozenset( - {"get_env_int", "get_env_float", "get_env_int_in_range", "get_env_int_or_none", "get_env_float_or_none"} + { + "get_env_int", + "get_env_float", + "get_env_int_in_range", + "get_env_float_in_range", + "get_env_int_or_none", + "get_env_float_or_none", + } ) diff --git a/tests/unit/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py index 7acacf05226..c5a0868f36d 100644 --- a/tests/unit/test_exception_header_preservation.py +++ b/tests/unit/test_exception_header_preservation.py @@ -13,6 +13,7 @@ from typing import Final import httpx import pytest +import litellm from litellm.exceptions import ( BadRequestError, ContentPolicyViolationError, @@ -418,3 +419,21 @@ class TestBedrockErrorHeaders: headers = get_response_headers(dict(_response_headers)) assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500" + + +@pytest.mark.parametrize( + "error_class", + (RateLimitError, ServiceUnavailableError, litellm.BadGatewayError, litellm.InternalServerError), +) +def test_error_built_from_an_unread_streaming_response_keeps_its_headers(error_class: type[Exception]) -> None: + unread: Final = httpx.Response( + 500, + headers={"x-request-id": "req-unread"}, + stream=httpx.ByteStream(b'{"error":"never read"}'), + ) + + error: Final = error_class(message="upstream failed", llm_provider="openai", model="gpt-5.6", response=unread) + + error_response: Final = getattr(error, "response") + assert error_response.text == "" + assert dict(error_response.headers) == {"x-request-id": "req-unread"} diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 05921aa70b3..e5ff4531f84 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -3,6 +3,8 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). """ import asyncio +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime from typing import Final import pytest @@ -10,13 +12,16 @@ import pytest import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_CACHE_KEY, ITPM_RESERVED_KEY, OTPM_CACHE_KEY, OTPM_RESERVED_KEY, + _explicit_max_tokens, _reservation_value, - _resolve_max_tokens, async_io_token_pre_call_check, async_io_token_reconcile_success, build_io_token_rate_limit_headers, @@ -30,7 +35,16 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage +from litellm.utils import get_utc_datetime + +MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7" +MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"] +UNMAPPED_MODEL = "openai/litellm-io-unmapped-model" +UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT = 4096 +AZURE_BASE_MODEL = "azure/gpt-5" +AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"] class TestIOTokenRateLimitHelpers: @@ -48,13 +62,23 @@ class TestIOTokenRateLimitHelpers: # A real non-zero estimate is reserved as-is. assert _reservation_value(42, 100) == 42 - def test_resolve_max_tokens_respects_explicit_zero(self): - deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} - # An explicit max_tokens=0 is honored, not replaced by the model default. - assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 - # max_completion_tokens is the fallback only when max_tokens is absent. - assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 - assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + @pytest.mark.parametrize( + "request_kwargs, expected", + [ + ({"max_tokens": 0}, 0), + ({"max_completion_tokens": 12}, 12), + ({"max_output_tokens": 9}, 9), + ({"max_tokens": 5, "max_completion_tokens": 12}, 5), + ({"max_tokens": None, "max_output_tokens": 9}, 9), + ({"max_tokens": [100]}, None), + ({"messages": []}, None), + (None, None), + ], + ) + def test_explicit_max_tokens_reads_every_output_cap_key( + self, request_kwargs: dict[str, object] | None, expected: int | None + ): + assert _explicit_max_tokens(request_kwargs) == expected def test_build_io_token_rate_limit_headers(self): headers = build_io_token_rate_limit_headers( @@ -386,6 +410,244 @@ class TestModelRateLimitingCheckIOTokens: normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + @pytest.mark.parametrize( + "model, otpm_limit", + [ + (MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT // 2), + (MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT - 1), + (UNMAPPED_MODEL, 1000), + (UNMAPPED_MODEL, 3), + ], + ) + async def test_request_without_max_tokens_is_admitted_below_model_max_output( + self, model: str, otpm_limit: int, use_sync_check: bool + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": otpm_limit}, + "model_info": {"id": f"io-no-max-tokens-{otpm_limit}"}, + "model_name": "no-max-tokens", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY] + assert isinstance(reserved, int) + assert 0 < reserved < otpm_limit + assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model, otpm_limit", + [(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)], + ) + async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot( + self, model: str, otpm_limit: int + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": otpm_limit}, + "model_info": {"id": "io-no-max-tokens-race-id"}, + "model_name": "no-max-tokens", + } + + async def _attempt() -> dict[str, object] | None: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + try: + await check.async_pre_call_check(deployment) + except litellm.RateLimitError: + return None + return request_kwargs["metadata"] + + admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None] + otpm_key: Final = admitted[0][OTPM_CACHE_KEY] + + assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) + assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key} + assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) + assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model, model_max_output", + [(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)], + ) + async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows( + self, model: str, model_max_output: int + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": model_max_output * 10}, + "model_info": {"id": "io-no-max-tokens-roomy-id"}, + "model_name": "roomy", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + await check.async_pre_call_check(deployment) + + assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == model_max_output + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10 + + async def _reserved(model_info: dict[str, str]) -> int: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + deployment: Final = { + "litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit}, + "model_info": model_info, + "model_name": "custom-name", + } + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + with_base_model: Final = await _reserved({"id": "io-with-base-model", "base_model": AZURE_BASE_MODEL}) + without_base_model: Final = await _reserved({"id": "io-without-base-model"}) + + assert with_base_model == without_base_model + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000}, + "model_info": {"id": "io-bad-max-tokens-id"}, + "model_name": "opus", + } + + async def _pre_call(request_kwargs: dict[str, object]) -> None: + set_io_token_rate_limit_request_kwargs(request_kwargs) + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + await _pre_call({"messages": [{"role": "user", "content": "hello there"}], "max_tokens": "abc", "metadata": {}}) + valid_kwargs: Final = { + "messages": [{"role": "user", "content": "hello there"}], + "max_tokens": 1, + "metadata": {}, + } + await _pre_call(valid_kwargs) + + itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY]) + assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY] + + @pytest.mark.asyncio + async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 100}, + "model_info": {"id": "io-explicit-max-tokens-id"}, + "model_name": "opus", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + await check.async_pre_call_check(deployment) + + assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100 + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100}, + "model_info": {"id": "io-deployment-max-tokens-id"}, + "model_name": "opus", + } + without_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}} + + for request_kwargs in (without_cap, with_cap): + set_io_token_rate_limit_request_kwargs(request_kwargs) + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + assert without_cap["metadata"][OTPM_RESERVED_KEY] == 100 + assert with_cap["metadata"][OTPM_RESERVED_KEY] == 300 + + @pytest.mark.asyncio + async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + chat: Final = {"messages": [{"role": "user", "content": "hello"}]} + responses: Final = {"input": "hello"} + + async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int: + request_kwargs: Final = {**request_shape, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check( + { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap}, + "model_info": {"id": f"io-sent-cap-{sorted(deployment_cap)}-{sorted(request_shape)}"}, + "model_name": "opus", + } + ) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + chat_default: Final = await reserved({}, chat) + responses_default: Final = await reserved({}, responses) + + assert chat_default == responses_default > 100 + assert await reserved({"max_tokens": 100}, chat) == 100 + assert await reserved({"max_completion_tokens": 100}, chat) == 100 + assert await reserved({"max_output_tokens": 100}, chat) == chat_default + assert await reserved({"max_output_tokens": 100}, responses) == 100 + assert await reserved({"max_tokens": 100}, responses) == responses_default + assert await reserved({"max_completion_tokens": 100}, responses) == responses_default + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5 + + async def reserved(deployment_cap: Mapping[str, int]) -> int: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + deployment: Final = { + "litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap}, + "model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"}, + "model_name": "unmapped", + } + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + default: Final = await reserved({}) + + assert await reserved({"max_tokens": default + 1}) == default + assert await reserved({"max_tokens": otpm_limit * 2}) == default + def test_sync_io_pre_call_reserves_and_reconciles(self): from litellm.utils import get_utc_datetime @@ -415,9 +677,9 @@ class TestModelRateLimitingCheckIOTokens: "standard_logging_object": { "model_id": "io-sync-id", "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), + "metadata": {}, }, - "metadata": request_kwargs["metadata"], + "litellm_params": {"metadata": request_kwargs["metadata"]}, } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], @@ -709,6 +971,32 @@ class TestModelRateLimitingCheckIOTokens: assert await dual_cache.async_get_cache(key=itpm_key) == 4 + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_reconcile", [False, True]) + async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in( + self, use_sync_reconcile: bool + ): + dual_cache: Final = DualCache() + counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm" + reserved_minute_key: Final = f"{counter}:99-99" + await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60) + kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}} + response: Final = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29), + ) + minute_before: Final = get_utc_datetime().strftime("%H-%M") + + if use_sync_reconcile: + io_token_reconcile_success(dual_cache, kwargs, response) + else: + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")} + charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes] + assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10 + assert sum(value or 0 for value in charged) == 25 - 10 + @pytest.mark.asyncio async def test_reconcile_missing_usage_keeps_reservation(self): dual_cache = DualCache() @@ -799,16 +1087,16 @@ class TestModelRateLimitingCheckIOTokens: assert await dual_cache.async_get_cache(key=itpm_key) == 4 - def test_sync_reconcile_anthropic_dict_usage(self): + def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self): dual_cache = DualCache() itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) dual_cache.set_cache(key=otpm_key, value=4, ttl=60) kwargs = { "metadata": { - ITPM_RESERVED_KEY: 6, + ITPM_RESERVED_KEY: 10, OTPM_RESERVED_KEY: 4, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key, @@ -816,12 +1104,17 @@ class TestModelRateLimitingCheckIOTokens: } response = { "type": "message", - "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + "usage": { + "input_tokens": 3, + "output_tokens": 2, + "cache_read_input_tokens": 1, + "cache_creation_input_tokens": 4, + }, } io_token_reconcile_success(dual_cache, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=itpm_key) == 3 + 4 assert dual_cache.get_cache(key=otpm_key) == 2 @pytest.mark.asyncio @@ -958,7 +1251,214 @@ class TestModelRateLimitingCheckIOTokens: assert current == 0 +ROUTER_IO_MODEL: Final = "anthropic/claude-haiku-5-5" +RouterCall = Callable[[Router, str | Exception], Awaitable[tuple[int, int]]] + + +def _router_with_io_limits() -> Router: + return Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": ROUTER_IO_MODEL, "api_key": "sk-test", "itpm": 50000, "otpm": 50000}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=1, + ) + + +def _current_minute() -> str: + return get_utc_datetime().strftime("%H-%M") + + +def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int: + deployment: Final = router.model_list[0] + return sum( + router.cache.get_cache( + key=template.value.format( + id=deployment["model_info"]["id"], + model=deployment["litellm_params"]["model"], + current_minute=minute, + ) + ) + or 0 + for minute in minutes + ) + + +async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, int]: + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + minutes: Final = frozenset({first_minute, _current_minute()}) + return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes) + + +async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.acompletion( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + mock_response=mock_response, + ) + return response.usage.prompt_tokens, response.usage.completion_tokens + + +async def _proxy_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.acompletion( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response.usage.prompt_tokens, response.usage.completion_tokens + + +async def _proxy_responses(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.aresponses( + model="limited", + input="hello there", + max_output_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response.usage.input_tokens, response.usage.output_tokens + + +async def _proxy_messages(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.aanthropic_messages( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response["usage"]["input_tokens"], response["usage"]["output_tokens"] + + +ProxyCallSharingLogging = Callable[[Router, str | Exception, LiteLLMLoggingObject], Awaitable[object]] +PROXY_MESSAGES: Final = [{"role": "user", "content": "hello there"}] + + +def _logging_obj_shared_across_attempts(original_function: str) -> LiteLLMLoggingObject: + logging_obj, _ = litellm.utils.function_setup( + original_function, + litellm.utils.Rules(), + datetime.now(), + model="limited", + messages=PROXY_MESSAGES, + litellm_call_id="proxy-call", + ) + return logging_obj + + +async def _proxy_chat_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.acompletion( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + +async def _proxy_responses_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.aresponses( + model="limited", + input="hello there", + max_output_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + +async def _proxy_messages_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.aanthropic_messages( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "call", + [_sdk_chat, _proxy_chat, _proxy_responses, _proxy_messages], + ids=["sdk_chat", "proxy_chat_metadata", "proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"], + ) + async def test_success_reconciles_io_counters_to_real_usage(self, call: RouterCall): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + usage: Final = await call(router, "hi") + + assert await _io_counters_since(router, first_minute) == usage + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call", "mock_error"), + [ + ( + _proxy_responses, + litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL), + ), + (_proxy_messages, "litellm.InternalServerError"), + ], + ids=["proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"], + ) + async def test_retried_failure_refunds_each_io_reservation_once( + self, call: RouterCall, mock_error: str | Exception + ): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + with pytest.raises(litellm.InternalServerError): + await call(router, mock_error) + + assert await _io_counters_since(router, first_minute) == (0, 0) + + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call", "original_function", "mock_error"), + [ + (_proxy_chat_sharing_logging, "acompletion", "litellm.InternalServerError"), + ( + _proxy_responses_sharing_logging, + "aresponses", + litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL), + ), + (_proxy_messages_sharing_logging, "anthropic_messages", "litellm.InternalServerError"), + ], + ids=["proxy_chat", "proxy_responses", "proxy_messages"], + ) + async def test_retried_failure_sharing_one_logging_object_refunds_every_attempt( + self, call: ProxyCallSharingLogging, original_function: str, mock_error: str | Exception + ): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + with pytest.raises(litellm.InternalServerError): + await call(router, mock_error, _logging_obj_shared_across_attempts(original_function)) + + assert await _io_counters_since(router, first_minute) == (0, 0) + @pytest.mark.asyncio async def test_model_group_info_aggregates_io_limits(self): router = Router(