fix: Keep chunk content history when thought summaries are enabled in the stream response

PiperOrigin-RevId: 775383091

Author: Annhiluc

google/genai/_extra_utils.py

@@ -25,6 +25,7 @@
 
 from . import _common
 from . import _mcp_utils
+from . import _transformers as t
 from . import errors
 from . import types
 from ._adapters import McpToGenAiToolAdapter
@@ -506,3 +507,15 @@ async def parse_config_for_mcp_sessions(
         parsed_config_copy.tools.append(tool)
 
   return parsed_config_copy, mcp_to_genai_tool_adapters
+
+
+def append_chunk_contents(
+    contents: Union[types.ContentListUnion, types.ContentListUnionDict],
+    chunk: types.GenerateContentResponse,
+) -> None:
+  """Appends the contents of the chunk to the contents list."""
+  if chunk is not None and chunk.candidates is not None:
+    chunk_content = chunk.candidates[0].content
+    contents = t.t_contents(contents)  # type: ignore[assignment]
+    if isinstance(contents, list) and chunk_content is not None:
+      contents.append(chunk_content)  # type: ignore[arg-type]

google/genai/models.py

@@ -6018,6 +6018,7 @@ def generate_content_stream(
         # Yield chunks only if there's no function response parts.
         for chunk in response:
           if not function_map:
+            _extra_utils.append_chunk_contents(contents, chunk)
             yield chunk
           else:
             if (
@@ -6030,6 +6031,7 @@ def generate_content_stream(
                 chunk, function_map
             )
             if not func_response_parts:
+              _extra_utils.append_chunk_contents(contents, chunk)
               yield chunk
 
       else:
@@ -6039,6 +6041,7 @@ def generate_content_stream(
             chunk.automatic_function_calling_history = (
                 automatic_function_calling_history
             )
+          _extra_utils.append_chunk_contents(contents, chunk)
           yield chunk
         if (
             chunk is None
@@ -7547,6 +7550,7 @@ async def async_generator(model, contents, config):  # type: ignore[no-untyped-d
           # Yield chunks only if there's no function response parts.
           async for chunk in response:  # type: ignore[attr-defined]
             if not function_map:
+              _extra_utils.append_chunk_contents(contents, chunk)
               yield chunk
             else:
               if (
@@ -7561,6 +7565,7 @@ async def async_generator(model, contents, config):  # type: ignore[no-untyped-d
                   )
               )
               if not func_response_parts:
+                _extra_utils.append_chunk_contents(contents, chunk)
                 yield chunk
 
         else:
@@ -7571,6 +7576,7 @@ async def async_generator(model, contents, config):  # type: ignore[no-untyped-d
               chunk.automatic_function_calling_history = (
                   automatic_function_calling_history
               )
+            _extra_utils.append_chunk_contents(contents, chunk)
             yield chunk
           if (
               chunk is None

google/genai/tests/afc/test_generate_content_stream_afc.py

@@ -326,6 +326,39 @@ def test_generate_content_stream_with_function_tools_used(
     ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
 
 
+def test_generate_content_stream_with_thought_summaries(
+    mock_generate_content_stream_with_afc,
+    mock_get_function_response_parts,
+):
+  """Test when function tools are provided and thought summaries are enabled.
+
+  Expected to answer weather based on function response.
+  """
+  models_instance = models.Models(api_client_=mock_api_client)
+  config = types.GenerateContentConfig(
+      tools=[get_current_weather],
+      thinking_config=types.ThinkingConfig(include_thoughts=True),
+  )
+  stream = models_instance.generate_content_stream(
+      model='test_model',
+      contents='what is the weather in San Francisco?',
+      config=config,
+  )
+
+  chunk = None
+  for chunk in stream:
+    assert chunk.text == TEST_AFC_TEXT_PART.text
+
+  assert mock_generate_content_stream_with_afc.call_count == 2
+  assert mock_get_function_response_parts.call_count == 2
+
+  assert chunk is not None
+  for i in range(len(chunk.automatic_function_calling_history)):
+    assert chunk.automatic_function_calling_history[i].model_dump(
+        exclude_none=True
+    ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
+
+
 @pytest.mark.asyncio
 async def test_generate_content_stream_no_function_map_async(
     mock_generate_content_stream_no_afc,
@@ -460,3 +493,38 @@ async def test_generate_content_stream_with_function_async_function_used_async(
     assert chunk.automatic_function_calling_history[i].model_dump(
         exclude_none=True
     ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
+
+
+@pytest.mark.asyncio
+async def test_generate_content_stream_with_thought_summaries_async(
+    mock_generate_content_stream_with_afc_async,
+    mock_get_function_response_parts_async,
+):
+  """Test when function tools are provided and thought summaries are enabled.
+
+  Expected to answer weather based on function response.
+  """
+  models_instance = models.AsyncModels(api_client_=mock_api_client)
+  config = types.GenerateContentConfig(
+      tools=[get_current_weather],
+      thinking_config=types.ThinkingConfig(include_thoughts=True),
+  )
+  stream = await models_instance.generate_content_stream(
+      model='test_model',
+      contents='what is the weather in San Francisco?',
+      config=config,
+  )
+
+  chunk = None
+  async for chunk in stream:
+    assert chunk.text == TEST_AFC_TEXT_PART.text
+
+  assert mock_generate_content_stream_with_afc_async.call_count == 2
+
+  assert mock_get_function_response_parts_async.call_count == 2
+
+  assert chunk is not None
+  for i in range(len(chunk.automatic_function_calling_history)):
+    assert chunk.automatic_function_calling_history[i].model_dump(
+        exclude_none=True
+    ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)

google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py

@@ -0,0 +1,77 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import pytest
+from ... import types
+from .. import pytest_helper
+
+
+def get_current_weather(location: str) -> str:
+  """Returns the current weather.
+
+  Args:
+    location: The location of a city and state, e.g. "San Francisco, CA".
+  """
+  return 'windy'
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.generate_content',
+)
+pytest_plugins = ('pytest_asyncio',)
+
+
+def test_generate_content_stream_with_function_and_thought_summaries(client):
+  """Test when function tools are provided and thought summaries are enabled.
+
+  Expected to answer weather based on function response.
+  """
+  config = types.GenerateContentConfig(
+      tools=[get_current_weather],
+      thinking_config=types.ThinkingConfig(include_thoughts=True),
+  )
+  stream = client.models.generate_content_stream(
+      model='gemini-2.5-flash',
+      contents='what is the weather in San Francisco, CA?',
+      config=config,
+  )
+
+  chunk = None
+  for chunk in stream:
+    assert chunk is not None
+
+
+@pytest.mark.asyncio
+async def test_generate_content_stream_with_function_and_thought_summaries_async(
+    client,
+):
+  """Test when function tools are provided and thought summaries are enabled.
+
+  Expected to answer weather based on function response.
+  """
+  config = types.GenerateContentConfig(
+      tools=[get_current_weather],
+      thinking_config=types.ThinkingConfig(include_thoughts=True),
+  )
+  stream = await client.aio.models.generate_content_stream(
+      model='gemini-2.5-flash',
+      contents='what is the weather in San Francisco, CA?',
+      config=config,
+  )
+
+  chunk = None
+  async for chunk in stream:
+    assert chunk is not None