fix: Set invocation_id and branch for event generated when both output_schema and tools are used

fixes https://github.com/google/adk-python/issues/2631

PiperOrigin-RevId: 797843460
This commit is contained in:
Xiang (Sean) Zhou
2025-08-21 11:24:43 -07:00
committed by Copybara-Service
parent 826f554789
commit 3f3aa7b32d
2 changed files with 9 additions and 1 deletions
@@ -78,7 +78,11 @@ def create_final_model_response_event(
from google.genai import types from google.genai import types
# Create a proper model response event # Create a proper model response event
final_event = Event(author=invocation_context.agent.name) final_event = Event(
author=invocation_context.agent.name,
invocation_id=invocation_context.invocation_id,
branch=invocation_context.branch,
)
final_event.content = types.Content( final_event.content = types.Content(
role='model', parts=[types.Part(text=json_response)] role='model', parts=[types.Part(text=json_response)]
) )
@@ -230,6 +230,8 @@ async def test_output_schema_helper_functions():
# Test create_final_model_response_event function # Test create_final_model_response_event function
final_event = create_final_model_response_event(invocation_context, test_json) final_event = create_final_model_response_event(invocation_context, test_json)
assert final_event.author == 'test_agent' assert final_event.author == 'test_agent'
assert final_event.invocation_id == invocation_context.invocation_id
assert final_event.branch == invocation_context.branch
assert final_event.content.role == 'model' assert final_event.content.role == 'model'
assert final_event.content.parts[0].text == test_json assert final_event.content.parts[0].text == test_json
@@ -348,6 +350,8 @@ async def test_flow_yields_both_events_for_set_model_response():
# Second event should be the final model response with JSON # Second event should be the final model response with JSON
second_event = events[1] second_event = events[1]
assert second_event.author == 'test_agent' assert second_event.author == 'test_agent'
assert second_event.invocation_id == invocation_context.invocation_id
assert second_event.branch == invocation_context.branch
assert second_event.content.role == 'model' assert second_event.content.role == 'model'
assert ( assert (
second_event.content.parts[0].text second_event.content.parts[0].text