From c4c96180e301179259236953ab0a62e3008fcd95 Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Wed, 16 Sep 2026 01:42:36 +0000
Subject: [PATCH] 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"