diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index fdc4b2375f..80acb9b2c0 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -1189,6 +1189,10 @@ def __build_response_event( parts=function_response_parts, ) part_function_response.function_response.id = tool_context.function_call_id + if tool.response_scheduling is not None: + part_function_response.function_response.scheduling = ( + tool.response_scheduling + ) content = types.Content( role='user', diff --git a/src/google/adk/tools/base_tool.py b/src/google/adk/tools/base_tool.py index 9139c23d57..fc14715e76 100644 --- a/src/google/adk/tools/base_tool.py +++ b/src/google/adk/tools/base_tool.py @@ -90,6 +90,19 @@ class BaseTool(ABC): NOTE: the entire dict must be JSON serializable. """ + response_scheduling: Optional[types.FunctionResponseScheduling] = None + """Controls when the model reacts to the tool's response. + + This is primarily used for Live (bidi) streaming: + - ``SILENT``: feeds the response back without triggering a model turn. + - ``WHEN_IDLE``: defers the reaction until the model is idle. + - ``INTERRUPT``: reacts immediately. + + When set, this value is applied to the emitted ``FunctionResponse``. It is + ignored by models that don't support response scheduling. + ``None`` preserves the default behavior. + """ + def __init__( self, *, @@ -97,12 +110,14 @@ def __init__( description, is_long_running: bool = False, custom_metadata: Optional[dict[str, Any]] = None, + response_scheduling: Optional[types.FunctionResponseScheduling] = None, ): self.name = name self.description = description self.is_long_running = is_long_running self._defers_response = False self.custom_metadata = custom_metadata + self.response_scheduling = response_scheduling def _get_declaration(self) -> Optional[types.FunctionDeclaration]: """Gets the OpenAPI specification of this tool in the form of a FunctionDeclaration. diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index 0663c1fe5a..b36b70e429 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -1559,3 +1559,64 @@ async def test_detection_exception_does_not_break_tool_call( assert len(recorded_calls) == 1 assert recorded_calls[0]['error_type'] is None assert recorded_calls[0]['error'] is None + + +@pytest.mark.asyncio +async def test_response_scheduling_applied_to_function_response(): + """response_scheduling on a tool is stamped onto the FunctionResponse part.""" + + def simple_fn(**kwargs) -> dict: + return {'result': 'test'} + + tool = FunctionTool(simple_fn) + tool.response_scheduling = types.FunctionResponseScheduling.SILENT + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + + function_call = types.FunctionCall(name=tool.name, args={}, id='fc_test') + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result_event = await handle_function_calls_async( + invocation_context, event, {tool.name: tool} + ) + + assert result_event is not None + function_response = result_event.content.parts[0].function_response + assert function_response.scheduling is types.FunctionResponseScheduling.SILENT + + +@pytest.mark.asyncio +async def test_response_scheduling_unset_by_default(): + """Without response_scheduling, the FunctionResponse part leaves it unset.""" + + def simple_fn(**kwargs) -> dict: + return {'result': 'test'} + + tool = FunctionTool(simple_fn) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + + function_call = types.FunctionCall(name=tool.name, args={}, id='fc_test') + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result_event = await handle_function_calls_async( + invocation_context, event, {tool.name: tool} + ) + + assert result_event is not None + function_response = result_event.content.parts[0].function_response + assert function_response.scheduling is None diff --git a/tests/unittests/tools/test_base_tool.py b/tests/unittests/tools/test_base_tool.py index 34e9269296..884e74c445 100644 --- a/tests/unittests/tools/test_base_tool.py +++ b/tests/unittests/tools/test_base_tool.py @@ -155,3 +155,15 @@ async def run_async(self, **kwargs): t2 = SimpleTool(name='test2', description='desc') t2._defers_response = True assert t2._defers_response is True + + +def test_response_scheduling_defaults_to_none(): + """response_scheduling defaults to None, preserving existing behavior.""" + + class SimpleTool(BaseTool): + + async def run_async(self, **kwargs): + pass + + t = SimpleTool(name='test', description='desc') + assert t.response_scheduling is None