mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(bedrock): answer a buffered invoke stream from the non-streaming reply
This commit is contained in:
parent
6b68ac0063
commit
7e9366f100
3 changed files with 50 additions and 2 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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'}"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue