You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
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:
committed by
Copybara-Service
parent
9b754564b3
commit
696852a280
@@ -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
|
||||
)
|
||||
Reference in New Issue
Block a user