From 279e4fedd0b1c0d1499c0f9a4454357af7da490e Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Fri, 15 Aug 2025 16:37:42 -0700 Subject: [PATCH] fix: Using base event's invocation id when merge multiple function response event fix https://github.com/google/adk-python/issues/1531 PiperOrigin-RevId: 795657473 --- src/google/adk/flows/llm_flows/functions.py | 2 +- .../flows/llm_flows/test_functions_simple.py | 116 ++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 0c8fa86a..b0700270 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -682,7 +682,7 @@ def merge_parallel_function_response_events( # Create the new merged event merged_event = Event( - invocation_id=Event.new_id(), + invocation_id=base_event.invocation_id, author=base_event.author, branch=base_event.branch, content=types.Content(role='user', parts=merged_parts), diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index 166800cc..b3aab2b6 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -19,6 +19,7 @@ from typing import Callable from google.adk.agents.llm_agent import Agent from google.adk.events.event import Event from google.adk.flows.llm_flows.functions import find_matching_function_call +from google.adk.flows.llm_flows.functions import merge_parallel_function_response_events from google.adk.tools.function_tool import FunctionTool from google.adk.tools.tool_context import ToolContext from google.genai import types @@ -931,6 +932,121 @@ async def test_async_function_without_yield_blocks_others(): assert execution_order == ['non_yield_A', 'non_yield_B', 'yield_C', 'yield_D'] +def test_merge_parallel_function_response_events_preserves_invocation_id(): + """Test that merge_parallel_function_response_events preserves the base event's invocation_id.""" + # Create multiple function response events with different invocation IDs + invocation_id = 'base_invocation_123' + + function_response1 = types.FunctionResponse( + id='func_123', name='test_function1', response={'result': 'success1'} + ) + + function_response2 = types.FunctionResponse( + id='func_456', name='test_function2', response={'result': 'success2'} + ) + + event1 = Event( + invocation_id=invocation_id, + author='test_agent', + content=types.Content( + role='user', parts=[types.Part(function_response=function_response1)] + ), + ) + + event2 = Event( + invocation_id='different_invocation_456', # Different invocation ID + author='test_agent', + content=types.Content( + role='user', parts=[types.Part(function_response=function_response2)] + ), + ) + + # Merge the events + merged_event = merge_parallel_function_response_events([event1, event2]) + + # Should preserve the base event's (first event's) invocation_id + assert merged_event.invocation_id == invocation_id + assert merged_event.invocation_id != 'different_invocation_456' + + # Should contain both function responses + assert len(merged_event.content.parts) == 2 + + # Verify the responses are preserved + response_ids = { + part.function_response.id for part in merged_event.content.parts + } + assert 'func_123' in response_ids + assert 'func_456' in response_ids + + +def test_merge_parallel_function_response_events_single_event(): + """Test that merge_parallel_function_response_events returns single event unchanged.""" + invocation_id = 'single_invocation_123' + + function_response = types.FunctionResponse( + id='func_123', name='test_function', response={'result': 'success'} + ) + + event = Event( + invocation_id=invocation_id, + author='test_agent', + content=types.Content( + role='user', parts=[types.Part(function_response=function_response)] + ), + ) + + # Merge single event + merged_event = merge_parallel_function_response_events([event]) + + # Should return the same event object + assert merged_event is event + assert merged_event.invocation_id == invocation_id + + +def test_merge_parallel_function_response_events_preserves_other_attributes(): + """Test that merge_parallel_function_response_events preserves other attributes from base event.""" + invocation_id = 'base_invocation_123' + base_author = 'base_agent' + base_branch = 'main_branch' + + function_response1 = types.FunctionResponse( + id='func_123', name='test_function1', response={'result': 'success1'} + ) + + function_response2 = types.FunctionResponse( + id='func_456', name='test_function2', response={'result': 'success2'} + ) + + event1 = Event( + invocation_id=invocation_id, + author=base_author, + branch=base_branch, + content=types.Content( + role='user', parts=[types.Part(function_response=function_response1)] + ), + ) + + event2 = Event( + invocation_id='different_invocation_456', + author='different_agent', # Different author + branch='different_branch', # Different branch + content=types.Content( + role='user', parts=[types.Part(function_response=function_response2)] + ), + ) + + # Merge the events + merged_event = merge_parallel_function_response_events([event1, event2]) + + # Should preserve base event's attributes + assert merged_event.invocation_id == invocation_id + assert merged_event.author == base_author + assert merged_event.branch == base_branch + + # Should contain both function responses + assert len(merged_event.content.parts) == 2 + + @pytest.mark.asyncio async def test_yielding_async_functions_run_concurrently(): """Test that async functions with proper yields run concurrently."""