diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 5012ed01c1e..c27a3c3528c 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -200,23 +200,12 @@ class SagemakerLLM(BaseAWSLLM): # Add model_id as InferenceComponentName header # boto3 doc: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) - sync_handler = _get_httpx_client() - sync_response = sync_handler.post( - url=prepared_request.url, + completion_stream = self.make_sync_call( + api_base=prepared_request.url, headers=prepared_request.headers, # type: ignore - data=prepared_request.body, - stream=stream, + data=cast(str, prepared_request.body), # cast-ok: signed body is a JSON str, mirrors async path + logging_obj=logging_obj, ) - - if sync_response.status_code != 200: - raise SagemakerError( - status_code=sync_response.status_code, - message=str(sync_response.read()), - ) - - decoder = AWSEventStreamDecoder(model="") - - completion_stream = decoder.iter_bytes(sync_response.iter_bytes()) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -334,6 +323,29 @@ class SagemakerLLM(BaseAWSLLM): litellm_params=litellm_params, ) + def make_sync_call( + self, + api_base: str, + headers: dict, + data: str, + logging_obj, + client=None, + ): + if client is None: + client = _get_httpx_client() + sync_response = client.post( + api_base, + headers=headers, + data=data, + stream=True, + ) + + if sync_response.status_code != 200: + raise SagemakerError(status_code=sync_response.status_code, message=str(sync_response.read())) + + decoder = AWSEventStreamDecoder(model="") + return decoder.iter_bytes(sync_response.iter_bytes()) + async def make_async_call( self, api_base: str, diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py index d43c07e09bb..c2cb46902b8 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -12,7 +12,7 @@ inflating TTFT and turning a steady provider stream into gap-then-burst delivery import binascii import json import struct -from typing import AsyncIterator +from typing import AsyncIterator, Iterator from unittest.mock import MagicMock import httpx @@ -55,6 +55,19 @@ def _make_frames(n: int) -> list[bytes]: return frames +class _CountingSyncStream(httpx.SyncByteStream): + """Yields provider frames one at a time and records how many have been pulled.""" + + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + def __iter__(self) -> Iterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + class _CountingAsyncStream(httpx.AsyncByteStream): """Yields provider frames one at a time and records how many have been pulled.""" @@ -68,6 +81,14 @@ class _CountingAsyncStream(httpx.AsyncByteStream): yield frame +class _FakeSyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + class _FakeAsyncClient: def __init__(self, response: httpx.Response) -> None: self._response = response @@ -76,6 +97,36 @@ class _FakeAsyncClient: return self._response +def test_sync_native_streaming_forwards_each_frame_incrementally(): + """Each token must be emitted after exactly one newly-pulled source frame. + + With the old `chunk_size=1024` the httpx chunker would swallow several small + frames before yielding, so the first token would arrive only after `consumed` + had already crossed multiple frames, and tokens would then replay in a burst. + """ + frames = _make_frames(24) + stream = _CountingSyncStream(frames) + response = httpx.Response(200, stream=stream) + + completion_stream = SagemakerLLM().make_sync_call( + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data="", + logging_obj=MagicMock(), + client=_FakeSyncClient(response), + ) + + consumed_at_token = [] + texts = [] + for chunk in completion_stream: + if chunk is not None and chunk["text"]: + consumed_at_token.append(stream.consumed) + texts.append(chunk["text"]) + + assert texts == [f"token{i} " for i in range(len(frames))] + assert consumed_at_token == list(range(1, len(frames) + 1)) + + @pytest.mark.asyncio async def test_async_native_streaming_forwards_each_frame_incrementally(): """Each token must be emitted after exactly one newly-pulled source frame.