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"
This commit is contained in:
mateo-berri 2026-09-18 16:03:45 -07:00
parent 66c01cf35c
commit a4624b6c6c
3 changed files with 36 additions and 41 deletions

View file

@ -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]:

View file

@ -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

View file

@ -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