diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 441070808a..06f587110f 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -1158,12 +1158,15 @@ async def _run_one_step_async( run_config = _require_run_config(invocation_context) # Preprocess before calling the LLM. + preprocess_yielded_final_response = False async with Aclosing( self._preprocess_async(invocation_context, llm_request) ) as agen: async for event in agen: + if event.is_final_response(): + preprocess_yielded_final_response = True yield event - if invocation_context.end_invocation: + if invocation_context.end_invocation or preprocess_yielded_final_response: return # Resume the LLM agent based on the last event from the current branch. diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index f46d6f97b9..738e1d7821 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -2902,6 +2902,48 @@ async def _drive_one_llm_call(flow, invocation_context): pass +@pytest.mark.asyncio +async def test_preprocess_final_response_skips_llm_call(): + """A final response from preprocessing must finish the current step.""" + agent = Agent( + name='root_agent', model=testing_utils.MockModel.create(responses=[]) + ) + flow = BaseLlmFlowForTesting() + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='resume' + ) + function_response_event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + role='user', + parts=[ + types.Part.from_function_response( + name='resumed_tool', response={'result': 'done'} + ) + ], + ), + ) + function_response_event.actions.skip_summarization = True + + async def mock_preprocess(_ctx, _request): + yield function_response_event + + async def fail_if_llm_called(*_args, **_kwargs): + raise AssertionError('LLM should not be called after a final response') + yield # pylint: disable=unreachable + + with ( + mock.patch.object(flow, '_preprocess_async', side_effect=mock_preprocess), + mock.patch.object( + flow, '_call_llm_async', side_effect=fail_if_llm_called + ), + ): + events = [event async for event in flow.run_async(invocation_context)] + + assert events == [function_response_event] + + @pytest.mark.asyncio async def test_cfc_llm_calls_are_counted_against_max_llm_calls(): """support_cfc must not exempt a run from the max_llm_calls spend cap."""