fix: append the current model chunk to contents in async streaming

Author: dechantoine

google/genai/_extra_utils.py

@@ -536,13 +536,14 @@ async def parse_config_for_mcp_sessions(
 def append_chunk_contents(
  contents: Union[types.ContentListUnion, types.ContentListUnionDict],
  chunk: types.GenerateContentResponse,
-) -> None:
- """Appends the contents of the chunk to the contents list."""
+) -> Union[types.ContentListUnion, types.ContentListUnionDict]:
+ """Appends the contents of the chunk to the contents list and returns it."""
  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]
+ return contents
 
 
 def prepare_resumable_upload(

google/genai/models.py

@@ -5168,7 +5168,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)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
  else:
  if (
@@ -5181,7 +5181,7 @@ def generate_content_stream(
  chunk, function_map
  )
  if not func_response_parts:
- _extra_utils.append_chunk_contents(contents, chunk)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
 
  else:
@@ -5191,7 +5191,7 @@ def generate_content_stream(
  chunk.automatic_function_calling_history = (
  automatic_function_calling_history
  )
- _extra_utils.append_chunk_contents(contents, chunk)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
  if (
  chunk is None
@@ -6938,7 +6938,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)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
  else:
  if (
@@ -6953,7 +6953,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)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
 
  else:
@@ -6964,7 +6964,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)
+ contents = _extra_utils.append_chunk_contents(contents, chunk)
  yield chunk
  if (
  chunk is None