From 7c670f638bc17374ceb08740bdd057e55c9c2e12 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Fri, 20 Jun 2025 17:14:47 -0700 Subject: [PATCH] chore: Send user message to the agent that returned a corresponding function call if user message is a function response PiperOrigin-RevId: 773895971 --- src/google/adk/runners.py | 43 +++ tests/unittests/test_runners.py | 481 ++++++++++++++++++++++++++++++++ 2 files changed, 524 insertions(+) create mode 100644 tests/unittests/test_runners.py diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 01412a2b..936bc520 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -337,6 +337,8 @@ class Runner: """Finds the agent to run to continue the session. A qualified agent must be either of: + - The agent that returned a function call and the last user message is a + function response to this function call. - The root agent; - An LlmAgent who replied last and is capable to transfer to any other agent in the agent hierarchy. @@ -348,6 +350,15 @@ class Runner: Returns: The agent of the last message in the session or the root agent. """ + # If the last event is a function response, should send this response to + # the agent that returned the corressponding function call regardless the + # type of the agent. e.g. a remote a2a agent may surface a credential + # request as a special long running function tool call. + event = _find_function_call_event_if_last_event_is_function_response( + session + ) + if event and event.author: + return root_agent.find_agent(event.author) for event in filter(lambda e: e.author != 'user', reversed(session.events)): if event.author == root_agent.name: # Found root agent. @@ -527,3 +538,35 @@ class InMemoryRunner(Runner): session_service=self._in_memory_session_service, memory_service=InMemoryMemoryService(), ) + + +def _find_function_call_event_if_last_event_is_function_response( + session: Session, +) -> Optional[Event]: + events = session.events + if not events: + return None + + last_event = events[-1] + if ( + last_event.content + and last_event.content.parts + and any(part.function_response for part in last_event.content.parts) + ): + + function_call_id = next( + part.function_response.id + for part in last_event.content.parts + if part.function_response + ) + for i in range(len(events) - 2, -1, -1): + event = events[i] + # looking for the system long running request euc function call + function_calls = event.get_function_calls() + if not function_calls: + continue + + for function_call in function_calls: + if function_call.id == function_call_id: + return event + return None diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py new file mode 100644 index 00000000..56d7667a --- /dev/null +++ b/tests/unittests/test_runners.py @@ -0,0 +1,481 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.llm_agent import LlmAgent +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.events.event import Event +from google.adk.runners import _find_function_call_event_if_last_event_is_function_response +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.session import Session +from google.genai import types + + +class MockAgent(BaseAgent): + """Mock agent for unit testing.""" + + def __init__( + self, + name: str, + parent_agent: Optional[BaseAgent] = None, + ): + super().__init__(name=name, sub_agents=[]) + # BaseAgent doesn't have disallow_transfer_to_parent field + # This is intentional as we want to test non-LLM agents + if parent_agent: + self.parent_agent = parent_agent + + async def _run_async_impl(self, invocation_context): + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="Test response")] + ), + ) + + +class MockLlmAgent(LlmAgent): + """Mock LLM agent for unit testing.""" + + def __init__( + self, + name: str, + disallow_transfer_to_parent: bool = False, + parent_agent: Optional[BaseAgent] = None, + ): + # Use a string model instead of mock + super().__init__(name=name, model="gemini-1.5-pro", sub_agents=[]) + self.disallow_transfer_to_parent = disallow_transfer_to_parent + self.parent_agent = parent_agent + + async def _run_async_impl(self, invocation_context): + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="Test LLM response")] + ), + ) + + +class TestFindFunctionCallEventIfLastEventIsFunctionResponse: + """Tests for _find_function_call_event_if_last_event_is_function_response function.""" + + def test_no_function_response_in_last_event(self): + """Test when last event has no function response.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="user", + content=types.Content( + role="user", parts=[types.Part(text="Hello")] + ), + ) + ], + ) + + result = _find_function_call_event_if_last_event_is_function_response( + session + ) + assert result is None + + def test_empty_session_events(self): + """Test when session has no events.""" + session = Session( + id="test_session", user_id="test_user", app_name="test_app", events=[] + ) + + result = _find_function_call_event_if_last_event_is_function_response( + session + ) + assert result is None + + def test_last_event_has_function_response_but_no_matching_call(self): + """Test when last event has function response but no matching call found.""" + # Create a function response + function_response = types.FunctionResponse( + id="func_123", name="test_func", response={} + ) + + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="agent1", + content=types.Content( + role="model", + parts=[types.Part(text="Some other response")], + ), + ), + Event( + invocation_id="inv2", + author="user", + content=types.Content( + role="user", + parts=[types.Part(function_response=function_response)], + ), + ), + ], + ) + + result = _find_function_call_event_if_last_event_is_function_response( + session + ) + assert result is None + + def test_last_event_has_function_response_with_matching_call(self): + """Test when last event has function response with matching function call.""" + # Create a function call + function_call = types.FunctionCall(id="func_123", name="test_func", args={}) + + # Create a function response with matching ID + function_response = types.FunctionResponse( + id="func_123", name="test_func", response={} + ) + + call_event = Event( + invocation_id="inv1", + author="agent1", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call)] + ), + ) + + response_event = Event( + invocation_id="inv2", + author="user", + content=types.Content( + role="user", parts=[types.Part(function_response=function_response)] + ), + ) + + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[call_event, response_event], + ) + + result = _find_function_call_event_if_last_event_is_function_response( + session + ) + assert result == call_event + + def test_last_event_has_multiple_function_responses(self): + """Test when last event has multiple function responses.""" + # Create function calls + function_call1 = types.FunctionCall( + id="func_123", name="test_func1", args={} + ) + function_call2 = types.FunctionCall( + id="func_456", name="test_func2", args={} + ) + + # Create function responses + function_response1 = types.FunctionResponse( + id="func_123", name="test_func1", response={} + ) + function_response2 = types.FunctionResponse( + id="func_456", name="test_func2", response={} + ) + + call_event1 = Event( + invocation_id="inv1", + author="agent1", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call1)] + ), + ) + + call_event2 = Event( + invocation_id="inv2", + author="agent2", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call2)] + ), + ) + + response_event = Event( + invocation_id="inv3", + author="user", + content=types.Content( + role="user", + parts=[ + types.Part(function_response=function_response1), + types.Part(function_response=function_response2), + ], + ), + ) + + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[call_event1, call_event2, response_event], + ) + + # Should return the first matching function call event found + result = _find_function_call_event_if_last_event_is_function_response( + session + ) + assert result == call_event1 # First match (func_123) + + +class TestRunnerFindAgentToRun: + """Tests for Runner._find_agent_to_run method.""" + + def setup_method(self): + """Set up test fixtures.""" + self.session_service = InMemorySessionService() + self.artifact_service = InMemoryArtifactService() + + # Create test agents + self.root_agent = MockLlmAgent("root_agent") + self.sub_agent1 = MockLlmAgent("sub_agent1", parent_agent=self.root_agent) + self.sub_agent2 = MockLlmAgent("sub_agent2", parent_agent=self.root_agent) + self.non_transferable_agent = MockLlmAgent( + "non_transferable", + disallow_transfer_to_parent=True, + parent_agent=self.root_agent, + ) + + self.root_agent.sub_agents = [ + self.sub_agent1, + self.sub_agent2, + self.non_transferable_agent, + ] + + self.runner = Runner( + app_name="test_app", + agent=self.root_agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + + def test_find_agent_to_run_with_function_response_scenario(self): + """Test finding agent when last event is function response.""" + # Create a function call from sub_agent1 + function_call = types.FunctionCall(id="func_123", name="test_func", args={}) + function_response = types.FunctionResponse( + id="func_123", name="test_func", response={} + ) + + call_event = Event( + invocation_id="inv1", + author="sub_agent1", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call)] + ), + ) + + response_event = Event( + invocation_id="inv2", + author="user", + content=types.Content( + role="user", parts=[types.Part(function_response=function_response)] + ), + ) + + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[call_event, response_event], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.sub_agent1 + + def test_find_agent_to_run_returns_root_agent_when_no_events(self): + """Test that root agent is returned when session has no non-user events.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="user", + content=types.Content( + role="user", parts=[types.Part(text="Hello")] + ), + ) + ], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.root_agent + + def test_find_agent_to_run_returns_root_agent_when_found_in_events(self): + """Test that root agent is returned when it's found in session events.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="root_agent", + content=types.Content( + role="model", parts=[types.Part(text="Root response")] + ), + ) + ], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.root_agent + + def test_find_agent_to_run_returns_transferable_sub_agent(self): + """Test that transferable sub agent is returned when found.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="sub_agent1", + content=types.Content( + role="model", parts=[types.Part(text="Sub agent response")] + ), + ) + ], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.sub_agent1 + + def test_find_agent_to_run_skips_non_transferable_agent(self): + """Test that non-transferable agent is skipped and root agent is returned.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="non_transferable", + content=types.Content( + role="model", + parts=[types.Part(text="Non-transferable response")], + ), + ) + ], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.root_agent + + def test_find_agent_to_run_skips_unknown_agent(self): + """Test that unknown agent is skipped and root agent is returned.""" + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[ + Event( + invocation_id="inv1", + author="unknown_agent", + content=types.Content( + role="model", + parts=[types.Part(text="Unknown agent response")], + ), + ), + Event( + invocation_id="inv2", + author="root_agent", + content=types.Content( + role="model", parts=[types.Part(text="Root response")] + ), + ), + ], + ) + + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.root_agent + + def test_find_agent_to_run_function_response_takes_precedence(self): + """Test that function response scenario takes precedence over other logic.""" + # Create a function call from sub_agent2 + function_call = types.FunctionCall(id="func_456", name="test_func", args={}) + function_response = types.FunctionResponse( + id="func_456", name="test_func", response={} + ) + + call_event = Event( + invocation_id="inv1", + author="sub_agent2", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call)] + ), + ) + + # Add another event from root_agent + root_event = Event( + invocation_id="inv2", + author="root_agent", + content=types.Content( + role="model", parts=[types.Part(text="Root response")] + ), + ) + + response_event = Event( + invocation_id="inv3", + author="user", + content=types.Content( + role="user", parts=[types.Part(function_response=function_response)] + ), + ) + + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[call_event, root_event, response_event], + ) + + # Should return sub_agent2 due to function response, not root_agent + result = self.runner._find_agent_to_run(session, self.root_agent) + assert result == self.sub_agent2 + + def test_is_transferable_across_agent_tree_with_llm_agent(self): + """Test _is_transferable_across_agent_tree with LLM agent.""" + result = self.runner._is_transferable_across_agent_tree(self.sub_agent1) + assert result is True + + def test_is_transferable_across_agent_tree_with_non_transferable_agent(self): + """Test _is_transferable_across_agent_tree with non-transferable agent.""" + result = self.runner._is_transferable_across_agent_tree( + self.non_transferable_agent + ) + assert result is False + + def test_is_transferable_across_agent_tree_with_non_llm_agent(self): + """Test _is_transferable_across_agent_tree with non-LLM agent.""" + non_llm_agent = MockAgent("non_llm_agent") + # MockAgent inherits from BaseAgent, not LlmAgent, so it should return False + result = self.runner._is_transferable_across_agent_tree(non_llm_agent) + assert result is False