scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(gemini): assert the served modelVersion reaches the assembled stream through CustomStreamWrapper

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-16 01:05:53 +00:00
parent d5a36eb2ca
commit 801f6a92b4
1 changed files with 36 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

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