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"