From a4624b6c6c4bb1ddf753bd57bbfd37b711598c90 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 18 Sep 2026 16:03:45 -0700
Subject: [PATCH] fix(gemini): derive the finish reason key set from the
 Candidates type

Candidates.finishReason listed eleven values while the mapping key set
carried twenty-one, so typed fixtures could not spell the reasons this
PR handles. GeminiFinishReason is now the one list, the key set derives
from it, and a test checks every documented reason has an explicit
mapping instead of falling through to "stop"
---
 .../vertex_and_google_ai_studio_gemini.py     | 29 ++------------
 litellm/types/llms/vertex_ai.py               | 39 ++++++++++++-------
 ...test_vertex_and_google_ai_studio_gemini.py |  9 ++++-
 3 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index a95b845718a..46f1b948026 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -6,7 +6,7 @@ import time
 from collections.abc import Callable, Mapping, Sequence
 from copy import deepcopy
 from functools import partial
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args
 
 import httpx
 
@@ -57,6 +57,7 @@ from litellm.types.llms.vertex_ai import (
     ContentType,
     FunctionCallingConfig,
     FunctionDeclaration,
+    GeminiFinishReason,
     GeminiThinkingConfig,
     GenerateContentResponseBody,
     HttpxPartType,
@@ -1330,31 +1331,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
             "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.",
         }
 
-    _GEMINI_FINISH_REASON_KEYS = frozenset(
-        {
-            "STOP",
-            "MAX_TOKENS",
-            "SAFETY",
-            "RECITATION",
-            "FINISH_REASON_UNSPECIFIED",
-            "MALFORMED_FUNCTION_CALL",
-            "LANGUAGE",
-            "OTHER",
-            "BLOCKLIST",
-            "PROHIBITED_CONTENT",
-            "SPII",
-            "IMAGE_SAFETY",
-            "IMAGE_PROHIBITED_CONTENT",
-            "TOO_MANY_TOOL_CALLS",
-            "MALFORMED_RESPONSE",
-            "NO_IMAGE",
-            "IMAGE_RECITATION",
-            "IMAGE_OTHER",
-            "ESCALATION",
-            "UNEXPECTED_TOOL_CALL",
-            "MISSING_THOUGHT_SIGNATURE",
-        }
-    )
+    _GEMINI_FINISH_REASON_KEYS: Final[frozenset[str]] = frozenset(get_args(GeminiFinishReason))
 
     @staticmethod
     def get_finish_reason_mapping() -> dict[str, OpenAIChatCompletionFinishReason]:
diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py
index 3b95b786631..ce51e46ef15 100644
--- a/litellm/types/llms/vertex_ai.py
+++ b/litellm/types/llms/vertex_ai.py
@@ -425,22 +425,35 @@ class UrlContextMetadata(TypedDict, total=False):
     urlMetadata: list[UrlMetadata]
 
 
+GeminiFinishReason = Literal[
+    "FINISH_REASON_UNSPECIFIED",
+    "STOP",
+    "MAX_TOKENS",
+    "SAFETY",
+    "RECITATION",
+    "LANGUAGE",
+    "OTHER",
+    "BLOCKLIST",
+    "PROHIBITED_CONTENT",
+    "SPII",
+    "MALFORMED_FUNCTION_CALL",
+    "IMAGE_SAFETY",
+    "IMAGE_PROHIBITED_CONTENT",
+    "TOO_MANY_TOOL_CALLS",
+    "MALFORMED_RESPONSE",
+    "NO_IMAGE",
+    "IMAGE_RECITATION",
+    "IMAGE_OTHER",
+    "ESCALATION",
+    "UNEXPECTED_TOOL_CALL",
+    "MISSING_THOUGHT_SIGNATURE",
+]
+
+
 class Candidates(TypedDict, total=False):
     index: int
     content: HttpxContentType
-    finishReason: Literal[
-        "FINISH_REASON_UNSPECIFIED",
-        "STOP",
-        "MAX_TOKENS",
-        "SAFETY",
-        "RECITATION",
-        "OTHER",
-        "BLOCKLIST",
-        "PROHIBITED_CONTENT",
-        "SPII",
-        "MALFORMED_FUNCTION_CALL",
-        "IMAGE_SAFETY",
-    ]
+    finishReason: GeminiFinishReason
     safetyRatings: list[SafetyRatings]
     citationMetadata: CitationMetadata
     groundingMetadata: GroundingMetadata
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 4ee199bd9af..6c818016c87 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -2,7 +2,7 @@ import asyncio
 import json
 import re
 from copy import deepcopy
-from typing import Final, List, cast
+from typing import Final, List, cast, get_args
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -18,7 +18,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     VertexGeminiConfig,
 )
-from litellm.types.llms.vertex_ai import UsageMetadata
+from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata
 from litellm.types.utils import ChoiceLogprobs, Usage
 from litellm.utils import CustomStreamWrapper
 
@@ -940,6 +940,11 @@ def test_check_finish_reason():
         )
 
 
+def test_every_documented_gemini_finish_reason_has_an_explicit_mapping():
+    documented: Final = frozenset(get_args(GeminiFinishReason))
+    assert set(VertexGeminiConfig.get_finish_reason_mapping()) == documented
+
+
 def test_finish_reason_unspecified_and_malformed_function_call():
     """
     Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL