fix(bedrock): answer a buffered invoke stream from the non-streaming reply

This commit is contained in:
mateo-berri 2026-11-06 05:47:56 -08:00
parent 6b68ac0063
commit 7e9366f100
3 changed files with 50 additions and 2 deletions

View file

@ -1056,7 +1056,7 @@ class BaseLLMHTTPHandler:
provider_config=provider_config,
fake_stream=fake_stream,
)
if provider_config.has_custom_stream_wrapper is True:
if provider_config.has_custom_stream_wrapper is True and not fake_stream:
return provider_config.get_sync_custom_stream_wrapper(
model=model,
custom_llm_provider=custom_llm_provider,
@ -1225,7 +1225,7 @@ class BaseLLMHTTPHandler:
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
):
if provider_config.has_custom_stream_wrapper is True:
if provider_config.has_custom_stream_wrapper is True and not fake_stream:
return await provider_config.get_async_custom_stream_wrapper(
model=model,
custom_llm_provider=custom_llm_provider,

View file

@ -640,6 +640,50 @@ async def test_a_denied_native_output_stream_is_sent_again_buffered_with_native_
assert "textFormat" in retry_body.get("outputConfig", {}) and "toolConfig" not in retry_body
_INVOKE_ANSWER: Final = json.dumps(
{
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-opus-5-5",
"content": [{"type": "text", "text": '{"answer": "hi"}'}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
).encode()
def _denies_invoke_streaming_then_answers(request: httpx.Request) -> httpx.Response:
if request.url.path.endswith("/invoke-with-response-stream"):
return httpx.Response(403, content=_STREAM_PERMISSION_DENIAL, request=request)
return httpx.Response(200, content=_INVOKE_ANSWER, headers={"content-type": "application/json"}, request=request)
@pytest.mark.asyncio
async def test_a_denied_native_output_stream_on_the_invoke_route_is_sent_again_buffered_to_invoke(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_denies_invoke_streaming_then_answers)
router: Final = _bedrock_router(
"bedrock/invoke/us.anthropic.claude-opus-5-5", timeout=60, stream_timeout=STREAM_TIMEOUT_SECONDS
)
response: Final = await router.acompletion(
model="bedrock",
messages=[{"role": "user", "content": "hi"}],
stream=True,
thinking={"type": "disabled"},
response_format={"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}},
)
assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == '{"answer": "hi"}'
assert _sent_paths(respx_mock) == ["invoke-with-response-stream", "invoke"]
retry_body: Final = json.loads(respx_mock.calls[1].request.content)
assert "format" in retry_body.get("output_config", {}) and "tools" not in retry_body
@pytest.mark.asyncio
async def test_failed_native_output_retry_reports_the_request_timeout_it_was_bounded_by(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch

View file

@ -4894,6 +4894,10 @@ class _StreamRouteConfig(_TransformRecordingConfig):
def __init__(self):
super().__init__(transform_async=False)
@property
def has_custom_stream_wrapper(self) -> bool:
return True
def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None):
return f"{api_base}/{'stream' if stream else 'answer'}"