diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 7bf95ca340..1517d1e776 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -78,9 +78,7 @@ class ExceptionCheckers: "is longer than the model's context length", "input tokens exceed the configured limit", "`inputs` tokens + `max_new_tokens` must be", - # Gemini pattern: "The input token count exceeds the maximum number of tokens allowed" - # See: https://github.com/BerriAI/litellm/issues/XXXX - "input token count exceeds the maximum number of tokens allowed", + "exceeds the maximum number of tokens allowed", # Gemini ] for substring in known_exception_substrings: if substring in _error_str_lowercase: @@ -1262,6 +1260,14 @@ def exception_type( # type: ignore # noqa: PLR0915 model=model, llm_provider=custom_llm_provider, ) + elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): + exception_mapping_worked = True + raise ContextWindowExceededError( + message=f"ContextWindowExceededError: {custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) elif ( "None Unknown Error." in error_str or "Content has no parts." in error_str diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index f69b9c3523..9e742a83c6 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -40,8 +40,7 @@ context_window_test_cases = [ "`inputs` tokens + `max_new_tokens` must be <= 4096", True, ), - # Gemini context window error format - # See: https://github.com/BerriAI/litellm/issues/XXXX + # Gemini 2.5/3 format ( "The input token count exceeds the maximum number of tokens allowed 1048576.", True, @@ -50,6 +49,15 @@ context_window_test_cases = [ "GeminiException BadRequestError - {\n \"error\": {\n \"code\": 400,\n \"message\": \"The input token count exceeds the maximum number of tokens allowed 1048576.\",\n \"status\": \"INVALID_ARGUMENT\"\n }\n}\n", True, ), + # Gemini 2.0 Flash format (includes input token count in message) + ( + "The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).", + True, + ), + ( + "GeminiException BadRequestError - {\n \"error\": {\n \"code\": 400,\n \"message\": \"The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).\",\n \"status\": \"INVALID_ARGUMENT\"\n }\n}\n", + True, + ), # Test case insensitivity ("ERROR: THIS MODEL'S MAXIMUM CONTEXT LENGTH IS 1024.", True), # Cerebras context window error format @@ -169,6 +177,54 @@ class TestExceptionCheckers: result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is False, f"Should NOT detect policy violation in: {error_str}" +gemini_context_window_test_cases = [ + # Gemini 2.0 Flash format (includes input token count in message) + ( + "The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).", + True, + ), + # Gemini 2.5/3 format + ( + "The input token count exceeds the maximum number of tokens allowed (1048576).", + True, + ), + ("A generic error occurred.", False), +] + + +@pytest.mark.parametrize( + "error_message, should_raise_context_window", gemini_context_window_test_cases +) +def test_gemini_context_window_error_mapping(error_message, should_raise_context_window): + """ + Tests that the exception_type function correctly maps Gemini's + context window exceeded errors to litellm.ContextWindowExceededError. + """ + model = "gemini/gemini-2.0-flash" + custom_llm_provider = "gemini" + + # Create a generic exception with the specific error message + original_exception = Exception(error_message) + + if should_raise_context_window: + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model=model, + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + ) + # Check if the raised exception is indeed a ContextWindowExceededError + assert isinstance(excinfo.value, litellm.ContextWindowExceededError) + else: + # For the negative case, we expect it to raise a generic APIConnectionError + with pytest.raises(litellm.APIConnectionError): + exception_type( + model=model, + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + ) + + # Test cases for Vertex AI RateLimitError mapping # As per https://github.com/BerriAI/litellm/issues/16189 vertex_rate_limit_test_cases = [