From 696852a28095a024cbe76413ee7617356e19a9e3 Mon Sep 17 00:00:00 2001 From: Ankur Sharma Date: Fri, 14 Nov 2025 11:02:22 -0800 Subject: [PATCH] chore: Add default retry options as fall back to llm_request that are made during evals In order to make evals more resilient to temporary model failures, we add retry options to llm_requests that are made during evals. Note that this is a fall back option, if the developer has already specified their own retry options, then those will be honored. Co-authored-by: Ankur Sharma PiperOrigin-RevId: 832383755 --- .../adk/evaluation/_retry_options_utils.py | 75 ++++++++++++++++++ .../adk/evaluation/evaluation_generator.py | 9 ++- .../adk/evaluation/hallucinations_v1.py | 3 + src/google/adk/evaluation/llm_as_judge.py | 2 + .../evaluation/llm_backed_user_simulator.py | 2 + .../evaluation/test_retry_options_utils.py | 78 +++++++++++++++++++ 6 files changed, 168 insertions(+), 1 deletion(-) create mode 100644 src/google/adk/evaluation/_retry_options_utils.py create mode 100644 tests/unittests/evaluation/test_retry_options_utils.py diff --git a/src/google/adk/evaluation/_retry_options_utils.py b/src/google/adk/evaluation/_retry_options_utils.py new file mode 100644 index 00000000..e5c83875 --- /dev/null +++ b/src/google/adk/evaluation/_retry_options_utils.py @@ -0,0 +1,75 @@ +# 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 __future__ import annotations + +from typing import Optional + +from google.genai import types +from typing_extensions import override + +from ..agents.callback_context import CallbackContext +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..plugins.base_plugin import BasePlugin + +_RETRY_HTTP_STATUS_CODES = ( + 408, # Request timeout. + 429, # Too many requests. + 500, # Internal server error. + 502, # Bad gateway. + 503, # Service unavailable. + 504, # Gateway timeout +) +_DEFAULT_HTTP_RETRY_OPTIONS = types.HttpRetryOptions( + attempts=7, + initial_delay=5.0, + max_delay=120, + exp_base=2.0, + http_status_codes=_RETRY_HTTP_STATUS_CODES, +) + + +def add_default_retry_options_if_not_present(llm_request: LlmRequest): + """Adds default HTTP Retry Options, if they are not present on the llm_request. + + NOTE: This implementation is intended for eval systems internal usage. Do not + take direct dependency on it. + """ + llm_request.config = llm_request.config or types.GenerateContentConfig() + + llm_request.config.http_options = ( + llm_request.config.http_options or types.HttpOptions() + ) + llm_request.config.http_options.retry_options = ( + llm_request.config.http_options.retry_options + or _DEFAULT_HTTP_RETRY_OPTIONS + ) + + +class EnsureRetryOptionsPlugin(BasePlugin): + """This plugin adds retry options to llm_request, if they are not present. + + This is done to ensure that temporary outages with the model provider don't + affect eval runs. + + NOTE: This implementation is intended for eval systems internal usage. Do not + take direct dependency on it. + """ + + @override + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + add_default_retry_options_if_not_present(llm_request) diff --git a/src/google/adk/evaluation/evaluation_generator.py b/src/google/adk/evaluation/evaluation_generator.py index 970ebd8f..e9c7dc54 100644 --- a/src/google/adk/evaluation/evaluation_generator.py +++ b/src/google/adk/evaluation/evaluation_generator.py @@ -35,6 +35,7 @@ from ..sessions.base_session_service import BaseSessionService from ..sessions.in_memory_session_service import InMemorySessionService from ..sessions.session import Session from ..utils.context_utils import Aclosing +from ._retry_options_utils import EnsureRetryOptionsPlugin from .app_details import AgentDetails from .app_details import AppDetails from .eval_case import EvalCase @@ -225,13 +226,19 @@ class EvaluationGenerator: request_intercepter_plugin = _RequestIntercepterPlugin( name="request_intercepter_plugin" ) + # We ensure that there is some kind of retries on the llm_requests that are + # generated from the Agent. This is done to make inferencing step of evals + # more resilient to temporary model failures. + ensure_retry_options_plugin = EnsureRetryOptionsPlugin( + name="ensure_retry_options" + ) async with Runner( app_name=app_name, agent=root_agent, artifact_service=artifact_service, session_service=session_service, memory_service=memory_service, - plugins=[request_intercepter_plugin], + plugins=[request_intercepter_plugin, ensure_retry_options_plugin], ) as runner: events = [] while True: diff --git a/src/google/adk/evaluation/hallucinations_v1.py b/src/google/adk/evaluation/hallucinations_v1.py index e2797dfd..cf03299f 100644 --- a/src/google/adk/evaluation/hallucinations_v1.py +++ b/src/google/adk/evaluation/hallucinations_v1.py @@ -32,6 +32,7 @@ from ..models.llm_response import LlmResponse from ..models.registry import LLMRegistry from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._retry_options_utils import add_default_retry_options_if_not_present from .app_details import AppDetails from .eval_case import Invocation from .eval_case import InvocationEvent @@ -526,6 +527,7 @@ class HallucinationsV1Evaluator(Evaluator): ], config=self._model_config, ) + add_default_retry_options_if_not_present(segmenter_llm_request) try: async with Aclosing( self._judge_model.generate_content_async(segmenter_llm_request) @@ -559,6 +561,7 @@ class HallucinationsV1Evaluator(Evaluator): ], config=self._model_config, ) + add_default_retry_options_if_not_present(validator_llm_request) try: async with Aclosing( self._judge_model.generate_content_async(validator_llm_request) diff --git a/src/google/adk/evaluation/llm_as_judge.py b/src/google/adk/evaluation/llm_as_judge.py index 652a12b5..633caeff 100644 --- a/src/google/adk/evaluation/llm_as_judge.py +++ b/src/google/adk/evaluation/llm_as_judge.py @@ -27,6 +27,7 @@ from ..models.llm_response import LlmResponse from ..models.registry import LLMRegistry from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._retry_options_utils import add_default_retry_options_if_not_present from .common import EvalBaseModel from .eval_case import Invocation from .eval_metrics import BaseCriterion @@ -142,6 +143,7 @@ class LlmAsJudge(Evaluator): ], config=self._judge_model_options.judge_model_config, ) + add_default_retry_options_if_not_present(llm_request) num_samples = self._judge_model_options.num_samples invocation_result_samples = [] for _ in range(num_samples): diff --git a/src/google/adk/evaluation/llm_backed_user_simulator.py b/src/google/adk/evaluation/llm_backed_user_simulator.py index 511d9f3c..2fbfcc44 100644 --- a/src/google/adk/evaluation/llm_backed_user_simulator.py +++ b/src/google/adk/evaluation/llm_backed_user_simulator.py @@ -27,6 +27,7 @@ from ..models.llm_request import LlmRequest from ..models.registry import LLMRegistry from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._retry_options_utils import add_default_retry_options_if_not_present from .conversation_scenarios import ConversationScenario from .evaluator import Evaluator from .user_simulator import BaseUserSimulatorConfig @@ -200,6 +201,7 @@ class LlmBackedUserSimulator(UserSimulator): ), ], ) + add_default_retry_options_if_not_present(llm_request) response = "" async with Aclosing(self._llm.generate_content_async(llm_request)) as agen: diff --git a/tests/unittests/evaluation/test_retry_options_utils.py b/tests/unittests/evaluation/test_retry_options_utils.py new file mode 100644 index 00000000..e3ff4f7c --- /dev/null +++ b/tests/unittests/evaluation/test_retry_options_utils.py @@ -0,0 +1,78 @@ +# 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 google.adk.agents.callback_context import CallbackContext +from google.adk.evaluation import _retry_options_utils +from google.adk.models.llm_request import LlmRequest +from google.genai import types +import pytest + + +def test_add_retry_options_with_default_request(): + request = LlmRequest() + _retry_options_utils.add_default_retry_options_if_not_present(request) + assert request.config.http_options is not None + assert ( + request.config.http_options.retry_options + == _retry_options_utils._DEFAULT_HTTP_RETRY_OPTIONS + ) + + +def test_add_retry_options_when_retry_options_is_none(): + request = LlmRequest() + request.config.http_options = types.HttpOptions(retry_options=None) + _retry_options_utils.add_default_retry_options_if_not_present(request) + assert ( + request.config.http_options.retry_options + == _retry_options_utils._DEFAULT_HTTP_RETRY_OPTIONS + ) + + +def test_add_retry_options_does_not_override_existing_options(): + my_retry_options = types.HttpRetryOptions(attempts=1) + request = LlmRequest() + request.config.http_options = types.HttpOptions( + retry_options=my_retry_options + ) + _retry_options_utils.add_default_retry_options_if_not_present(request) + assert request.config.http_options.retry_options == my_retry_options + + +def test_add_retry_options_when_config_is_none(): + request = LlmRequest() + request.config = None + _retry_options_utils.add_default_retry_options_if_not_present(request) + assert request.config is not None + assert request.config.http_options is not None + assert ( + request.config.http_options.retry_options + == _retry_options_utils._DEFAULT_HTTP_RETRY_OPTIONS + ) + + +@pytest.mark.asyncio +async def test_ensure_retry_options_plugin(mocker): + request = LlmRequest() + plugin = _retry_options_utils.EnsureRetryOptionsPlugin(name="test_plugin") + mock_invocation_context = mocker.MagicMock() + mock_invocation_context.session.state = {} + callback_context = CallbackContext(mock_invocation_context) + await plugin.before_model_callback( + callback_context=callback_context, llm_request=request + ) + assert request.config.http_options is not None + assert ( + request.config.http_options.retry_options + == _retry_options_utils._DEFAULT_HTTP_RETRY_OPTIONS + )