From d5a36eb2cabfbcd0530b4a43576c9ac7d6c530a7 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 00:36:07 +0000 Subject: [PATCH 1/4] fix(gemini): propagate provider modelVersion onto model responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 10 ++- ...test_vertex_and_google_ai_studio_gemini.py | 62 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 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 d113b2b4f6b..eb03b17435c 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 @@ -2434,7 +2434,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_response = GenerateContentResponseBody(**completion_response) ## GET MODEL ## - model_response.model = model + model_version: Final = completion_response.get("modelVersion") + model_response.model = model_version if isinstance(model_version, str) else model ## CHECK IF RESPONSE FLAGGED if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]: @@ -3264,7 +3265,12 @@ class ModelResponseIterator: processed_chunk: Final = GenerateContentResponseBody(**chunk) response_id: Final = processed_chunk.get("responseId") - model_response = ModelResponseStream(choices=[], id=response_id) + chunk_model_version: Final = processed_chunk.get("modelVersion") + model_response = ModelResponseStream( + choices=[], + id=response_id, + model=chunk_model_version if isinstance(chunk_model_version, str) else None, + ) # Check if prompt is blocked due to content filtering blocked_response: Final = VertexGeminiConfig._check_prompt_level_content_filter( 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 101f6e6fa5d..29a47cd54f2 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 @@ -5836,3 +5836,65 @@ def test_supported_reasoning_efforts_still_map(model): drop_params=False, ) assert "thinkingConfig" in result + + +def _generate_content_body() -> dict: + return { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 7, + "totalTokenCount": 12, + }, + } + + +def test_generate_content_transform_uses_reported_model_version(): + """The served modelVersion must win over the requested name so downstream + pricing sees what actually ran.""" + import httpx + + body = {**_generate_content_body(), "modelVersion": "gemini-x-served"} + response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response=body, + model_response=ModelResponse(), + model="gemini-x", + logging_obj=MagicMock(), + raw_response=httpx.Response(200, headers={}), + ) + + assert response.model == "gemini-x-served" + + +def test_generate_content_transform_falls_back_to_requested_model(): + import httpx + + response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response=_generate_content_body(), + model_response=ModelResponse(), + model="gemini-x", + logging_obj=MagicMock(), + raw_response=httpx.Response(200, headers={}), + ) + + assert response.model == "gemini-x" + + +def test_streaming_chunk_carries_model_version(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk = {**_generate_content_body(), "modelVersion": "gemini-x-served"} + iterator: Final = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=MagicMock() + ) + streaming_chunk: Final = iterator.chunk_parser(chunk) + + assert streaming_chunk.model == "gemini-x-served" From 801f6a92b4a2d0e55daa738ff9fd3db0098382c6 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:05:53 +0000 Subject: [PATCH 2/4] test(gemini): assert the served modelVersion reaches the assembled stream through CustomStreamWrapper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...test_vertex_and_google_ai_studio_gemini.py | 39 +++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) 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 29a47cd54f2..3fb00a7d962 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 @@ -5892,9 +5892,42 @@ def test_streaming_chunk_carries_model_version(): ) chunk = {**_generate_content_body(), "modelVersion": "gemini-x-served"} - iterator: Final = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) streaming_chunk: Final = iterator.chunk_parser(chunk) assert streaming_chunk.model == "gemini-x-served" + + +def test_served_model_version_reaches_assembled_stream_through_custom_stream_wrapper(): + """Greptile claimed CustomStreamWrapper drops the served modelVersion before + pricing. It does not: chunk_creator stashes the parser chunk's model into each + yielded chunk's _hidden_params["provider_response_model"], and + stream_chunk_builder carries it onto the assembled response. The wrapper's + trailing bookkeeping chunk is the one emission that legitimately omits it.""" + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + served_model: Final = "gemini-3.8-flash-001" + iterator: Final = ModelResponseIterator( + streaming_response=iter( + [json.dumps({**_generate_content_body(), "modelVersion": served_model}) for _ in range(3)] + ), + sync_stream=True, + logging_obj=MagicMock(), + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(iterator), + model="gemini/gemini-3.8-flash", + custom_llm_provider="gemini", + logging_obj=MagicMock(), + ) + + chunks: Final = list(wrapper) + + assert len(chunks) >= 3 + for chunk in chunks[:-1]: + assert chunk._hidden_params["provider_response_model"] == served_model + assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}]) + assert assembled._hidden_params["provider_response_model"] == served_model From 44a6d1988958d3614f602d49746e83ae603ccb3b Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:14:37 +0000 Subject: [PATCH 3/4] test(gemini): drop review narration from the wrapper test docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini/test_vertex_and_google_ai_studio_gemini.py | 5 ----- 1 file changed, 5 deletions(-) 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 3fb00a7d962..cd0ed0057c0 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 @@ -5899,11 +5899,6 @@ def test_streaming_chunk_carries_model_version(): def test_served_model_version_reaches_assembled_stream_through_custom_stream_wrapper(): - """Greptile claimed CustomStreamWrapper drops the served modelVersion before - pricing. It does not: chunk_creator stashes the parser chunk's model into each - yielded chunk's _hidden_params["provider_response_model"], and - stream_chunk_builder carries it onto the assembled response. The wrapper's - trailing bookkeeping chunk is the one emission that legitimately omits it.""" from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, From c4c96180e301179259236953ab0a62e3008fcd95 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:42:36 +0000 Subject: [PATCH 4/4] fix(gemini): strip version suffix from modelVersion and keep it on blocked streams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 18 +++++++--- ...test_vertex_and_google_ai_studio_gemini.py | 33 +++++++++++++++++++ 2 files changed, 46 insertions(+), 5 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 eb03b17435c..b4712fd376b 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 @@ -124,6 +124,12 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr ) +def _served_model_name(model_version: object) -> str | None: + if not isinstance(model_version, str) or not model_version: + return None + return model_version.split("@", 1)[0] + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -1951,6 +1957,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _check_prompt_level_content_filter( processed_chunk: GenerateContentResponseBody, response_id: str | None, + model: str | None = None, ) -> Optional["ModelResponseStream"]: """ Check if prompt is blocked due to content filtering at the prompt level. @@ -1990,7 +1997,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): enhancements=None, ) - model_response: Final = ModelResponseStream(choices=[choice], id=response_id) + model_response: Final = ModelResponseStream(choices=[choice], id=response_id, model=model) return model_response return None @@ -2434,8 +2441,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_response = GenerateContentResponseBody(**completion_response) ## GET MODEL ## - model_version: Final = completion_response.get("modelVersion") - model_response.model = model_version if isinstance(model_version, str) else model + served: Final = _served_model_name(completion_response.get("modelVersion")) + model_response.model = served if served is not None else model ## CHECK IF RESPONSE FLAGGED if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]: @@ -3265,17 +3272,18 @@ class ModelResponseIterator: processed_chunk: Final = GenerateContentResponseBody(**chunk) response_id: Final = processed_chunk.get("responseId") - chunk_model_version: Final = processed_chunk.get("modelVersion") + served: Final = _served_model_name(processed_chunk.get("modelVersion")) model_response = ModelResponseStream( choices=[], id=response_id, - model=chunk_model_version if isinstance(chunk_model_version, str) else None, + model=served, ) # Check if prompt is blocked due to content filtering blocked_response: Final = VertexGeminiConfig._check_prompt_level_content_filter( processed_chunk=processed_chunk, response_id=response_id, + model=served, ) if blocked_response is not None: model_response = blocked_response 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 cd0ed0057c0..001105fc53d 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 @@ -5926,3 +5926,36 @@ def test_served_model_version_reaches_assembled_stream_through_custom_stream_wra assert chunk._hidden_params["provider_response_model"] == served_model assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}]) assert assembled._hidden_params["provider_response_model"] == served_model + + +def test_generate_content_transform_strips_version_suffix_from_model_version(): + import httpx + + body: Final = {**_generate_content_body(), "modelVersion": "gemini-3.8-flash-001@default"} + response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response=body, + model_response=ModelResponse(), + model="gemini-3.8-flash", + logging_obj=MagicMock(), + raw_response=httpx.Response(200, headers={}), + ) + + assert response.model == "gemini-3.8-flash-001" + + +def test_prompt_blocked_chunk_keeps_served_model_version(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk: Final = { + "promptFeedback": {"blockReason": "SAFETY", "blockReasonMessage": "prompt was blocked"}, + "modelVersion": "gemini-3.8-flash-001", + "responseId": "resp-1", + } + iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) + + streaming_chunk: Final = iterator.chunk_parser(chunk) + + assert streaming_chunk.model == "gemini-3.8-flash-001" + assert streaming_chunk.choices[0].finish_reason == "content_filter"