diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 402934bf26d..2f0dceb9528 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -492,8 +492,13 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # The SDK stream owns the HTTP response; close it as soon as the consumer stops early # instead of leaving the response open until the generator is garbage collected. try: + # Once the server names a version, later chunks that omit it keep that version + # instead of reverting to the requested alias. + fallback_model = model async for chunk in response_stream: - yield self._process_chunk(chunk) + if chunk.model_version: + fallback_model = chunk.model_version + yield self._process_chunk(chunk, fallback_model) finally: aclose = getattr(response_stream, "aclose", None) if aclose is not None: @@ -1071,13 +1076,15 @@ def _process_generate_response( raw_representation=response, ) - def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUpdate: + def _process_chunk(self, chunk: types.GenerateContentResponse, model: str) -> ChatResponseUpdate: """Convert a single streaming chunk to a framework ChatResponseUpdate. Usage details are attached only to the final chunk, identified by a non-None finish reason. Args: chunk: A streaming ``GenerateContentResponse`` chunk from the Gemini API. + model: The model to report when the chunk carries no model version: the last version the + server reported in this stream, or the model resolved for this request. Returns: A ``ChatResponseUpdate`` with parsed contents, finish reason, and model ID. @@ -1096,7 +1103,7 @@ def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUp return ChatResponseUpdate( contents=contents, - model=chunk.model_version, + model=chunk.model_version or model, finish_reason=finish_reason, raw_representation=chunk, ) diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 2e3ffbfbe70..62f50e815a5 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -547,6 +547,74 @@ async def test_get_response_model_from_response() -> None: assert response.model == "gemini-2.5-pro-002" +async def test_get_response_streaming_model_from_chunk() -> None: + """Populates the update model from the model_version field on each chunk.""" + client, mock = _make_gemini_client() + chunks = [_make_response([_make_part(text="Hi")], model_version="gemini-2.5-pro-002")] + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks)) + + stream = client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Hi")])], + stream=True, + ) + updates = [update async for update in stream] + + assert updates[0].model == "gemini-2.5-pro-002" + + +async def test_get_response_streaming_model_falls_back_to_configured_model() -> None: + """Falls back to the configured model when a chunk carries no model_version.""" + client, mock = _make_gemini_client() + chunk = _make_response([_make_part(text="Hi")]) + chunk.model_version = None + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([chunk])) + + stream = client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Hi")])], + stream=True, + ) + updates = [update async for update in stream] + + assert updates[0].model == "gemini-2.5-flash" + + +async def test_get_response_streaming_model_falls_back_to_requested_model() -> None: + """Falls back to the model requested in options, not the client default.""" + client, mock = _make_gemini_client(model="gemini-2.5-flash") + chunk = _make_response([_make_part(text="Hi")]) + chunk.model_version = None + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([chunk])) + + stream = client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Hi")])], + options={"model": "gemini-2.5-pro"}, + stream=True, + ) + updates = [update async for update in stream] + + assert updates[0].model == "gemini-2.5-pro" + + +async def test_get_response_streaming_keeps_model_version_seen_earlier() -> None: + """A later chunk without model_version must not replace a version already seen in the stream.""" + client, mock = _make_gemini_client(model="gemini-2.5-flash") + first = _make_response([_make_part(text="Hi")], finish_reason=None, model_version="gemini-2.5-pro-002") + last = _make_response([_make_part(text="!")], finish_reason="STOP") + last.model_version = None + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([first, last])) + + stream = client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Hi")])], + options={"model": "gemini-2.5-pro"}, + stream=True, + ) + updates = [update async for update in stream] + final = await stream.get_final_response() + + assert [update.model for update in updates] == ["gemini-2.5-pro-002", "gemini-2.5-pro-002"] + assert final.model == "gemini-2.5-pro-002" + + async def test_get_response_uses_model_from_options() -> None: """Uses the model specified in options, overriding the client's default.""" client, mock = _make_gemini_client(model="gemini-2.5-flash")