From 801f6a92b4a2d0e55daa738ff9fd3db0098382c6 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:05:53 +0000 Subject: [PATCH] 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