wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
1867 lines
59 KiB
Python
1867 lines
59 KiB
Python
# Copyright 2026 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.
|
||
|
||
import asyncio
|
||
from typing import Any
|
||
from typing import Callable
|
||
from unittest import mock
|
||
|
||
from fastapi.openapi.models import HTTPBearer
|
||
from google.adk.agents.live_request_queue import LiveRequestQueue
|
||
from google.adk.agents.llm_agent import Agent
|
||
from google.adk.auth.auth_tool import AuthConfig
|
||
from google.adk.events.event import Event
|
||
from google.adk.events.event_actions import EventActions
|
||
from google.adk.events.ui_widget import UiWidget
|
||
from google.adk.flows.llm_flows.functions import find_matching_function_call
|
||
from google.adk.flows.llm_flows.functions import handle_function_calls_async
|
||
from google.adk.flows.llm_flows.functions import handle_function_calls_live
|
||
from google.adk.flows.llm_flows.functions import merge_parallel_function_response_events
|
||
from google.adk.tools.base_tool import BaseTool
|
||
from google.adk.tools.computer_use.computer_use_tool import ComputerUseTool
|
||
from google.adk.tools.function_tool import FunctionTool
|
||
from google.adk.tools.tool_confirmation import ToolConfirmation
|
||
from google.adk.tools.tool_context import ToolContext
|
||
from google.genai import types
|
||
import pytest
|
||
|
||
from ... import testing_utils
|
||
|
||
|
||
def test_simple_function():
|
||
function_call_1 = types.Part.from_function_call(
|
||
name='increase_by_one', args={'x': 1}
|
||
)
|
||
function_responses_2 = types.Part.from_function_response(
|
||
name='increase_by_one', response={'result': 2}
|
||
)
|
||
responses: list[types.Content] = [
|
||
function_call_1,
|
||
'response1',
|
||
'response2',
|
||
'response3',
|
||
'response4',
|
||
]
|
||
function_called = 0
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
def increase_by_one(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x + 1
|
||
|
||
agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one])
|
||
runner = testing_utils.InMemoryRunner(agent)
|
||
assert testing_utils.simplify_events(runner.run('test')) == [
|
||
('root_agent', function_call_1),
|
||
('root_agent', function_responses_2),
|
||
('root_agent', 'response1'),
|
||
]
|
||
|
||
# Asserts the requests.
|
||
assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
|
||
('user', 'test')
|
||
]
|
||
assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
|
||
('user', 'test'),
|
||
('model', function_call_1),
|
||
('user', function_responses_2),
|
||
]
|
||
|
||
# Asserts the function calls.
|
||
assert function_called == 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_function():
|
||
function_calls = [
|
||
types.Part.from_function_call(name='increase_by_one', args={'x': 1}),
|
||
types.Part.from_function_call(name='multiple_by_two', args={'x': 2}),
|
||
types.Part.from_function_call(name='multiple_by_two_sync', args={'x': 3}),
|
||
]
|
||
function_responses = [
|
||
types.Part.from_function_response(
|
||
name='increase_by_one', response={'result': 2}
|
||
),
|
||
types.Part.from_function_response(
|
||
name='multiple_by_two', response={'result': 4}
|
||
),
|
||
types.Part.from_function_response(
|
||
name='multiple_by_two_sync', response={'result': 6}
|
||
),
|
||
]
|
||
|
||
responses: list[types.Content] = [
|
||
function_calls,
|
||
'response1',
|
||
'response2',
|
||
'response3',
|
||
'response4',
|
||
]
|
||
function_called = 0
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
async def increase_by_one(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x + 1
|
||
|
||
async def multiple_by_two(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x * 2
|
||
|
||
def multiple_by_two_sync(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x * 2
|
||
|
||
agent = Agent(
|
||
name='root_agent',
|
||
model=mock_model,
|
||
tools=[increase_by_one, multiple_by_two, multiple_by_two_sync],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
assert testing_utils.simplify_events(events) == [
|
||
('root_agent', function_calls),
|
||
('root_agent', function_responses),
|
||
('root_agent', 'response1'),
|
||
]
|
||
|
||
# Asserts the requests.
|
||
assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
|
||
('user', 'test')
|
||
]
|
||
assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
|
||
('user', 'test'),
|
||
('model', function_calls),
|
||
('user', function_responses),
|
||
]
|
||
|
||
# Asserts the function calls.
|
||
assert function_called == 3
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_function_tool():
|
||
function_calls = [
|
||
types.Part.from_function_call(name='increase_by_one', args={'x': 1}),
|
||
types.Part.from_function_call(name='multiple_by_two', args={'x': 2}),
|
||
types.Part.from_function_call(name='multiple_by_two_sync', args={'x': 3}),
|
||
]
|
||
function_responses = [
|
||
types.Part.from_function_response(
|
||
name='increase_by_one', response={'result': 2}
|
||
),
|
||
types.Part.from_function_response(
|
||
name='multiple_by_two', response={'result': 4}
|
||
),
|
||
types.Part.from_function_response(
|
||
name='multiple_by_two_sync', response={'result': 6}
|
||
),
|
||
]
|
||
|
||
responses: list[types.Content] = [
|
||
function_calls,
|
||
'response1',
|
||
'response2',
|
||
'response3',
|
||
'response4',
|
||
]
|
||
function_called = 0
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
async def increase_by_one(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x + 1
|
||
|
||
async def multiple_by_two(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x * 2
|
||
|
||
def multiple_by_two_sync(x: int) -> int:
|
||
nonlocal function_called
|
||
function_called += 1
|
||
return x * 2
|
||
|
||
class TestTool(FunctionTool):
|
||
|
||
def __init__(self, func: Callable[..., Any]):
|
||
super().__init__(func=func)
|
||
|
||
wrapped_increase_by_one = TestTool(func=increase_by_one)
|
||
agent = Agent(
|
||
name='root_agent',
|
||
model=mock_model,
|
||
tools=[wrapped_increase_by_one, multiple_by_two, multiple_by_two_sync],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
assert testing_utils.simplify_events(events) == [
|
||
('root_agent', function_calls),
|
||
('root_agent', function_responses),
|
||
('root_agent', 'response1'),
|
||
]
|
||
|
||
# Asserts the requests.
|
||
assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
|
||
('user', 'test')
|
||
]
|
||
assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
|
||
('user', 'test'),
|
||
('model', function_calls),
|
||
('user', function_responses),
|
||
]
|
||
|
||
# Asserts the function calls.
|
||
assert function_called == 3
|
||
|
||
|
||
def test_update_state():
|
||
mock_model = testing_utils.MockModel.create(
|
||
responses=[
|
||
types.Part.from_function_call(name='update_state', args={}),
|
||
'response1',
|
||
]
|
||
)
|
||
|
||
def update_state(tool_context: ToolContext):
|
||
tool_context.state['x'] = 1
|
||
|
||
agent = Agent(name='root_agent', model=mock_model, tools=[update_state])
|
||
runner = testing_utils.InMemoryRunner(agent)
|
||
runner.run('test')
|
||
assert runner.session.state['x'] == 1
|
||
|
||
|
||
def test_function_call_id():
|
||
responses = [
|
||
types.Part.from_function_call(name='increase_by_one', args={'x': 1}),
|
||
'response1',
|
||
]
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
def increase_by_one(x: int) -> int:
|
||
return x + 1
|
||
|
||
agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one])
|
||
runner = testing_utils.InMemoryRunner(agent)
|
||
events = runner.run('test')
|
||
for request in mock_model.requests:
|
||
for content in request.contents:
|
||
for part in content.parts:
|
||
if part.function_call:
|
||
assert part.function_call.id is None
|
||
if part.function_response:
|
||
assert part.function_response.id is None
|
||
assert events[0].content.parts[0].function_call.id.startswith('adk-')
|
||
assert events[1].content.parts[0].function_response.id.startswith('adk-')
|
||
|
||
|
||
def test_find_function_call_event_no_function_response_in_last_event():
|
||
"""Test when last event has no function response."""
|
||
events = [
|
||
Event(
|
||
invocation_id='inv1',
|
||
author='user',
|
||
content=types.Content(role='user', parts=[types.Part(text='Hello')]),
|
||
)
|
||
]
|
||
|
||
result = find_matching_function_call(events)
|
||
assert result is None
|
||
|
||
|
||
def test_find_function_call_event_empty_session_events():
|
||
"""Test when session has no events."""
|
||
events = []
|
||
|
||
result = find_matching_function_call(events)
|
||
assert result is None
|
||
|
||
|
||
def test_find_function_call_event_function_response_but_no_matching_call():
|
||
"""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={}
|
||
)
|
||
|
||
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_matching_function_call(events)
|
||
assert result is None
|
||
|
||
|
||
def test_find_function_call_event_function_response_with_matching_call():
|
||
"""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)]
|
||
),
|
||
)
|
||
|
||
events = [call_event, response_event]
|
||
|
||
result = find_matching_function_call(events)
|
||
assert result == call_event
|
||
|
||
|
||
def test_find_function_call_event_multiple_function_responses():
|
||
"""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),
|
||
],
|
||
),
|
||
)
|
||
|
||
events = [call_event1, call_event2, response_event]
|
||
|
||
# Should return the first matching function call event found
|
||
result = find_matching_function_call(events)
|
||
assert result == call_event1 # First match (func_123)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_function_call_args_not_modified():
|
||
"""Test that function_call.args is not modified when making a copy."""
|
||
|
||
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=''
|
||
)
|
||
|
||
# Create original args that we want to ensure are not modified
|
||
original_args = {'param1': 'value1', 'param2': 42}
|
||
function_call = types.FunctionCall(name=tool.name, args=original_args)
|
||
content = types.Content(parts=[types.Part(function_call=function_call)])
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content,
|
||
)
|
||
tools_dict = {tool.name: tool}
|
||
|
||
# Test handle_function_calls_async
|
||
result_async = await handle_function_calls_async(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Verify original args are not modified
|
||
assert function_call.args == original_args
|
||
assert function_call.args is not original_args # Should be a copy
|
||
|
||
# Test handle_function_calls_live
|
||
result_live = await handle_function_calls_live(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Verify original args are still not modified
|
||
assert function_call.args == original_args
|
||
assert function_call.args is not original_args # Should be a copy
|
||
|
||
# Both should return valid results
|
||
assert result_async is not None
|
||
assert result_live is not None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_function_call_args_none_handling():
|
||
"""Test that function_call.args=None is handled correctly."""
|
||
|
||
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=''
|
||
)
|
||
|
||
# Create function call with None args
|
||
function_call = types.FunctionCall(name=tool.name, args=None)
|
||
content = types.Content(parts=[types.Part(function_call=function_call)])
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content,
|
||
)
|
||
tools_dict = {tool.name: tool}
|
||
|
||
# Test handle_function_calls_async
|
||
result_async = await handle_function_calls_async(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Test handle_function_calls_live
|
||
result_live = await handle_function_calls_live(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Both should return valid results even with None args
|
||
assert result_async is not None
|
||
assert result_live is not None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_function_call_args_copy_behavior():
|
||
"""Test that modifying the copied args doesn't affect the original."""
|
||
|
||
def simple_fn(test_param: str, other_param: int) -> dict:
|
||
# Modify the args to test that the copy prevents affecting the original
|
||
return {
|
||
'result': 'test',
|
||
'received_args': {'test_param': test_param, 'other_param': other_param},
|
||
}
|
||
|
||
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=''
|
||
)
|
||
|
||
# Create original args
|
||
original_args = {'test_param': 'original_value', 'other_param': 123}
|
||
function_call = types.FunctionCall(name=tool.name, args=original_args)
|
||
content = types.Content(parts=[types.Part(function_call=function_call)])
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content,
|
||
)
|
||
tools_dict = {tool.name: tool}
|
||
|
||
# Test handle_function_calls_async
|
||
result_async = await handle_function_calls_async(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Verify original args are unchanged
|
||
assert function_call.args == original_args
|
||
assert function_call.args['test_param'] == 'original_value'
|
||
|
||
# Verify the tool received the args correctly
|
||
assert result_async is not None
|
||
response = result_async.content.parts[0].function_response.response
|
||
|
||
# Check if the response has the expected structure
|
||
assert 'received_args' in response
|
||
received_args = response['received_args']
|
||
assert 'test_param' in received_args
|
||
assert received_args['test_param'] == 'original_value'
|
||
assert received_args['other_param'] == 123
|
||
assert (
|
||
function_call.args['test_param'] == 'original_value'
|
||
) # Original unchanged
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_function_call_args_deep_copy_behavior():
|
||
"""Test that deep copy behavior works correctly with nested structures."""
|
||
|
||
def simple_fn(nested_dict: dict, list_param: list) -> dict:
|
||
# Modify the nested structures to test deep copy
|
||
nested_dict['inner']['value'] = 'modified'
|
||
list_param.append('new_item')
|
||
return {
|
||
'result': 'test',
|
||
'received_nested': nested_dict,
|
||
'received_list': list_param,
|
||
}
|
||
|
||
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=''
|
||
)
|
||
|
||
# Create original args with nested structures
|
||
original_nested_dict = {'inner': {'value': 'original'}}
|
||
original_list = ['item1', 'item2']
|
||
original_args = {
|
||
'nested_dict': original_nested_dict,
|
||
'list_param': original_list,
|
||
}
|
||
|
||
function_call = types.FunctionCall(name=tool.name, args=original_args)
|
||
content = types.Content(parts=[types.Part(function_call=function_call)])
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content,
|
||
)
|
||
tools_dict = {tool.name: tool}
|
||
|
||
# Test handle_function_calls_async
|
||
result_async = await handle_function_calls_async(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
# Verify original args are completely unchanged
|
||
assert function_call.args == original_args
|
||
assert function_call.args['nested_dict']['inner']['value'] == 'original'
|
||
assert function_call.args['list_param'] == ['item1', 'item2']
|
||
|
||
# Verify the tool received the modified nested structures
|
||
assert result_async is not None
|
||
response = result_async.content.parts[0].function_response.response
|
||
|
||
# Check that the tool received modified versions
|
||
assert 'received_nested' in response
|
||
assert 'received_list' in response
|
||
assert response['received_nested']['inner']['value'] == 'modified'
|
||
assert 'new_item' in response['received_list']
|
||
|
||
# Verify original is still unchanged
|
||
assert function_call.args['nested_dict']['inner']['value'] == 'original'
|
||
assert function_call.args['list_param'] == ['item1', 'item2']
|
||
|
||
|
||
def test_shallow_vs_deep_copy_demonstration():
|
||
"""Demonstrate why deep copy is necessary vs shallow copy."""
|
||
import copy
|
||
|
||
# Original nested structure
|
||
original = {
|
||
'nested_dict': {'inner': {'value': 'original'}},
|
||
'list_param': ['item1', 'item2'],
|
||
}
|
||
|
||
# Shallow copy (what dict() does)
|
||
shallow_copy = dict(original)
|
||
|
||
# Deep copy (what copy.deepcopy() does)
|
||
deep_copy = copy.deepcopy(original)
|
||
|
||
# Modify the shallow copy
|
||
shallow_copy['nested_dict']['inner']['value'] = 'modified'
|
||
shallow_copy['list_param'].append('new_item')
|
||
|
||
# Check that shallow copy affects the original
|
||
assert (
|
||
original['nested_dict']['inner']['value'] == 'modified'
|
||
) # Original is affected!
|
||
assert 'new_item' in original['list_param'] # Original is affected!
|
||
|
||
# Reset original for deep copy test
|
||
original = {
|
||
'nested_dict': {'inner': {'value': 'original'}},
|
||
'list_param': ['item1', 'item2'],
|
||
}
|
||
|
||
# Modify the deep copy
|
||
deep_copy['nested_dict']['inner']['value'] = 'modified'
|
||
deep_copy['list_param'].append('new_item')
|
||
|
||
# Check that deep copy does NOT affect the original
|
||
assert (
|
||
original['nested_dict']['inner']['value'] == 'original'
|
||
) # Original unchanged
|
||
assert 'new_item' not in original['list_param'] # Original unchanged
|
||
assert (
|
||
deep_copy['nested_dict']['inner']['value'] == 'modified'
|
||
) # Copy is modified
|
||
assert 'new_item' in deep_copy['list_param'] # Copy is modified
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_parallel_function_execution_timing():
|
||
"""Test that multiple function calls are executed in parallel, not sequentially."""
|
||
execution_order = []
|
||
|
||
async def slow_function_1(delay: float = 0.1) -> dict:
|
||
execution_order.append('start_1')
|
||
await asyncio.sleep(delay)
|
||
execution_order.append('end_1')
|
||
return {'result': 'function_1_result'}
|
||
|
||
async def slow_function_2(delay: float = 0.1) -> dict:
|
||
execution_order.append('start_2')
|
||
await asyncio.sleep(delay)
|
||
execution_order.append('end_2')
|
||
return {'result': 'function_2_result'}
|
||
|
||
# Create function calls
|
||
function_calls = [
|
||
types.Part.from_function_call(
|
||
name='slow_function_1', args={'delay': 0.1}
|
||
),
|
||
types.Part.from_function_call(
|
||
name='slow_function_2', args={'delay': 0.1}
|
||
),
|
||
]
|
||
|
||
function_responses = [
|
||
types.Part.from_function_response(
|
||
name='slow_function_1', response={'result': 'function_1_result'}
|
||
),
|
||
types.Part.from_function_response(
|
||
name='slow_function_2', response={'result': 'function_2_result'}
|
||
),
|
||
]
|
||
|
||
responses: list[types.Content] = [
|
||
function_calls,
|
||
'response1',
|
||
]
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[slow_function_1, slow_function_2],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# Parallel execution means both functions start before either one finishes.
|
||
assert set(execution_order) == {'start_1', 'start_2', 'end_1', 'end_2'}
|
||
last_start = max(
|
||
execution_order.index('start_1'), execution_order.index('start_2')
|
||
)
|
||
first_end = min(
|
||
execution_order.index('end_1'), execution_order.index('end_2')
|
||
)
|
||
assert (
|
||
last_start < first_end
|
||
), f'Functions did not overlap; execution was sequential: {execution_order}'
|
||
|
||
# Verify the results are correct
|
||
assert testing_utils.simplify_events(events) == [
|
||
('test_agent', function_calls),
|
||
('test_agent', function_responses),
|
||
('test_agent', 'response1'),
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_parallel_state_modifications_thread_safety():
|
||
"""Test that parallel function calls modifying state are thread-safe."""
|
||
state_modifications = []
|
||
|
||
def modify_state_1(tool_context: ToolContext) -> dict:
|
||
# Track when this function modifies state
|
||
current_state = dict(tool_context.state.to_dict())
|
||
state_modifications.append(('func_1_start', current_state))
|
||
|
||
tool_context.state['counter'] = tool_context.state.get('counter', 0) + 1
|
||
tool_context.state['func_1_executed'] = True
|
||
|
||
final_state = dict(tool_context.state.to_dict())
|
||
state_modifications.append(('func_1_end', final_state))
|
||
return {'result': 'modified_state_1'}
|
||
|
||
def modify_state_2(tool_context: ToolContext) -> dict:
|
||
# Track when this function modifies state
|
||
current_state = dict(tool_context.state.to_dict())
|
||
state_modifications.append(('func_2_start', current_state))
|
||
|
||
tool_context.state['counter'] = tool_context.state.get('counter', 0) + 1
|
||
tool_context.state['func_2_executed'] = True
|
||
|
||
final_state = dict(tool_context.state.to_dict())
|
||
state_modifications.append(('func_2_end', final_state))
|
||
return {'result': 'modified_state_2'}
|
||
|
||
# Create function calls
|
||
function_calls = [
|
||
types.Part.from_function_call(name='modify_state_1', args={}),
|
||
types.Part.from_function_call(name='modify_state_2', args={}),
|
||
]
|
||
|
||
responses: list[types.Content] = [
|
||
function_calls,
|
||
'response1',
|
||
]
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[modify_state_1, modify_state_2],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# Verify the parallel execution worked correctly by checking the events
|
||
# The function response event should have the merged state_delta
|
||
function_response_event = events[
|
||
1
|
||
] # Second event should be the function response
|
||
assert function_response_event.actions.state_delta['counter'] == 2
|
||
assert function_response_event.actions.state_delta['func_1_executed'] is True
|
||
assert function_response_event.actions.state_delta['func_2_executed'] is True
|
||
|
||
# Verify both functions were called
|
||
assert len(state_modifications) == 4 # 2 functions × 2 events each
|
||
|
||
# Extract function names from modifications
|
||
func_names = [mod[0] for mod in state_modifications]
|
||
assert 'func_1_start' in func_names
|
||
assert 'func_1_end' in func_names
|
||
assert 'func_2_start' in func_names
|
||
assert 'func_2_end' in func_names
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_sync_function_blocks_async_functions():
|
||
"""Test that sync functions block async functions from running concurrently."""
|
||
execution_order = []
|
||
|
||
def blocking_sync_function() -> dict:
|
||
execution_order.append('sync_A')
|
||
# Simulate CPU-intensive work that blocks the event loop
|
||
result = 0
|
||
for i in range(1000000): # This blocks the event loop
|
||
result += i
|
||
execution_order.append('sync_B')
|
||
return {'result': 'sync_done'}
|
||
|
||
async def yielding_async_function() -> dict:
|
||
execution_order.append('async_C')
|
||
await asyncio.sleep(
|
||
0.001
|
||
) # This should yield, but can't if event loop is blocked
|
||
execution_order.append('async_D')
|
||
return {'result': 'async_done'}
|
||
|
||
# Create function calls - these should run "in parallel"
|
||
function_calls = [
|
||
types.Part.from_function_call(name='blocking_sync_function', args={}),
|
||
types.Part.from_function_call(name='yielding_async_function', args={}),
|
||
]
|
||
|
||
responses: list[types.Content] = [function_calls, 'response1']
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[blocking_sync_function, yielding_async_function],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# With blocking sync function, execution should be sequential: A, B, C, D
|
||
# The sync function blocks, preventing the async function from yielding properly
|
||
assert execution_order == ['sync_A', 'sync_B', 'async_C', 'async_D']
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_function_without_yield_blocks_others():
|
||
"""Test that async functions without yield statements block other functions."""
|
||
execution_order = []
|
||
|
||
async def non_yielding_async_function() -> dict:
|
||
execution_order.append('non_yield_A')
|
||
# CPU-intensive work without any await statements - blocks like sync function
|
||
result = 0
|
||
for i in range(1000000): # No await here, so this blocks the event loop
|
||
result += i
|
||
execution_order.append('non_yield_B')
|
||
return {'result': 'non_yielding_done'}
|
||
|
||
async def yielding_async_function() -> dict:
|
||
execution_order.append('yield_C')
|
||
await asyncio.sleep(
|
||
0.001
|
||
) # This should yield, but can't if event loop is blocked
|
||
execution_order.append('yield_D')
|
||
return {'result': 'yielding_done'}
|
||
|
||
# Create function calls
|
||
function_calls = [
|
||
types.Part.from_function_call(
|
||
name='non_yielding_async_function', args={}
|
||
),
|
||
types.Part.from_function_call(name='yielding_async_function', args={}),
|
||
]
|
||
|
||
responses: list[types.Content] = [function_calls, 'response1']
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[non_yielding_async_function, yielding_async_function],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# Non-yielding async function blocks, so execution is sequential: A, B, C, D
|
||
assert execution_order == ['non_yield_A', 'non_yield_B', 'yield_C', 'yield_D']
|
||
|
||
|
||
def test_merge_parallel_function_response_events_preserves_invocation_id():
|
||
"""Test that merge_parallel_function_response_events preserves the base event's invocation_id."""
|
||
# Create multiple function response events with different invocation IDs
|
||
invocation_id = 'base_invocation_123'
|
||
|
||
function_response1 = types.FunctionResponse(
|
||
id='func_123', name='test_function1', response={'result': 'success1'}
|
||
)
|
||
|
||
function_response2 = types.FunctionResponse(
|
||
id='func_456', name='test_function2', response={'result': 'success2'}
|
||
)
|
||
|
||
event1 = Event(
|
||
invocation_id=invocation_id,
|
||
author='test_agent',
|
||
content=types.Content(
|
||
role='user', parts=[types.Part(function_response=function_response1)]
|
||
),
|
||
)
|
||
|
||
event2 = Event(
|
||
invocation_id='different_invocation_456', # Different invocation ID
|
||
author='test_agent',
|
||
content=types.Content(
|
||
role='user', parts=[types.Part(function_response=function_response2)]
|
||
),
|
||
)
|
||
|
||
# Merge the events
|
||
merged_event = merge_parallel_function_response_events([event1, event2])
|
||
|
||
# Should preserve the base event's (first event's) invocation_id
|
||
assert merged_event.invocation_id == invocation_id
|
||
assert merged_event.invocation_id != 'different_invocation_456'
|
||
|
||
# Should contain both function responses
|
||
assert len(merged_event.content.parts) == 2
|
||
|
||
# Verify the responses are preserved
|
||
response_ids = {
|
||
part.function_response.id for part in merged_event.content.parts
|
||
}
|
||
assert 'func_123' in response_ids
|
||
assert 'func_456' in response_ids
|
||
|
||
|
||
def test_merge_parallel_function_response_events_single_event():
|
||
"""Test that merge_parallel_function_response_events returns single event unchanged."""
|
||
invocation_id = 'single_invocation_123'
|
||
|
||
function_response = types.FunctionResponse(
|
||
id='func_123', name='test_function', response={'result': 'success'}
|
||
)
|
||
|
||
event = Event(
|
||
invocation_id=invocation_id,
|
||
author='test_agent',
|
||
content=types.Content(
|
||
role='user', parts=[types.Part(function_response=function_response)]
|
||
),
|
||
)
|
||
|
||
# Merge single event
|
||
merged_event = merge_parallel_function_response_events([event])
|
||
|
||
# Should return the same event object
|
||
assert merged_event is event
|
||
assert merged_event.invocation_id == invocation_id
|
||
|
||
|
||
def test_merge_parallel_function_response_events_preserves_other_attributes():
|
||
"""Test that merge_parallel_function_response_events preserves other attributes from base event."""
|
||
invocation_id = 'base_invocation_123'
|
||
base_author = 'base_agent'
|
||
base_branch = 'main_branch'
|
||
|
||
function_response1 = types.FunctionResponse(
|
||
id='func_123', name='test_function1', response={'result': 'success1'}
|
||
)
|
||
|
||
function_response2 = types.FunctionResponse(
|
||
id='func_456', name='test_function2', response={'result': 'success2'}
|
||
)
|
||
|
||
event1 = Event(
|
||
invocation_id=invocation_id,
|
||
author=base_author,
|
||
branch=base_branch,
|
||
content=types.Content(
|
||
role='user', parts=[types.Part(function_response=function_response1)]
|
||
),
|
||
)
|
||
|
||
event2 = Event(
|
||
invocation_id='different_invocation_456',
|
||
author='different_agent', # Different author
|
||
branch='different_branch', # Different branch
|
||
content=types.Content(
|
||
role='user', parts=[types.Part(function_response=function_response2)]
|
||
),
|
||
)
|
||
|
||
# Merge the events
|
||
merged_event = merge_parallel_function_response_events([event1, event2])
|
||
|
||
# Should preserve base event's attributes
|
||
assert merged_event.invocation_id == invocation_id
|
||
assert merged_event.author == base_author
|
||
assert merged_event.branch == base_branch
|
||
|
||
# Should contain both function responses
|
||
assert len(merged_event.content.parts) == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_yielding_async_functions_run_concurrently():
|
||
"""Test that async functions with proper yields run concurrently."""
|
||
execution_order = []
|
||
|
||
async def yielding_async_function_1() -> dict:
|
||
execution_order.append('func1_A')
|
||
await asyncio.sleep(0.001) # Yield control
|
||
execution_order.append('func1_B')
|
||
return {'result': 'func1_done'}
|
||
|
||
async def yielding_async_function_2() -> dict:
|
||
execution_order.append('func2_C')
|
||
await asyncio.sleep(0.001) # Yield control
|
||
execution_order.append('func2_D')
|
||
return {'result': 'func2_done'}
|
||
|
||
# Create function calls
|
||
function_calls = [
|
||
types.Part.from_function_call(name='yielding_async_function_1', args={}),
|
||
types.Part.from_function_call(name='yielding_async_function_2', args={}),
|
||
]
|
||
|
||
responses: list[types.Content] = [function_calls, 'response1']
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[yielding_async_function_1, yielding_async_function_2],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# With proper yielding, execution should interleave: A, C, B, D
|
||
# Both functions start, yield, then complete
|
||
assert execution_order == ['func1_A', 'func2_C', 'func1_B', 'func2_D']
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mixed_function_types_execution_order():
|
||
"""Test execution order with all three types of functions."""
|
||
execution_order = []
|
||
|
||
def sync_function() -> dict:
|
||
execution_order.append('sync_A')
|
||
# Small amount of blocking work
|
||
result = sum(range(100000))
|
||
execution_order.append('sync_B')
|
||
return {'result': 'sync_done'}
|
||
|
||
async def non_yielding_async() -> dict:
|
||
execution_order.append('non_yield_C')
|
||
# CPU work without yield
|
||
result = sum(range(100000))
|
||
execution_order.append('non_yield_D')
|
||
return {'result': 'non_yield_done'}
|
||
|
||
async def yielding_async() -> dict:
|
||
execution_order.append('yield_E')
|
||
await asyncio.sleep(0.001) # Proper yield
|
||
execution_order.append('yield_F')
|
||
return {'result': 'yield_done'}
|
||
|
||
# Create function calls
|
||
function_calls = [
|
||
types.Part.from_function_call(name='sync_function', args={}),
|
||
types.Part.from_function_call(name='non_yielding_async', args={}),
|
||
types.Part.from_function_call(name='yielding_async', args={}),
|
||
]
|
||
|
||
responses: list[types.Content] = [function_calls, 'response1']
|
||
mock_model = testing_utils.MockModel.create(responses=responses)
|
||
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=mock_model,
|
||
tools=[sync_function, non_yielding_async, yielding_async],
|
||
)
|
||
runner = testing_utils.TestInMemoryRunner(agent)
|
||
events = await runner.run_async_with_new_session('test')
|
||
|
||
# All blocking functions run sequentially, then the yielding one
|
||
# Expected order: sync_A, sync_B, non_yield_C, non_yield_D, yield_E, yield_F
|
||
assert execution_order == [
|
||
'sync_A',
|
||
'sync_B',
|
||
'non_yield_C',
|
||
'non_yield_D',
|
||
'yield_E',
|
||
'yield_F',
|
||
]
|
||
|
||
|
||
def test_merge_parallel_function_response_events_merges_ui_widgets():
|
||
"""Test that merge_parallel_function_response_events merges render_ui_widgets."""
|
||
invocation_id = 'base_invocation_123'
|
||
|
||
widget1 = UiWidget(
|
||
id='widget_1', provider='mcp', payload={'resource_uri': 'ui://widget1'}
|
||
)
|
||
widget2 = UiWidget(
|
||
id='widget_2', provider='mcp', payload={'resource_uri': 'ui://widget2'}
|
||
)
|
||
widget3 = UiWidget(
|
||
id='widget_3', provider='mcp', payload={'resource_uri': 'ui://widget3'}
|
||
)
|
||
|
||
event1 = Event(
|
||
invocation_id=invocation_id,
|
||
author='test_agent',
|
||
actions=EventActions(render_ui_widgets=[widget1]),
|
||
)
|
||
|
||
event2 = Event(
|
||
invocation_id='different_invocation_456',
|
||
author='different_agent',
|
||
actions=EventActions(render_ui_widgets=[widget2, widget3]),
|
||
)
|
||
|
||
# Merge the events
|
||
merged_event = merge_parallel_function_response_events([event1, event2])
|
||
|
||
# Should contain all ui widgets
|
||
assert merged_event.actions.render_ui_widgets is not None
|
||
assert len(merged_event.actions.render_ui_widgets) == 3
|
||
|
||
widget_ids = {widget.id for widget in merged_event.actions.render_ui_widgets}
|
||
assert 'widget_1' in widget_ids
|
||
assert 'widget_2' in widget_ids
|
||
assert 'widget_3' in widget_ids
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
async def test_computer_use_tool_decoding_behavior(handle_function_calls):
|
||
"""Tests that computer use tools automatically decode base64 images."""
|
||
valid_b64 = 'R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7'
|
||
|
||
# make the tool return a dictionary with the image
|
||
async def mock_run(*args, **kwargs):
|
||
return {
|
||
'image': {'data': valid_b64, 'mimetype': 'image/png'},
|
||
'url': 'https://example.com',
|
||
}
|
||
|
||
# create a ComputerUseTool
|
||
tool = ComputerUseTool(func=mock_run, screen_size=(1024, 768))
|
||
|
||
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=''
|
||
)
|
||
|
||
# Create function call
|
||
function_call = types.FunctionCall(name=tool.name, args={})
|
||
content = types.Content(parts=[types.Part(function_call=function_call)])
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content,
|
||
)
|
||
tools_dict = {tool.name: tool}
|
||
|
||
result = await handle_function_calls(
|
||
invocation_context,
|
||
event,
|
||
tools_dict,
|
||
)
|
||
|
||
assert result is not None
|
||
response_part = result.content.parts[0].function_response
|
||
|
||
# Verify original image data is removed from the dict response
|
||
assert 'image' not in response_part.response
|
||
assert 'url' in response_part.response
|
||
# Verify the image was converted to a blob
|
||
assert len(response_part.parts) == 1
|
||
assert response_part.parts[0].inline_data is not None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_handle_function_calls_live_preserves_live_session_id():
|
||
"""Tests that handle_function_calls_live preserves live_session_id for single call."""
|
||
|
||
def simple_fn() -> dict[str, str]:
|
||
return {'result': 'test'}
|
||
|
||
tool1 = FunctionTool(simple_fn)
|
||
model = testing_utils.MockModel.create(responses=[])
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=model,
|
||
tools=[tool1],
|
||
)
|
||
invocation_context = await testing_utils.create_invocation_context(
|
||
agent=agent, user_content=''
|
||
)
|
||
|
||
function_call1 = types.FunctionCall(id='call_1', name=tool1.name, args={})
|
||
content1 = types.Content(parts=[types.Part(function_call=function_call1)])
|
||
event1 = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content1,
|
||
live_session_id='test-live-session-id',
|
||
)
|
||
tools_dict = {tool1.name: tool1}
|
||
|
||
result_single = await handle_function_calls_live(
|
||
invocation_context,
|
||
event1,
|
||
tools_dict,
|
||
)
|
||
|
||
assert result_single is not None
|
||
assert result_single.live_session_id == 'test-live-session-id'
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_handle_function_calls_live_parallel_preserves_live_session_id():
|
||
"""Tests that handle_function_calls_live preserves live_session_id for parallel calls."""
|
||
|
||
def simple_fn() -> dict[str, str]:
|
||
return {'result': 'test'}
|
||
|
||
tool1 = FunctionTool(simple_fn)
|
||
tool2 = FunctionTool(simple_fn)
|
||
model = testing_utils.MockModel.create(responses=[])
|
||
agent = Agent(
|
||
name='test_agent',
|
||
model=model,
|
||
tools=[tool1, tool2],
|
||
)
|
||
invocation_context = await testing_utils.create_invocation_context(
|
||
agent=agent, user_content=''
|
||
)
|
||
|
||
function_call1 = types.FunctionCall(id='call_1', name=tool1.name, args={})
|
||
function_call2 = types.FunctionCall(id='call_2', name=tool2.name, args={})
|
||
content2 = types.Content(
|
||
parts=[
|
||
types.Part(function_call=function_call1),
|
||
types.Part(function_call=function_call2),
|
||
]
|
||
)
|
||
event2 = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=content2,
|
||
live_session_id='test-live-session-id-parallel',
|
||
)
|
||
tools_dict = {tool1.name: tool1, tool2.name: tool2}
|
||
|
||
result_parallel = await handle_function_calls_live(
|
||
invocation_context,
|
||
event2,
|
||
tools_dict,
|
||
)
|
||
|
||
assert result_parallel is not None
|
||
assert result_parallel.live_session_id == 'test-live-session-id-parallel'
|
||
|
||
|
||
class _MockControlSignalTool(BaseTool):
|
||
"""A tool that simulates requesting confirmation or OAuth authentication."""
|
||
|
||
def __init__(self, name: str, behavior: str):
|
||
super().__init__(name=name, description='Simulated control tool')
|
||
self.behavior = behavior
|
||
|
||
async def run_async(self, *, args, tool_context):
|
||
if self.behavior == 'confirm':
|
||
tool_context.actions.requested_tool_confirmations = {
|
||
'fc_test_confirm': ToolConfirmation(hint='Authorize execution?')
|
||
}
|
||
return {'error': 'This tool requires user approval.'}
|
||
elif self.behavior == 'auth':
|
||
tool_context.actions.requested_auth_configs = {
|
||
'fc_test_auth': AuthConfig(auth_scheme=HTTPBearer())
|
||
}
|
||
return {'error': 'Please complete OAuth setup.'}
|
||
|
||
def _detect_error_in_response(self, response: Any) -> str | None:
|
||
if isinstance(response, dict) and 'error' in response:
|
||
return 'TOOL_ERROR'
|
||
return None
|
||
|
||
|
||
class _ErrorDetectingTool(BaseTool):
|
||
"""A test tool whose _detect_error_in_response raises an exception."""
|
||
|
||
async def run_async(self, *, args, tool_context):
|
||
return {'result': 'result'}
|
||
|
||
def _detect_error_in_response(self, response: Any) -> str | None:
|
||
raise RuntimeError('detection exploded')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
@pytest.mark.parametrize(
|
||
'mock_response,expected_error_type',
|
||
[
|
||
({'error': 'Internal component timeout'}, 'TOOL_ERROR'),
|
||
({'result': 'Execution succeeded'}, None),
|
||
],
|
||
ids=['dict_error_recorded', 'success_dict_ignored'],
|
||
)
|
||
async def test_e2e_telemetry_error_classification(
|
||
monkeypatch, handle_function_calls, mock_response, expected_error_type
|
||
):
|
||
"""E2E: asserts that tool outputs successfully translate to targeted OTel span error attributes."""
|
||
recorded_calls = []
|
||
|
||
# Intercept trace_tool_call to capture final telemetry state
|
||
monkeypatch.setattr(
|
||
'google.adk.telemetry._instrumentation.tracing.trace_tool_call',
|
||
lambda **kw: recorded_calls.append(kw),
|
||
)
|
||
|
||
tool = FunctionTool(func=lambda: mock_response)
|
||
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)]),
|
||
)
|
||
|
||
await handle_function_calls(invocation_context, event, {tool.name: tool})
|
||
|
||
assert len(recorded_calls) == 1
|
||
assert recorded_calls[0]['error_type'] == expected_error_type
|
||
assert recorded_calls[0]['error'] is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
async def test_exception_takes_precedence_over_dict_error(
|
||
monkeypatch, handle_function_calls
|
||
):
|
||
"""End-to-end integration: exception takes strict precedence over manual dict error_type."""
|
||
recorded_calls = []
|
||
monkeypatch.setattr(
|
||
'google.adk.telemetry._instrumentation.tracing.trace_tool_call',
|
||
lambda **kw: recorded_calls.append(kw),
|
||
)
|
||
|
||
def mock_crashing_func():
|
||
raise ValueError('Fatal arithmetic error')
|
||
|
||
tool = FunctionTool(func=mock_crashing_func)
|
||
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_exception'
|
||
)
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
with pytest.raises(ValueError, match='Fatal arithmetic error'):
|
||
await handle_function_calls(invocation_context, event, {tool.name: tool})
|
||
|
||
assert len(recorded_calls) == 1
|
||
assert isinstance(recorded_calls[0]['error'], ValueError)
|
||
assert recorded_calls[0]['error_type'] is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
async def test_detection_skipped_when_confirmation_requested(
|
||
monkeypatch, handle_function_calls
|
||
):
|
||
"""E2E confirmation verification: control prompt avoids polluting telemetry with TOOL_ERROR."""
|
||
recorded_calls = []
|
||
monkeypatch.setattr(
|
||
'google.adk.telemetry._instrumentation.tracing.trace_tool_call',
|
||
lambda **kw: recorded_calls.append(kw),
|
||
)
|
||
|
||
tool = _MockControlSignalTool(name='confirm_tool', behavior='confirm')
|
||
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_confirm'
|
||
)
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
await handle_function_calls(invocation_context, event, {tool.name: tool})
|
||
|
||
assert len(recorded_calls) == 1
|
||
assert recorded_calls[0]['error_type'] is None
|
||
assert recorded_calls[0]['error'] is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
async def test_detection_skipped_when_auth_requested(
|
||
monkeypatch, handle_function_calls
|
||
):
|
||
"""E2E OAuth verification: authenticate control prompt avoids polluting telemetry with TOOL_ERROR."""
|
||
recorded_calls = []
|
||
monkeypatch.setattr(
|
||
'google.adk.telemetry._instrumentation.tracing.trace_tool_call',
|
||
lambda **kw: recorded_calls.append(kw),
|
||
)
|
||
|
||
tool = _MockControlSignalTool(name='auth_tool', behavior='auth')
|
||
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_auth')
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
await handle_function_calls(invocation_context, event, {tool.name: tool})
|
||
|
||
assert len(recorded_calls) == 1
|
||
assert recorded_calls[0]['error_type'] is None
|
||
assert recorded_calls[0]['error'] is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
'handle_function_calls',
|
||
[
|
||
(handle_function_calls_async),
|
||
(handle_function_calls_live),
|
||
],
|
||
)
|
||
async def test_detection_exception_does_not_break_tool_call(
|
||
monkeypatch, handle_function_calls
|
||
):
|
||
"""Safety Verification: telemetry errors during error parsing are caught cleanly, not crashing tool calls."""
|
||
recorded_calls = []
|
||
monkeypatch.setattr(
|
||
'google.adk.telemetry._instrumentation.tracing.trace_tool_call',
|
||
lambda **kw: recorded_calls.append(kw),
|
||
)
|
||
|
||
tool = _ErrorDetectingTool(
|
||
name='buggy_telemetry_tool', description='raises on tel'
|
||
)
|
||
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_buggy'
|
||
)
|
||
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(
|
||
invocation_context, event, {tool.name: tool}
|
||
)
|
||
|
||
assert result_event is not None
|
||
assert result_event.content.parts[0].function_response.response == {
|
||
'result': 'result'
|
||
}
|
||
|
||
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
|
||
|
||
|
||
async def _drain_live_function_responses(
|
||
live_request_queue: LiveRequestQueue,
|
||
count: int,
|
||
) -> list[types.Content]:
|
||
"""Drains ``count`` contents from a live request queue, ignoring acks."""
|
||
contents = []
|
||
while len(contents) < count:
|
||
request = await asyncio.wait_for(live_request_queue._queue.get(), timeout=5)
|
||
if request.content is not None:
|
||
contents.append(request.content)
|
||
return contents
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_with_scheduling_emits_function_response():
|
||
"""A streaming tool with response_scheduling relays yields as FunctionResponses."""
|
||
|
||
async def streaming_fn(**kwargs):
|
||
yield 'first'
|
||
yield 'second'
|
||
|
||
tool = FunctionTool(streaming_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=''
|
||
)
|
||
invocation_context.live_request_queue = LiveRequestQueue()
|
||
|
||
function_call = types.FunctionCall(name=tool.name, args={}, id='fc_stream')
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
await handle_function_calls_live(invocation_context, event, {tool.name: tool})
|
||
contents = await _drain_live_function_responses(
|
||
invocation_context.live_request_queue, count=2
|
||
)
|
||
|
||
responses = [content.parts[0].function_response for content in contents]
|
||
assert [r.response for r in responses] == [
|
||
{'result': 'first'},
|
||
{'result': 'second'},
|
||
]
|
||
assert all(r.id == 'fc_stream' for r in responses)
|
||
assert all(
|
||
r.scheduling is types.FunctionResponseScheduling.SILENT for r in responses
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_without_scheduling_emits_function_response():
|
||
"""A streaming tool without response_scheduling still relays yields as
|
||
FunctionResponses, leaving scheduling unset."""
|
||
|
||
async def streaming_fn(**kwargs):
|
||
yield 'hello'
|
||
|
||
tool = FunctionTool(streaming_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=''
|
||
)
|
||
invocation_context.live_request_queue = LiveRequestQueue()
|
||
|
||
function_call = types.FunctionCall(name=tool.name, args={}, id='fc_stream')
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
await handle_function_calls_live(invocation_context, event, {tool.name: tool})
|
||
contents = await _drain_live_function_responses(
|
||
invocation_context.live_request_queue, count=1
|
||
)
|
||
|
||
function_response = contents[0].parts[0].function_response
|
||
assert function_response.response == {'result': 'hello'}
|
||
assert function_response.id == 'fc_stream'
|
||
assert function_response.scheduling is None
|
||
|
||
|
||
async def test_non_blocking_tool_handled_asynchronously():
|
||
"""Tests that a NON_BLOCKING tool returns None inline and pushes to live request queue."""
|
||
import asyncio
|
||
|
||
async def slow_fn() -> dict[str, str]:
|
||
await asyncio.sleep(0.1)
|
||
return {'result': 'done'}
|
||
|
||
tool = FunctionTool(slow_fn)
|
||
tool.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE
|
||
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=''
|
||
)
|
||
invocation_context.live_request_queue = LiveRequestQueue()
|
||
|
||
function_call = types.FunctionCall(
|
||
name=tool.name, args={}, id='fc_non_blocking'
|
||
)
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
# Inline call should return None, indicating it is handled asynchronously
|
||
result = await handle_function_calls_live(
|
||
invocation_context, event, {tool.name: tool}
|
||
)
|
||
assert result is None
|
||
# Verify strong reference is kept in active_non_blocking_tool_tasks while running
|
||
assert invocation_context.active_non_blocking_tool_tasks is not None
|
||
task_key = f'{tool.name}_fc_non_blocking'
|
||
assert task_key in invocation_context.active_non_blocking_tool_tasks
|
||
|
||
# Wait for the background task to complete and push to the queue
|
||
contents = await _drain_live_function_responses(
|
||
invocation_context.live_request_queue, count=1
|
||
)
|
||
|
||
function_response = contents[0].parts[0].function_response
|
||
assert function_response.response == {'result': 'done'}
|
||
assert function_response.id == 'fc_non_blocking'
|
||
assert (
|
||
function_response.scheduling == types.FunctionResponseScheduling.WHEN_IDLE
|
||
)
|
||
# Give event loop a tick for finally block to clean up the completed task
|
||
await asyncio.sleep(0)
|
||
assert task_key not in invocation_context.active_non_blocking_tool_tasks
|
||
|
||
|
||
async def test_non_blocking_tool_exception_handling_and_cleanup():
|
||
"""Tests that an exception in a NON_BLOCKING tool is logged and cleaned up."""
|
||
import asyncio
|
||
|
||
async def failing_fn() -> dict[str, str]:
|
||
await asyncio.sleep(0.05)
|
||
raise RuntimeError('Tool execution failed purposely')
|
||
|
||
tool = FunctionTool(failing_fn)
|
||
tool.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE
|
||
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=''
|
||
)
|
||
invocation_context.live_request_queue = LiveRequestQueue()
|
||
|
||
function_call = types.FunctionCall(
|
||
name=tool.name, args={}, id='fc_failing_non_blocking'
|
||
)
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(parts=[types.Part(function_call=function_call)]),
|
||
)
|
||
|
||
result = await handle_function_calls_live(
|
||
invocation_context, event, {tool.name: tool}
|
||
)
|
||
assert result is None
|
||
task_key = f'{tool.name}_fc_failing_non_blocking'
|
||
assert (
|
||
invocation_context.active_non_blocking_tool_tasks
|
||
and task_key in invocation_context.active_non_blocking_tool_tasks
|
||
)
|
||
|
||
# Allow background task to execute and raise its exception
|
||
await asyncio.sleep(0.1)
|
||
# Give event loop a tick for finally block to clean up the completed task
|
||
await asyncio.sleep(0)
|
||
assert task_key not in invocation_context.active_non_blocking_tool_tasks
|
||
|
||
|
||
async def test_parallel_non_blocking_tools():
|
||
"""Tests that multiple NON_BLOCKING tools execute in parallel and clean up independently."""
|
||
import asyncio
|
||
|
||
async def slow_fn_1() -> dict[str, str]:
|
||
await asyncio.sleep(0.05)
|
||
return {'result': 'done_1'}
|
||
|
||
async def slow_fn_2() -> dict[str, str]:
|
||
await asyncio.sleep(0.05)
|
||
return {'result': 'done_2'}
|
||
|
||
tool1 = FunctionTool(slow_fn_1)
|
||
tool1.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE
|
||
tool2 = FunctionTool(slow_fn_2)
|
||
tool2.response_scheduling = types.FunctionResponseScheduling.SILENT
|
||
|
||
model = testing_utils.MockModel.create(responses=[])
|
||
agent = Agent(name='test_agent', model=model, tools=[tool1, tool2])
|
||
invocation_context = await testing_utils.create_invocation_context(
|
||
agent=agent, user_content=''
|
||
)
|
||
invocation_context.live_request_queue = LiveRequestQueue()
|
||
|
||
fc1 = types.FunctionCall(name=tool1.name, args={}, id='fc_1')
|
||
fc2 = types.FunctionCall(name=tool2.name, args={}, id='fc_2')
|
||
event = Event(
|
||
invocation_id=invocation_context.invocation_id,
|
||
author=agent.name,
|
||
content=types.Content(
|
||
parts=[
|
||
types.Part(function_call=fc1),
|
||
types.Part(function_call=fc2),
|
||
]
|
||
),
|
||
)
|
||
|
||
result = await handle_function_calls_live(
|
||
invocation_context, event, {tool1.name: tool1, tool2.name: tool2}
|
||
)
|
||
assert result is None
|
||
assert len(invocation_context.active_non_blocking_tool_tasks) == 2
|
||
key1 = f'{tool1.name}_fc_1'
|
||
key2 = f'{tool2.name}_fc_2'
|
||
assert key1 in invocation_context.active_non_blocking_tool_tasks
|
||
assert key2 in invocation_context.active_non_blocking_tool_tasks
|
||
|
||
contents = await _drain_live_function_responses(
|
||
invocation_context.live_request_queue, count=2
|
||
)
|
||
responses = {
|
||
c.parts[0].function_response.id: c.parts[0].function_response
|
||
for c in contents
|
||
}
|
||
assert responses['fc_1'].response == {'result': 'done_1'}
|
||
assert responses['fc_2'].response == {'result': 'done_2'}
|
||
|
||
await asyncio.sleep(0)
|
||
assert len(invocation_context.active_non_blocking_tool_tasks) == 0
|