wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
669 lines
23 KiB
Python
669 lines
23 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.
|
|
|
|
from typing import Any
|
|
from unittest import IsolatedAsyncioTestCase
|
|
from unittest.mock import Mock
|
|
|
|
from google.adk.agents.llm_agent import LlmAgent
|
|
from google.adk.plugins.reflect_retry_tool_plugin import REFLECT_AND_RETRY_RESPONSE_TYPE
|
|
from google.adk.plugins.reflect_retry_tool_plugin import ReflectAndRetryToolPlugin
|
|
from google.adk.tools.base_tool import BaseTool
|
|
from google.adk.tools.tool_context import ToolContext
|
|
from google.genai import types
|
|
|
|
from .. import testing_utils
|
|
|
|
|
|
class MockTool(BaseTool):
|
|
"""Mock tool for testing purposes."""
|
|
|
|
def __init__(self, name: str = "mock_tool"):
|
|
self.name = name
|
|
self.description = f"Mock tool named {name}"
|
|
|
|
async def run(self, **kwargs) -> Any:
|
|
return "mock result"
|
|
|
|
|
|
class CustomErrorExtractionPlugin(ReflectAndRetryToolPlugin):
|
|
"""Custom plugin for testing error extraction from tool responses."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.error_conditions = {}
|
|
|
|
def set_error_condition(self, condition_func):
|
|
"""Set a custom error condition function for testing."""
|
|
self.error_condition = condition_func
|
|
|
|
async def extract_error_from_result(
|
|
self, *, tool, tool_args, tool_context, result
|
|
):
|
|
"""Extract error based on custom conditions set for testing."""
|
|
if hasattr(self, "error_condition"):
|
|
return self.error_condition(result)
|
|
return None
|
|
|
|
|
|
# Inheriting from IsolatedAsyncioTestCase ensures consistent behavior.
|
|
# See https://github.com/pytest-dev/pytest-asyncio/issues/1039
|
|
class TestReflectAndRetryToolPlugin(IsolatedAsyncioTestCase):
|
|
"""Comprehensive tests for ReflectAndRetryToolPlugin focusing on behavior."""
|
|
|
|
def get_plugin(self):
|
|
"""Create a default plugin instance for testing."""
|
|
return ReflectAndRetryToolPlugin()
|
|
|
|
def get_custom_plugin(self):
|
|
"""Create a plugin with custom parameters."""
|
|
return ReflectAndRetryToolPlugin(
|
|
name="custom_plugin",
|
|
max_retries=5,
|
|
throw_exception_if_retry_exceeded=False,
|
|
)
|
|
|
|
def get_mock_tool(self):
|
|
"""Create a mock tool for testing."""
|
|
return MockTool("test_tool_id")
|
|
|
|
def get_mock_tool_context(self):
|
|
"""Create a mock tool context."""
|
|
return Mock(spec=ToolContext)
|
|
|
|
def get_custom_error_plugin(self):
|
|
"""Create a custom error extraction plugin for testing."""
|
|
return CustomErrorExtractionPlugin(max_retries=3)
|
|
|
|
def get_sample_tool_args(self):
|
|
"""Sample tool arguments for testing."""
|
|
return {"param1": "value1", "param2": 42, "param3": True}
|
|
|
|
async def test_plugin_initialization_default(self):
|
|
"""Test plugin initialization with default parameters."""
|
|
plugin = self.get_plugin()
|
|
|
|
self.assertEqual(plugin.name, "reflect_retry_tool_plugin")
|
|
self.assertEqual(plugin.max_retries, 3)
|
|
self.assertIs(plugin.throw_exception_if_retry_exceeded, True)
|
|
|
|
async def test_plugin_initialization_custom(self):
|
|
"""Test plugin initialization with custom parameters."""
|
|
plugin = ReflectAndRetryToolPlugin(
|
|
name="custom_name",
|
|
max_retries=10,
|
|
throw_exception_if_retry_exceeded=False,
|
|
)
|
|
|
|
self.assertEqual(plugin.name, "custom_name")
|
|
self.assertEqual(plugin.max_retries, 10)
|
|
self.assertIsNot(plugin.throw_exception_if_retry_exceeded, True)
|
|
|
|
async def test_after_tool_callback_successful_call(self):
|
|
"""Test after_tool_callback with successful tool call."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
result = {"success": True, "data": "test_data"}
|
|
|
|
callback_result = await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=result,
|
|
)
|
|
|
|
# Should return None for successful calls
|
|
self.assertIsNone(callback_result)
|
|
|
|
async def test_after_tool_callback_ignore_retry_response(self):
|
|
"""Test that retry responses are ignored in after_tool_callback."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
retry_result = {"response_type": REFLECT_AND_RETRY_RESPONSE_TYPE}
|
|
|
|
callback_result = await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=retry_result,
|
|
)
|
|
|
|
# Retry responses should be ignored
|
|
self.assertIsNone(callback_result)
|
|
|
|
async def test_on_tool_error_callback_max_retries_zero(self):
|
|
"""Test error callback when max_retries is 0.
|
|
|
|
This should return None so that the exception is rethrown
|
|
"""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = ReflectAndRetryToolPlugin(max_retries=0)
|
|
error = ValueError("Test error")
|
|
|
|
with self.assertRaises(ValueError) as cm:
|
|
await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Should re-raise the original exception when max_retries is 0
|
|
self.assertIs(cm.exception, error)
|
|
|
|
async def test_on_tool_error_callback_max_retries_zero_without_exception(
|
|
self,
|
|
):
|
|
"""Test error callback when max_retries is 0 and exception is disabled."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = ReflectAndRetryToolPlugin(
|
|
max_retries=0, throw_exception_if_retry_exceeded=False
|
|
)
|
|
error = ValueError("Test error")
|
|
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Should return a retry exceeded message instead of raising
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
|
|
self.assertEqual(result["error_type"], "ValueError")
|
|
self.assertEqual(result["retry_count"], 0)
|
|
self.assertIn(
|
|
"the retry limit has been exceeded", result["reflection_guidance"]
|
|
)
|
|
|
|
async def test_on_tool_error_callback_max_retries_zero_with_dict_error(self):
|
|
"""Test error callback when max_retries is 0 and error is a dict."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = CustomErrorExtractionPlugin(
|
|
max_retries=0, throw_exception_if_retry_exceeded=True
|
|
)
|
|
dict_error = {"status": "error", "message": "Custom dict error"}
|
|
plugin.set_error_condition(lambda result: dict_error)
|
|
|
|
with self.assertRaises(Exception) as cm:
|
|
await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result={"some": "result"},
|
|
)
|
|
|
|
# Should raise an Exception wrapping the dict
|
|
self.assertNotIsInstance(cm.exception, TypeError)
|
|
self.assertIn("Custom dict error", str(cm.exception))
|
|
|
|
async def test_on_tool_error_callback_first_failure(self):
|
|
"""Test first tool failure creates reflection response."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
error = ValueError("Test error message")
|
|
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
|
|
self.assertEqual(result["error_type"], "ValueError")
|
|
self.assertEqual(result["error_details"], "Test error message")
|
|
self.assertEqual(result["retry_count"], 1)
|
|
self.assertIn("test_tool_id", result["reflection_guidance"])
|
|
self.assertIn("Test error message", result["reflection_guidance"])
|
|
|
|
async def test_retry_behavior_with_consecutive_failures(self):
|
|
"""Test the retry behavior with consecutive failures."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
error = RuntimeError("Runtime error")
|
|
|
|
# First failure
|
|
result1 = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
self.assertEqual(result1["retry_count"], 1)
|
|
|
|
# Second failure - should have different retry count based on plugin logic
|
|
result2 = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
# The plugin's internal logic determines the exact retry count
|
|
self.assertIsNotNone(result2)
|
|
self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
|
|
self.assertEqual(result2["retry_count"], 2)
|
|
|
|
async def test_different_tools_behavior(self):
|
|
"""Test behavior when using different tools."""
|
|
plugin = self.get_plugin()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
tool1 = MockTool("tool1")
|
|
tool2 = MockTool("tool2")
|
|
error = ValueError("Test error")
|
|
|
|
# First failure on tool1
|
|
result1 = await plugin.on_tool_error_callback(
|
|
tool=tool1,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
self.assertEqual(result1["retry_count"], 1)
|
|
|
|
# Failure on tool2
|
|
result2 = await plugin.on_tool_error_callback(
|
|
tool=tool2,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
# Since tool is different, retry count should start over.
|
|
self.assertIsNotNone(result2)
|
|
self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
|
|
self.assertEqual(result2["retry_count"], 1)
|
|
|
|
async def test_max_retries_exceeded_with_exception(self):
|
|
"""Test that original exception is raised when max retries exceeded."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = ReflectAndRetryToolPlugin(
|
|
max_retries=1, throw_exception_if_retry_exceeded=True
|
|
)
|
|
error = ConnectionError("Connection failed")
|
|
|
|
# First call should succeed and return a retry response
|
|
await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Second call should exceed max_retries and raise
|
|
with self.assertRaises(ConnectionError) as cm:
|
|
await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Verify exception properties
|
|
self.assertIs(cm.exception, error)
|
|
|
|
async def test_max_retries_exceeded_with_dict_error(self):
|
|
"""Test that Exception is raised when max retries exceeded with dict error."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = CustomErrorExtractionPlugin(
|
|
max_retries=1, throw_exception_if_retry_exceeded=True
|
|
)
|
|
dict_error = {"status": "error", "message": "Custom dict error"}
|
|
plugin.set_error_condition(lambda result: dict_error)
|
|
|
|
# First call should fail and return a retry response
|
|
result1 = await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result={"some": "result"},
|
|
)
|
|
self.assertIsNotNone(result1)
|
|
self.assertEqual(result1["retry_count"], 1)
|
|
|
|
# Second call should exceed max_retries and raise
|
|
with self.assertRaises(Exception) as cm:
|
|
await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result={"some": "result"},
|
|
)
|
|
|
|
# Verify exception properties
|
|
self.assertNotIsInstance(cm.exception, TypeError)
|
|
self.assertIn("Custom dict error", str(cm.exception))
|
|
|
|
async def test_max_retries_exceeded_without_exception(self):
|
|
"""Test max retries exceeded returns failure message when exception is disabled."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = ReflectAndRetryToolPlugin(
|
|
max_retries=2, throw_exception_if_retry_exceeded=False
|
|
)
|
|
error = TimeoutError("Timeout occurred")
|
|
|
|
# Call until we exceed the retry limit
|
|
result = None
|
|
for _ in range(3):
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Should get a retry exceeded message on the last call
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
|
|
self.assertEqual(result["error_type"], "TimeoutError")
|
|
self.assertIn(
|
|
"the retry limit has been exceeded", result["reflection_guidance"]
|
|
)
|
|
self.assertIn("Do not attempt to use the", result["reflection_guidance"])
|
|
|
|
async def test_successful_call_resets_retry_behavior(self):
|
|
"""Test that successful calls reset the retry behavior."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
error = ValueError("Test error")
|
|
|
|
# First failure
|
|
result1 = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
self.assertEqual(result1["retry_count"], 1)
|
|
|
|
# Successful call
|
|
await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result={"success": True},
|
|
)
|
|
|
|
# Next failure should start fresh
|
|
result2 = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
self.assertEqual(result2["retry_count"], 1) # Should restart from 1
|
|
|
|
async def test_none_result_handling(self):
|
|
"""Test handling of None results in after_tool_callback."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
|
|
# None result should be handled gracefully
|
|
callback_result = await plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=None,
|
|
)
|
|
|
|
self.assertIsNone(callback_result)
|
|
|
|
async def test_empty_tool_args_handling(self):
|
|
"""Test handling of empty tool arguments."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
empty_args = {}
|
|
error = ValueError("Test error")
|
|
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=empty_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
# Empty args should be represented in the response
|
|
self.assertIn("{}", result["reflection_guidance"])
|
|
|
|
async def test_retry_count_progression(self):
|
|
"""Test that retry counts progress correctly for the same tool."""
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
plugin = ReflectAndRetryToolPlugin(max_retries=5)
|
|
error = ValueError("Test error")
|
|
tool = MockTool("single_tool")
|
|
|
|
for i in range(1, 4):
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
self.assertEqual(result["retry_count"], i)
|
|
|
|
async def test_max_retries_parameter_behavior(self):
|
|
"""Test that max_retries parameter affects behavior correctly."""
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
# Test with very low max_retries
|
|
plugin = ReflectAndRetryToolPlugin(
|
|
max_retries=1, throw_exception_if_retry_exceeded=False
|
|
)
|
|
error = ValueError("Test error")
|
|
|
|
# First call is fine
|
|
await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Second call exceeds limit
|
|
result = await plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=error,
|
|
)
|
|
|
|
# Should hit max retries quickly with max_retries=1
|
|
self.assertIn(
|
|
"the retry limit has been exceeded.", result["reflection_guidance"]
|
|
)
|
|
|
|
async def test_default_extract_error_returns_none(self):
|
|
"""Test that default extract_error_from_result returns None."""
|
|
plugin = self.get_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
result = {"status": "success", "data": "some data"}
|
|
|
|
error = await plugin.extract_error_from_result(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=result,
|
|
)
|
|
self.assertIsNone(error)
|
|
|
|
async def test_custom_error_detection_and_success_handling(self):
|
|
"""Test custom error detection, success handling, and retry progression."""
|
|
custom_error_plugin = self.get_custom_error_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
custom_error_plugin.set_error_condition(
|
|
lambda result: result if result.get("status") == "error" else None
|
|
)
|
|
|
|
# Test error detection
|
|
error_result = {"status": "error", "message": "Something went wrong"}
|
|
callback_result = await custom_error_plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=error_result,
|
|
)
|
|
self.assertIsNotNone(callback_result)
|
|
self.assertEqual(
|
|
callback_result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE
|
|
)
|
|
self.assertEqual(callback_result["retry_count"], 1)
|
|
|
|
# Test success handling
|
|
success_result = {"status": "success", "data": "operation completed"}
|
|
callback_result = await custom_error_plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=success_result,
|
|
)
|
|
self.assertIsNone(callback_result)
|
|
|
|
async def test_retry_state_management(self):
|
|
"""Test retry state management with custom errors and mixed error types."""
|
|
custom_error_plugin = self.get_custom_error_plugin()
|
|
mock_tool = self.get_mock_tool()
|
|
mock_tool_context = self.get_mock_tool_context()
|
|
sample_tool_args = self.get_sample_tool_args()
|
|
custom_error_plugin.set_error_condition(
|
|
lambda result: result if result.get("failed") else None
|
|
)
|
|
|
|
# Custom error followed by exception
|
|
custom_error = {"failed": True, "reason": "Network timeout"}
|
|
result1 = await custom_error_plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=custom_error,
|
|
)
|
|
self.assertEqual(result1["retry_count"], 1)
|
|
|
|
# Exception should increment retry count
|
|
exception = ValueError("Invalid parameter")
|
|
result2 = await custom_error_plugin.on_tool_error_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
error=exception,
|
|
)
|
|
self.assertEqual(result2["retry_count"], 2)
|
|
|
|
# Success should reset
|
|
success = {"result": "success"}
|
|
result3 = await custom_error_plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=success,
|
|
)
|
|
self.assertIsNone(result3)
|
|
|
|
# Next error should start fresh
|
|
result4 = await custom_error_plugin.after_tool_callback(
|
|
tool=mock_tool,
|
|
tool_args=sample_tool_args,
|
|
tool_context=mock_tool_context,
|
|
result=custom_error,
|
|
)
|
|
self.assertEqual(result4["retry_count"], 1)
|
|
|
|
async def test_hallucinating_tool_name(self):
|
|
"""Test that hallucinating tool name is handled correctly."""
|
|
wrong_function_call = types.Part.from_function_call(
|
|
name="increase_by_one", args={"x": 1}
|
|
)
|
|
correct_function_call = types.Part.from_function_call(
|
|
name="increase", args={"x": 1}
|
|
)
|
|
responses: list[types.Content] = [
|
|
wrong_function_call,
|
|
correct_function_call,
|
|
"response1",
|
|
]
|
|
mock_model = testing_utils.MockModel.create(responses=responses)
|
|
|
|
function_called = 0
|
|
|
|
def increase(x: int) -> int:
|
|
nonlocal function_called
|
|
function_called += 1
|
|
return x + 1
|
|
|
|
agent = LlmAgent(name="root_agent", model=mock_model, tools=[increase])
|
|
runner = testing_utils.TestInMemoryRunner(
|
|
agent=agent, plugins=[self.get_plugin()]
|
|
)
|
|
|
|
events = await runner.run_async_with_new_session("test")
|
|
# Filter out agent_state events (no content).
|
|
events = [e for e in events if e.content is not None]
|
|
|
|
# Assert that the first event is a function call with the wrong name
|
|
assert events[0].content.parts[0].function_call.name == "increase_by_one"
|
|
|
|
# Assert that the second event is a function response with the
|
|
# reflection_guidance
|
|
assert (
|
|
events[1].content.parts[0].function_response.response["error_type"]
|
|
== "ValueError"
|
|
)
|
|
assert (
|
|
events[1].content.parts[0].function_response.response["retry_count"]
|
|
== 1
|
|
)
|
|
assert (
|
|
"Wrong Function Name"
|
|
in events[1]
|
|
.content.parts[0]
|
|
.function_response.response["reflection_guidance"]
|
|
)
|
|
|
|
# Assert that the third event is a function call with the correct name
|
|
assert events[2].content.parts[0].function_call.name == "increase"
|
|
self.assertEqual(function_called, 1)
|