Skip to content
12 changes: 5 additions & 7 deletions src/google/adk/flows/llm_flows/_output_schema_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,13 @@ def create_final_model_response_event(


def get_structured_model_response(function_response_event: Event) -> str | None:
"""Check if function response contains set_model_response and extract JSON.
"""Check if function response contains a validated set_model_response result.

Args:
function_response_event: The function response event to check.

Returns:
JSON response string if set_model_response was called, None otherwise.
JSON response string if set_model_response succeeded, None otherwise.
"""
if (
not function_response_event
Expand All @@ -110,11 +110,9 @@ def get_structured_model_response(function_response_event: Event) -> str | None:

for func_response in function_response_event.get_function_responses():
if func_response.name == 'set_model_response':
# Extract the actual result from the wrapped response.
# Tool results are wrapped as {'result': ...} when not already a dict.
response = func_response.response
if isinstance(response, dict) and 'result' in response:
response = response['result']
response = function_response_event.actions.set_model_response
if response is None:
return None
return json.dumps(response, ensure_ascii=False)

return None
Expand Down
47 changes: 30 additions & 17 deletions src/google/adk/tools/set_model_response_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

from google.genai import types
from pydantic import TypeAdapter
from pydantic import ValidationError
from typing_extensions import override

from ..utils._schema_utils import get_list_inner_type
Expand Down Expand Up @@ -150,27 +151,39 @@ async def run_async(
tool_context: Tool execution context.

Returns:
The validated response. Type depends on the output_schema:
The validated response, or validation feedback for the model to retry.
Type depends on the output_schema:
- dict for BaseModel
- list of dicts for list[BaseModel]
- raw value for other schema types (list[str], dict, etc.)
- dict with an error message when Pydantic validation fails
"""
if self._is_basemodel:
# For regular BaseModel, validate directly
validated_response = self.output_schema.model_validate(args)
result = validated_response.model_dump(exclude_none=True)
elif self._is_list_of_basemodel:
# For list[BaseModel], extract and validate the 'items' field
items = args.get('items', [])
type_adapter = TypeAdapter(self.output_schema)
validated_response = type_adapter.validate_python(items)
result = [
item.model_dump(exclude_none=True) for item in validated_response
]
else:
# For other schema types (list[str], dict, etc.),
# return the value directly without pydantic validation
result = args.get('response')
try:
if self._is_basemodel:
# For regular BaseModel, validate directly
validated_response = self.output_schema.model_validate(args)
result = validated_response.model_dump(exclude_none=True)
elif self._is_list_of_basemodel:
# For list[BaseModel], extract and validate the 'items' field
items = args.get('items', [])
type_adapter = TypeAdapter(self.output_schema)
validated_response = type_adapter.validate_python(items)
result = [
item.model_dump(exclude_none=True) for item in validated_response
]
else:
# For other schema types (list[str], dict, etc.),
# return the value directly without pydantic validation
result = args.get('response')
except ValidationError as e:
return {
'error': (
f'Validation Error found:\n{e}\n'
'Recall the set_model_response function correctly, fix the'
' errors, and call it again with all required fields using the'
' correct types.'
)
}

tool_context.actions.set_model_response = result
return result
93 changes: 93 additions & 0 deletions tests/unittests/flows/llm_flows/test_output_schema_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.run_config import RunConfig
from google.adk.events.event_actions import EventActions
from google.adk.flows.llm_flows.single_flow import SingleFlow
from google.adk.models.llm_request import LlmRequest
from google.adk.sessions.in_memory_session_service import InMemorySessionService
Expand Down Expand Up @@ -245,6 +246,7 @@ async def test_output_schema_helper_functions():
# Create a function response event with set_model_response
function_response_event = Event(
author='test_agent',
actions=EventActions(set_model_response=test_dict),
content=types.Content(
role='user',
parts=[
Expand Down Expand Up @@ -302,6 +304,7 @@ async def test_get_structured_model_response_with_non_ascii():
# Create a function response event
function_response_event = Event(
author='test_agent',
actions=EventActions(set_model_response=test_dict),
content=types.Content(
role='user',
parts=[
Expand Down Expand Up @@ -344,6 +347,7 @@ async def test_get_structured_model_response_with_wrapped_result():
# Create a function response event with wrapped result
function_response_event = Event(
author='test_agent',
actions=EventActions(set_model_response=wrapped_response['result']),
content=types.Content(
role='user',
parts=[
Expand All @@ -363,6 +367,38 @@ async def test_get_structured_model_response_with_wrapped_result():
assert extracted_json == expected_json


@pytest.mark.asyncio
async def test_get_structured_model_response_skips_error_response():
"""Test set_model_response error payloads are not treated as final output."""
from google.adk.events.event import Event
from google.adk.flows.llm_flows._output_schema_processor import get_structured_model_response
from google.genai import types

function_response_event = Event(
author='test_agent',
content=types.Content(
role='user',
parts=[
types.Part(
function_response=types.FunctionResponse(
name='set_model_response',
response={
'error': (
'Validation Error found:\nage\n'
'Input should be a valid integer'
)
},
)
)
],
),
)

extracted_json = get_structured_model_response(function_response_event)

assert extracted_json is None


@pytest.mark.asyncio
async def test_end_to_end_integration():
"""Test the complete output schema with tools integration."""
Expand Down Expand Up @@ -468,6 +504,63 @@ async def test_flow_yields_both_events_for_set_model_response():
)


@pytest.mark.asyncio
async def test_flow_yields_error_response_for_invalid_set_model_response():
"""Test invalid set_model_response args are sent back without finalizing."""
from google.adk.events.event import Event
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
from google.adk.tools.set_model_response_tool import SetModelResponseTool
from google.genai import types

agent = LlmAgent(
name='test_agent',
model='gemini-2.5-flash',
output_schema=PersonSchema,
tools=[],
)

invocation_context = await _create_invocation_context(agent)
flow = BaseLlmFlow()

set_response_tool = SetModelResponseTool(PersonSchema)
llm_request = LlmRequest()
llm_request.tools_dict['set_model_response'] = set_response_tool

function_call_event = Event(
author='test_agent',
content=types.Content(
role='model',
parts=[
types.Part(
function_call=types.FunctionCall(
name='set_model_response',
args={
'name': 'Test User',
'age': 'not-an-int',
# Missing city.
},
)
)
],
),
)

events = []
async for event in flow._postprocess_handle_function_calls_async(
invocation_context, function_call_event, llm_request
):
events.append(event)

assert len(events) == 1
function_response = events[0].get_function_responses()[0]
assert function_response.name == 'set_model_response'
assert 'error' in function_response.response
assert 'Validation Error found' in function_response.response['error']
assert 'age' in function_response.response['error']
assert 'city' in function_response.response['error']
assert events[0].actions.set_model_response is None


@pytest.mark.asyncio
async def test_flow_yields_only_function_response_for_normal_tools():
"""Test that the flow yields only function response event for non-set_model_response tools."""
Expand Down
66 changes: 43 additions & 23 deletions tests/unittests/tools/test_set_model_response_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@
from google.genai import types
from pydantic import BaseModel
from pydantic import Field
from pydantic import ValidationError
import pytest


Expand Down Expand Up @@ -161,40 +160,53 @@ async def test_run_async_complex_schema():
assert result['tags'] == ['tag1', 'tag2']
assert result['metadata'] == {'key': 'value'}
assert result['is_active'] is False
assert tool_context.actions.set_model_response == result


@pytest.mark.asyncio
async def test_run_async_validation_error():
"""Test tool execution with invalid data raises validation error."""
"""Test tool execution with invalid data returns validation feedback."""
tool = SetModelResponseTool(PersonSchema)

agent = LlmAgent(name='test_agent', model='gemini-2.5-flash')
invocation_context = await _create_invocation_context(agent)
tool_context = ToolContext(invocation_context)

# Execute with invalid data (wrong type for age)
with pytest.raises(ValidationError):
await tool.run_async(
args={'name': 'Bob', 'age': 'not_a_number', 'city': 'Portland'},
tool_context=tool_context,
)
result = await tool.run_async(
args={'name': 'Bob', 'age': 'not_a_number', 'city': 'Portland'},
tool_context=tool_context,
)

assert result is not None
assert 'error' in result
assert 'Validation Error found' in result['error']
assert 'age' in result['error']
assert 'int_parsing' in result['error']
assert tool_context.actions.set_model_response is None


@pytest.mark.asyncio
async def test_run_async_missing_required_field():
"""Test tool execution with missing required field."""
"""Test tool execution with missing required field returns feedback."""
tool = SetModelResponseTool(PersonSchema)

agent = LlmAgent(name='test_agent', model='gemini-2.5-flash')
invocation_context = await _create_invocation_context(agent)
tool_context = ToolContext(invocation_context)

# Execute with missing required field
with pytest.raises(ValidationError):
await tool.run_async(
args={'name': 'Charlie', 'city': 'Denver'}, # Missing age
tool_context=tool_context,
)
result = await tool.run_async(
args={'name': 'Charlie', 'city': 'Denver'}, # Missing age
tool_context=tool_context,
)

assert result is not None
assert 'error' in result
assert 'Validation Error found' in result['error']
assert 'age' in result['error']
assert 'Field required' in result['error']
assert tool_context.actions.set_model_response is None


@pytest.mark.asyncio
Expand All @@ -216,6 +228,7 @@ async def test_session_state_storage_key():
assert result['name'] == 'Diana'
assert result['age'] == 35
assert result['city'] == 'Miami'
assert tool_context.actions.set_model_response == result


@pytest.mark.asyncio
Expand Down Expand Up @@ -357,27 +370,34 @@ async def test_run_async_list_schema_empty_list():
assert result is not None
assert isinstance(result, list)
assert len(result) == 0
assert tool_context.actions.set_model_response == result


@pytest.mark.asyncio
async def test_run_async_list_schema_validation_error():
"""Test tool execution with invalid list data raises validation error."""
"""Test tool execution with invalid list data returns validation feedback."""
tool = SetModelResponseTool(list[ItemSchema])

agent = LlmAgent(name='test_agent', model='gemini-2.5-flash')
invocation_context = await _create_invocation_context(agent)
tool_context = ToolContext(invocation_context)

# Execute with invalid data (wrong type for id)
with pytest.raises(ValidationError):
await tool.run_async(
args={
'items': [
{'id': 'not_a_number', 'name': 'Item 1'},
]
},
tool_context=tool_context,
)
result = await tool.run_async(
args={
'items': [
{'id': 'not_a_number', 'name': 'Item 1'},
]
},
tool_context=tool_context,
)

assert result is not None
assert 'error' in result
assert 'Validation Error found' in result['error']
assert '0.id' in result['error']
assert 'int_parsing' in result['error']
assert tool_context.actions.set_model_response is None


# Tests for other schema types (list[str], dict, etc.)
Expand Down
Loading