From 3f3aa7b32d63cae5750d71bc586c088427c979ea Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Thu, 21 Aug 2025 11:24:05 -0700 Subject: [PATCH] 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 --- src/google/adk/flows/llm_flows/_output_schema_processor.py | 6 +++++- .../flows/llm_flows/test_output_schema_processor.py | 4 ++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/google/adk/flows/llm_flows/_output_schema_processor.py b/src/google/adk/flows/llm_flows/_output_schema_processor.py index 16638702..da793f8f 100644 --- a/src/google/adk/flows/llm_flows/_output_schema_processor.py +++ b/src/google/adk/flows/llm_flows/_output_schema_processor.py @@ -78,7 +78,11 @@ def create_final_model_response_event( from google.genai import types # 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( role='model', parts=[types.Part(text=json_response)] ) diff --git a/tests/unittests/flows/llm_flows/test_output_schema_processor.py b/tests/unittests/flows/llm_flows/test_output_schema_processor.py index 42bfa880..70eedcc6 100644 --- a/tests/unittests/flows/llm_flows/test_output_schema_processor.py +++ b/tests/unittests/flows/llm_flows/test_output_schema_processor.py @@ -230,6 +230,8 @@ async def test_output_schema_helper_functions(): # Test create_final_model_response_event function final_event = create_final_model_response_event(invocation_context, test_json) 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.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 = events[1] 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.parts[0].text