From 9efd35bf83ce8e0d1f0e7ae03d39eb3a97f753d7 Mon Sep 17 00:00:00 2001 From: Ayoub Date: Wed, 7 Oct 2026 10:08:24 +0100 Subject: [PATCH 1/3] Python: Honor the configured model in streamed Gemini chunks _process_chunk passed chunk.model_version straight into ChatResponseUpdate.model, while _process_generate_response falls back to self.model for the same field 33 lines above. Both take the same types.GenerateContentResponse, and model_version is Optional[str] in the pinned google-genai. ChatResponse.model is built only from the updates, so a stream whose chunks carry no model_version ended up with model=None. The observability layer reads that field for gen_ai.response.model, which is a metric dimension of the token and duration histograms, so those lost their per-model breakdown on streams. Apply the same expression the non-streaming path already uses. A chunk that carries model_version is untouched. --- .../agent_framework_gemini/_chat_client.py | 2 +- .../gemini/tests/test_gemini_client.py | 31 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 793dc891c4f..2ba8de533b2 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -1088,7 +1088,7 @@ def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUp return ChatResponseUpdate( contents=contents, - model=chunk.model_version, + model=chunk.model_version or self.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 cceb1d847b8..8fffe252ffb 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -547,6 +547,37 @@ 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_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") From 7e1373369a645f498651d17363b62214879d518e Mon Sep 17 00:00:00 2001 From: Ayoub Date: Wed, 7 Oct 2026 10:29:05 +0100 Subject: [PATCH 2/3] Python: Use the resolved request model as the streaming fallback _process_chunk fell back to self.model, which is the client default. A call that overrides the model through options was then labeled with the client default instead of the model that actually ran, including in the observability metrics. _prepare_request already resolves the model and it was in scope at the call site, so _process_chunk now takes it and uses it as the fallback. Adds a streaming test with an options override. --- .../agent_framework_gemini/_chat_client.py | 7 ++++--- .../packages/gemini/tests/test_gemini_client.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 2ba8de533b2..2c5a19ffa68 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -489,7 +489,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: contents=contents, config=config, ): - yield self._process_chunk(chunk) + yield self._process_chunk(chunk, model) except AgentFrameworkException: raise except Exception as ex: @@ -1063,13 +1063,14 @@ 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 resolved for this request, used when the chunk carries no model version. Returns: A ``ChatResponseUpdate`` with parsed contents, finish reason, and model ID. @@ -1088,7 +1089,7 @@ def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUp return ChatResponseUpdate( contents=contents, - model=chunk.model_version or self.model, + 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 8fffe252ffb..0d66628a6de 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -578,6 +578,23 @@ async def test_get_response_streaming_model_falls_back_to_configured_model() -> 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_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") From 3f9d32ce08303882dfe024b753ef889683eec20a Mon Sep 17 00:00:00 2001 From: Ayoub Date: Mon, 12 Oct 2026 14:07:51 +0100 Subject: [PATCH 3/3] Python: Scope the streaming model fallback to the stream The fallback was applied per chunk, so a chunk with no model_version reverted the update to the requested alias even after the server had already named a concrete version earlier in the same stream. ChatResponse.from_updates keeps the last non-None model, so a trailing usage chunk was enough to turn an observed gemini-2.5-pro-002 back into the requested gemini-2.5-pro in the aggregated response. The streaming loop now tracks the last model_version the server reported and uses it for later chunks that omit one, falling back to the resolved request model only until the server names a version. --- .../agent_framework_gemini/_chat_client.py | 10 ++++++++-- .../gemini/tests/test_gemini_client.py | 20 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 3859830ef8d..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, model) + 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: @@ -1078,7 +1083,8 @@ def _process_chunk(self, chunk: types.GenerateContentResponse, model: str) -> Ch Args: chunk: A streaming ``GenerateContentResponse`` chunk from the Gemini API. - model: The model resolved for this request, used when the chunk carries no model version. + 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. diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 79aa7a6f3b0..62f50e815a5 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -595,6 +595,26 @@ async def test_get_response_streaming_model_falls_back_to_requested_model() -> N 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")