From 9f602c9657bde9134bdfd1176907e90ff9bb538e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 3 Aug 2026 20:19:06 -0700 Subject: [PATCH] test(bedrock): cover all four Converse send closures, and warn when the retry fires Only async_streaming had closure-level coverage, so a wiring mistake in the other three (a wrong `request_data` binding, a closure passing the outer body instead of its argument) would have passed the suite. All four now run through a fake transport that rejects the first body the way Bedrock does and accepts the second. The fakes subclass HTTPHandler and AsyncHTTPHandler because the handler discards any client that is not one, which otherwise sends the test to real AWS. The retry also logs a warning naming the provider error. The pre-call log for a retried request records the payload that was refused rather than the one that worked, which matches the upstream retry loop but leaves an operator reading logs with no sign the retry happened; re-firing pre_call would double-count in spend logs and callbacks, so the warning carries that signal instead. It also names the cost-map flag that removes the extra round trip. Folded the decide-log-serialise step into one helper so both retry wrappers share it rather than repeating the sequence. --- litellm/llms/bedrock/chat/converse_handler.py | 34 +++++-- ...test_converse_rejected_tool_field_retry.py | 88 ++++++++++++++++--- 2 files changed, 105 insertions(+), 17 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 69647f28088..020a3959693 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -5,6 +5,7 @@ from typing import Any, TypeVar import httpx import litellm +from litellm._logging import verbose_logger from litellm.anthropic_beta_headers_manager import ( update_headers_with_filtered_beta, ) @@ -29,6 +30,29 @@ from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_ca _SendResultT = TypeVar("_SendResultT") +def _retry_body_without_rejected_tool_fields( + request_data: Mapping[str, Any], err: BedrockError | httpx.HTTPStatusError +) -> str | None: + """ + Serialised retry body with the rejected ``toolSpec`` members gone, or ``None``. + + Warns when it fires: the retry costs a round trip on every request until the model's + ``bedrock_converse_supports_strict_tools`` entry is set, and the pre-call log for + this request records the payload that was refused rather than the one that worked, + so an operator reading logs needs the retry itself to be visible. + """ + retried = drop_bedrock_rejected_tool_fields(request_data, _provider_error_text(err)) + if retried is None: + return None + verbose_logger.warning( + "Bedrock Converse rejected tool fields; retrying once without them. " + "Set bedrock_converse_supports_strict_tools for this model to avoid the extra round trip. " + "Provider error: %s", + _provider_error_text(err), + ) + return json.dumps(retried) + + def _provider_error_text(err: BedrockError | httpx.HTTPStatusError) -> str: """ Read the provider's error body off either shape Converse raises. @@ -152,10 +176,9 @@ class BedrockConverseLLM(BaseAWSLLM): try: return await send(data, headers), data except (BedrockError, httpx.HTTPStatusError) as err: - retried = drop_bedrock_rejected_tool_fields(request_data, _provider_error_text(err)) - if retried is None: + body = _retry_body_without_rejected_tool_fields(request_data, err) + if body is None: raise - body = json.dumps(retried) return await send(body, sign(body)), body def _send_with_tool_field_retry( @@ -171,10 +194,9 @@ class BedrockConverseLLM(BaseAWSLLM): try: return send(data, headers), data except (BedrockError, httpx.HTTPStatusError) as err: - retried = drop_bedrock_rejected_tool_fields(request_data, _provider_error_text(err)) - if retried is None: + body = _retry_body_without_rejected_tool_fields(request_data, err) + if body is None: raise - body = json.dumps(retried) return send(body, sign(body)), body async def async_streaming( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py b/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py index 369d4519046..7dc426933f2 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py @@ -12,6 +12,7 @@ import pytest from litellm.llms.base_llm.base_utils import parse_rejected_tool_fields from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.bedrock.common_utils import BedrockError, drop_bedrock_rejected_tool_fields +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler _STRICT_REJECTION = ( '{"message":"The model returned the following errors: ' @@ -253,7 +254,7 @@ _DESCRIPTION_REJECTION = ( ) -class _RejectThenAcceptClient: +class _RejectThenAcceptTransport: """Fake transport: rejects the first body the way Bedrock does, accepts the second. Drives the real ``_send``/``_send_stream`` closures inside the handler rather than @@ -277,12 +278,32 @@ class _RejectThenAcceptClient: raise httpx.HTTPStatusError( "400", request=request, response=httpx.Response(400, text=_DESCRIPTION_REJECTION, request=request) ) - return httpx.Response(200, json=_CONVERSE_OK) + return httpx.Response(200, json=_CONVERSE_OK, request=httpx.Request("POST", "https://x/y")) + + +class _RejectThenAcceptClient(HTTPHandler, _RejectThenAcceptTransport): + """Injectable sync client. The handler discards anything that is not an HTTPHandler, + so the fake has to be one or the call silently goes to real AWS.""" + + def __init__(self) -> None: + HTTPHandler.__init__(self) + self.posts = [] def post(self, *args, **kwargs) -> httpx.Response: return self._respond(kwargs.get("data") or "") +class _AsyncRejectThenAcceptClient(AsyncHTTPHandler, _RejectThenAcceptTransport): + """Injectable async twin of ``_RejectThenAcceptClient``.""" + + def __init__(self) -> None: + AsyncHTTPHandler.__init__(self) + self.posts = [] + + async def post(self, *args, **kwargs) -> httpx.Response: + return self._respond(kwargs.get("data") or "") + + def _converse_kwargs(client, **overrides): return { "model": "us.anthropic.claude-sonnet-5", @@ -321,6 +342,39 @@ def _raw_credentials(): return Credentials(access_key="AKIAEXAMPLE", secret_key="secret", token=None) +def _assert_retried_without_description(client: _RejectThenAcceptTransport) -> None: + assert len(client.posts) == 2, f"expected exactly one retry, saw {len(client.posts)} request(s)" + assert '"description"' in client.posts[0] + assert '"description"' not in client.posts[1] + assert '"get_weather"' in client.posts[1] + + +def _completion_kwargs(client, stream: bool): + import litellm + + return { + "model": "us.anthropic.claude-sonnet-5", + "messages": [{"role": "user", "content": "hi"}], + "api_base": None, + "custom_prompt_dict": {}, + "model_response": litellm.ModelResponse(), + "encoding": None, + "logging_obj": _logging_obj(), + "optional_params": { + **_converse_kwargs(client)["optional_params"], + "stream": stream, + "fake_stream": True, + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + }, + "acompletion": False, + "timeout": None, + "litellm_params": {"aws_region_name": "us-east-1"}, + "client": client, + "api_key": None, + } + + def _logging_obj(): from unittest.mock import MagicMock @@ -332,14 +386,26 @@ def _logging_obj(): @pytest.mark.asyncio async def test_async_streaming_closure_retries_and_resends_without_the_field() -> None: """Covers the `_send` closure inside async_streaming, not just the wrapper.""" - class _AsyncClient(_RejectThenAcceptClient): - async def post(self, *args, **kwargs): - return self._respond(kwargs.get("data") or "") + client = _AsyncRejectThenAcceptClient() + await BedrockConverseLLM().async_streaming(**_converse_kwargs(client)) + _assert_retried_without_description(client) - async_client = _AsyncClient() - await BedrockConverseLLM().async_streaming(**_converse_kwargs(async_client)) - assert len(async_client.posts) == 2 - assert '"description"' in async_client.posts[0] - assert '"description"' not in async_client.posts[1] - assert '"get_weather"' in async_client.posts[1] +@pytest.mark.asyncio +async def test_async_completion_closure_retries_and_resends_without_the_field() -> None: + """Covers the `_send` closure inside async_completion.""" + client = _AsyncRejectThenAcceptClient() + kwargs = _converse_kwargs(client) + kwargs.pop("fake_stream", None) + kwargs.pop("stream_chunk_size", None) + kwargs["stream"] = False + await BedrockConverseLLM().async_completion(**kwargs) + _assert_retried_without_description(client) + + +@pytest.mark.parametrize("stream", [True, False], ids=["sync streaming", "sync completion"]) +def test_sync_closures_retry_and_resend_without_the_field(stream: bool) -> None: + """Covers the `_send_stream` and `_send` closures inside the sync completion path.""" + client = _RejectThenAcceptClient() + BedrockConverseLLM().completion(**_completion_kwargs(client, stream=stream)) + _assert_retried_without_description(client)