mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Send user message to the agent that returned a corresponding function call if user message is a function response
PiperOrigin-RevId: 773895971
This commit is contained in:
Xiang (Sean) Zhou
committed by
Copybara-Service
Copybara-Service
parent
fb13963ded
commit
7c670f638b
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user