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 <ankusharma@google.com>
PiperOrigin-RevId: 832383755
This commit is contained in:
Ankur Sharma
2025-11-14 11:02:22 -08:00
committed by Copybara-Service
parent 9b754564b3
commit 696852a280
6 changed files with 168 additions and 1 deletions
@@ -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)
@@ -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:
@@ -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)
@@ -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):
@@ -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:
@@ -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
)