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 e426801518..314a480118 100644 --- a/src/google/adk/flows/llm_flows/_output_schema_processor.py +++ b/src/google/adk/flows/llm_flows/_output_schema_processor.py @@ -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 @@ -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 diff --git a/src/google/adk/tools/set_model_response_tool.py b/src/google/adk/tools/set_model_response_tool.py index 3e80ba64f7..0203efd940 100644 --- a/src/google/adk/tools/set_model_response_tool.py +++ b/src/google/adk/tools/set_model_response_tool.py @@ -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 @@ -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 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 9bde17344a..7cc30db2be 100644 --- a/tests/unittests/flows/llm_flows/test_output_schema_processor.py +++ b/tests/unittests/flows/llm_flows/test_output_schema_processor.py @@ -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 @@ -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=[ @@ -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=[ @@ -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=[ @@ -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.""" @@ -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.""" diff --git a/tests/unittests/tools/test_set_model_response_tool.py b/tests/unittests/tools/test_set_model_response_tool.py index 54ff459eea..507e97a1b0 100644 --- a/tests/unittests/tools/test_set_model_response_tool.py +++ b/tests/unittests/tools/test_set_model_response_tool.py @@ -28,7 +28,6 @@ from google.genai import types from pydantic import BaseModel from pydantic import Field -from pydantic import ValidationError import pytest @@ -161,11 +160,12 @@ 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') @@ -173,16 +173,22 @@ async def test_run_async_validation_error(): 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') @@ -190,11 +196,17 @@ async def test_run_async_missing_required_field(): 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 @@ -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 @@ -357,11 +370,12 @@ 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') @@ -369,15 +383,21 @@ async def test_run_async_list_schema_validation_error(): 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.)