mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(bedrock): report the retried body, not the rejected one, after a successful retry
The retry wrappers returned only the response, so the two non-streaming Converse paths kept passing their original `data` into `_transform_response`. A request that succeeded on the second attempt therefore logged the payload Bedrock had just rejected, claiming a field was sent that was in fact removed. That is exactly backwards for the debugging this feature exists to support, and it diverges from the established pattern in `llm_http_handler.py`, where the retry loop rebinds `data` so everything downstream sees the corrected body. The wrappers now return the body that actually produced the response alongside the result, and every call site binds it. The streaming paths were already correct, since they pass the retried body through to `make_call` themselves; this makes all four consistent.
This commit is contained in:
parent
37c8efde0f
commit
5ee1b09bd5
2 changed files with 32 additions and 12 deletions
|
|
@ -149,16 +149,19 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
caller_headers: Mapping[str, str],
|
||||
endpoint_url: str,
|
||||
api_key: str | None,
|
||||
) -> _SendResultT:
|
||||
) -> tuple[_SendResultT, str]:
|
||||
"""
|
||||
Send once, and if Bedrock rejects extra ``toolSpec`` members, drop them and send again.
|
||||
|
||||
``send`` owns the transport and the provider-error contract, so a request that
|
||||
fails for any other reason raises exactly what it raised before. The retry is
|
||||
single-shot: a second rejection surfaces rather than looping.
|
||||
|
||||
Returns the result together with the body that actually produced it, so callers
|
||||
log and transform against what was sent rather than the payload that was rejected.
|
||||
"""
|
||||
try:
|
||||
return await send(data, headers)
|
||||
return await send(data, headers), data
|
||||
except (BedrockError, httpx.HTTPStatusError) as err:
|
||||
retry = self._resign_without_rejected_tool_fields(
|
||||
request_data=request_data,
|
||||
|
|
@ -172,7 +175,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
if retry is None:
|
||||
raise
|
||||
return await send(*retry)
|
||||
retry_data, retry_headers = retry
|
||||
return await send(retry_data, retry_headers), retry_data
|
||||
|
||||
def _send_retrying_rejected_tool_fields(
|
||||
self,
|
||||
|
|
@ -186,10 +190,10 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
caller_headers: Mapping[str, str],
|
||||
endpoint_url: str,
|
||||
api_key: str | None,
|
||||
) -> _SendResultT:
|
||||
) -> tuple[_SendResultT, str]:
|
||||
"""Synchronous twin of ``_asend_retrying_rejected_tool_fields``."""
|
||||
try:
|
||||
return send(data, headers)
|
||||
return send(data, headers), data
|
||||
except (BedrockError, httpx.HTTPStatusError) as err:
|
||||
retry = self._resign_without_rejected_tool_fields(
|
||||
request_data=request_data,
|
||||
|
|
@ -203,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
if retry is None:
|
||||
raise
|
||||
return send(*retry)
|
||||
retry_data, retry_headers = retry
|
||||
return send(retry_data, retry_headers), retry_data
|
||||
|
||||
async def async_streaming(
|
||||
self,
|
||||
|
|
@ -270,7 +275,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
stream_chunk_size=stream_chunk_size,
|
||||
)
|
||||
|
||||
completion_stream = await self._asend_retrying_rejected_tool_fields(
|
||||
completion_stream, data = await self._asend_retrying_rejected_tool_fields(
|
||||
send=_send,
|
||||
request_data=request_data,
|
||||
data=data,
|
||||
|
|
@ -364,7 +369,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
response = await self._asend_retrying_rejected_tool_fields(
|
||||
response, data = await self._asend_retrying_rejected_tool_fields(
|
||||
send=_send,
|
||||
request_data=request_data,
|
||||
data=data,
|
||||
|
|
@ -606,7 +611,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
stream_chunk_size=stream_chunk_size,
|
||||
)
|
||||
|
||||
completion_stream = self._send_retrying_rejected_tool_fields(
|
||||
completion_stream, data = self._send_retrying_rejected_tool_fields(
|
||||
send=_send_stream,
|
||||
request_data=_data,
|
||||
data=data,
|
||||
|
|
@ -643,7 +648,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
response = self._send_retrying_rejected_tool_fields(
|
||||
response, data = self._send_retrying_rejected_tool_fields(
|
||||
send=_send,
|
||||
request_data=_data,
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -125,10 +125,12 @@ def test_sync_retry_resends_without_the_rejected_field_and_resigns(raised: Excep
|
|||
raise raised
|
||||
return "ok"
|
||||
|
||||
result = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
|
||||
result, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
|
||||
|
||||
assert result == "ok"
|
||||
assert len(attempts) == 2
|
||||
assert sent_body == attempts[1][0]
|
||||
assert '"strict"' not in sent_body
|
||||
|
||||
first_body, first_headers = attempts[0]
|
||||
retry_body, retry_headers = attempts[1]
|
||||
|
|
@ -178,12 +180,25 @@ async def test_async_retry_resends_without_the_rejected_field_and_resigns() -> N
|
|||
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
|
||||
return "ok"
|
||||
|
||||
result = await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
|
||||
result, sent_body = await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(
|
||||
send=send, **_retry_kwargs()
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert len(attempts) == 2
|
||||
assert '"strict"' not in attempts[1][0]
|
||||
assert attempts[1][1]["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert sent_body == attempts[1][0]
|
||||
|
||||
|
||||
def test_reported_body_is_the_original_when_no_retry_happens() -> None:
|
||||
"""A request that succeeds first time reports exactly what it sent."""
|
||||
|
||||
def send(body: str, headers) -> str:
|
||||
return "ok"
|
||||
|
||||
_, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
|
||||
assert sent_body == "original-body"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue