feat: Support streaming function call arguments in progressive SSE streaming feature

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 837172244

src/google/adk/utils/streaming_utils.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import Any
 from typing import AsyncGenerator
 from typing import Optional
 
@@ -43,6 +44,12 @@ class StreamingResponseAggregator:
     self._current_text_is_thought: Optional[bool] = None
     self._finish_reason: Optional[types.FinishReason] = None
 
+    # For streaming function call arguments
+    self._current_fc_name: Optional[str] = None
+    self._current_fc_args: dict[str, Any] = {}
+    self._current_fc_id: Optional[str] = None
+    self._current_thought_signature: Optional[str] = None
+
   def _flush_text_buffer_to_sequence(self):
     """Flush current text buffer to parts sequence.
 
@@ -61,6 +68,171 @@ class StreamingResponseAggregator:
       self._current_text_buffer = ''
       self._current_text_is_thought = None
 
+  def _get_value_from_partial_arg(
+      self, partial_arg: types.PartialArg, json_path: str
+  ):
+    """Extract value from a partial argument.
+
+    Args:
+      partial_arg: The partial argument object
+      json_path: JSONPath for this argument
+
+    Returns:
+      Tuple of (value, has_value) where has_value indicates if a value exists
+    """
+    value = None
+    has_value = False
+
+    if partial_arg.string_value is not None:
+      # For streaming strings, append chunks to existing value
+      string_chunk = partial_arg.string_value
+      has_value = True
+
+      # Get current value for this path (if any)
+      path_without_prefix = (
+          json_path[2:] if json_path.startswith('$.') else json_path
+      )
+      path_parts = path_without_prefix.split('.')
+
+      # Try to get existing value
+      existing_value = self._current_fc_args
+      for part in path_parts:
+        if isinstance(existing_value, dict) and part in existing_value:
+          existing_value = existing_value[part]
+        else:
+          existing_value = None
+          break
+
+      # Append to existing string or set new value
+      if isinstance(existing_value, str):
+        value = existing_value + string_chunk
+      else:
+        value = string_chunk
+
+    elif partial_arg.number_value is not None:
+      value = partial_arg.number_value
+      has_value = True
+    elif partial_arg.bool_value is not None:
+      value = partial_arg.bool_value
+      has_value = True
+    elif partial_arg.null_value is not None:
+      value = None
+      has_value = True
+
+    return value, has_value
+
+  def _set_value_by_json_path(self, json_path: str, value: Any):
+    """Set a value in _current_fc_args using JSONPath notation.
+
+    Args:
+      json_path: JSONPath string like "$.location" or "$.location.latitude"
+      value: The value to set
+    """
+    # Remove leading "$." from jsonPath
+    if json_path.startswith('$.'):
+      path = json_path[2:]
+    else:
+      path = json_path
+
+    # Split path into components
+    path_parts = path.split('.')
+
+    # Navigate to the correct location and set the value
+    current = self._current_fc_args
+    for part in path_parts[:-1]:
+      if part not in current:
+        current[part] = {}
+      current = current[part]
+
+    # Set the final value
+    current[path_parts[-1]] = value
+
+  def _flush_function_call_to_sequence(self):
+    """Flush current function call to parts sequence.
+
+    This creates a complete FunctionCall part from accumulated partial args.
+    """
+    if self._current_fc_name:
+      # Create function call part with accumulated args
+      fc_part = types.Part.from_function_call(
+          name=self._current_fc_name,
+          args=self._current_fc_args.copy(),
+      )
+
+      # Set the ID if provided (directly on the function_call object)
+      if self._current_fc_id and fc_part.function_call:
+        fc_part.function_call.id = self._current_fc_id
+
+      # Set thought_signature if provided (on the Part, not FunctionCall)
+      if self._current_thought_signature:
+        fc_part.thought_signature = self._current_thought_signature
+
+      self._parts_sequence.append(fc_part)
+
+      # Reset FC state
+      self._current_fc_name = None
+      self._current_fc_args = {}
+      self._current_fc_id = None
+      self._current_thought_signature = None
+
+  def _process_streaming_function_call(self, fc: types.FunctionCall):
+    """Process a streaming function call with partialArgs.
+
+    Args:
+      fc: The function call object with partial_args
+    """
+    # Save function name if present (first chunk)
+    if fc.name:
+      self._current_fc_name = fc.name
+    if fc.id:
+      self._current_fc_id = fc.id
+
+    # Process each partial argument
+    for partial_arg in getattr(fc, 'partial_args', []):
+      json_path = partial_arg.json_path
+      if not json_path:
+        continue
+
+      # Extract value from partial arg
+      value, has_value = self._get_value_from_partial_arg(
+          partial_arg, json_path
+      )
+
+      # Set the value using JSONPath (only if a value was provided)
+      if has_value:
+        self._set_value_by_json_path(json_path, value)
+
+    # Check if function call is complete
+    fc_will_continue = getattr(fc, 'will_continue', False)
+    if not fc_will_continue:
+      # Function call complete, flush it
+      self._flush_text_buffer_to_sequence()
+      self._flush_function_call_to_sequence()
+
+  def _process_function_call_part(self, part: types.Part):
+    """Process a function call part (streaming or non-streaming).
+
+    Args:
+      part: The part containing a function call
+    """
+    fc = part.function_call
+
+    # Check if this is a streaming FC (has partialArgs)
+    if hasattr(fc, 'partial_args') and fc.partial_args:
+      # Streaming function call arguments
+
+      # Save thought_signature from the part (first chunk should have it)
+      if part.thought_signature and not self._current_thought_signature:
+        self._current_thought_signature = part.thought_signature
+      self._process_streaming_function_call(fc)
+    else:
+      # Non-streaming function call (standard format with args)
+      # Skip empty function calls (used as streaming end markers)
+      if fc.name:
+        # Flush any buffered text first, then add the FC part
+        self._flush_text_buffer_to_sequence()
+        self._parts_sequence.append(part)
+
   async def process_response(
       self, response: types.GenerateContentResponse
   ) -> AsyncGenerator[LlmResponse, None]:
@@ -101,8 +273,12 @@ class StreamingResponseAggregator:
             if not self._current_text_buffer:
               self._current_text_is_thought = part.thought
             self._current_text_buffer += part.text
+          elif part.function_call:
+            # Process function call (handles both streaming Args and
+            # non-streaming Args)
+            self._process_function_call_part(part)
           else:
-            # Non-text part (function_call, bytes, etc.)
+            # Other non-text parts (bytes, etc.)
             # Flush any buffered text first, then add the non-text part
             self._flush_text_buffer_to_sequence()
             self._parts_sequence.append(part)
@@ -155,8 +331,9 @@ class StreamingResponseAggregator:
     if is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING):
       # Always generate final aggregated response in progressive mode
       if self._response and self._response.candidates:
-        # Flush any remaining text buffer to complete the sequence
+        # Flush any remaining buffers to complete the sequence
         self._flush_text_buffer_to_sequence()
+        self._flush_function_call_to_sequence()
 
         # Use the parts sequence which preserves original ordering
         final_parts = self._parts_sequence