Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/google/adk/flows/llm_flows/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
116 changes: 116 additions & 0 deletions tests/unittests/flows/llm_flows/test_functions_simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
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
Expand Down Expand Up @@ -931,6 +932,121 @@ async def yielding_async_function() -> dict:
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."""
Expand Down