Merge pull request #18283 from lucasrothman/fix/gemini-context-window-error

fix(gemini): properly catch context window exceeded errors

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

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 = [