From af4d4ab2ee4f410cb1d54d9dc5b77b21bad27d72 Mon Sep 17 00:00:00 2001 From: harish876 Date: Fri, 10 Apr 2026 06:54:08 +0000 Subject: [PATCH] Introduced Content-Length response headers into the streaming response. This provides a 1:1 behaviour mapping similar to the non streaming behaviour. --- litellm/files/main.py | 41 ++++++++++++++----- litellm/files/streaming.py | 7 +++- litellm/llms/openai/openai.py | 40 ++++++++++++------ .../openai_files_endpoints/files_endpoints.py | 35 +++++++++------- .../test_openai_file_content_streaming.py | 21 +++++++--- .../test_files_endpoint.py | 7 +++- 6 files changed, 105 insertions(+), 46 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 865b679e04f..2c06f9459e4 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -36,7 +36,10 @@ FileContentProvider = Literal[ import litellm from litellm import get_secret_str -from litellm.files.streaming import FileContentStreamingResponse +from litellm.files.streaming import ( + FileContentStreamingResponse, + FileContentStreamingResult, +) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials @@ -990,7 +993,7 @@ async def afile_content_streaming( extra_body: Optional[Dict[str, str]] = None, chunk_size: int = 1024 * 1024, **kwargs, -) -> Union[Iterator[bytes], AsyncIterator[bytes]]: +) -> FileContentStreamingResult: """ Async wrapper for file_content_streaming. """ @@ -1034,7 +1037,7 @@ def file_content_streaming( extra_body: Optional[Dict[str, str]] = None, chunk_size: int = 1024 * 1024, **kwargs, -) -> Union[Iterator[bytes], AsyncIterator[bytes]]: +) -> Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]]: """ Prototype API: Returns a byte iterator for file contents. @@ -1080,7 +1083,23 @@ def file_content_streaming( litellm_params["api_base"] = optional_params.api_base logging_obj.model_call_details["litellm_params"] = litellm_params - response = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) + def _wrap_streaming_result( + response: FileContentStreamingResult, + ) -> FileContentStreamingResult: + return FileContentStreamingResult( + stream_iterator=FileContentStreamingResponse( + stream_iterator=response.stream_iterator, + file_id=file_id, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ), + headers=response.headers, + ) + + response: Union[ + FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult] + ] = FileContentStreamingResult(stream_iterator=iter(()), headers={}) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: openai_creds = get_openai_credentials( api_base=optional_params.api_base, @@ -1115,12 +1134,12 @@ def file_content_streaming( ), ) - return FileContentStreamingResponse( - stream_iterator=response, - file_id=file_id, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging_obj, - ) + if asyncio.iscoroutine(response): + async def _await_and_wrap() -> FileContentStreamingResult: + return _wrap_streaming_result(await response) + + return _await_and_wrap() + + return _wrap_streaming_result(response) except Exception as e: raise e \ No newline at end of file diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index 60b8df2294b..ab149cdb4c3 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -1,6 +1,6 @@ import datetime import traceback -from typing import AsyncIterator, Dict, Iterator, Literal, Optional, Union, cast +from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Optional, Union, cast from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, @@ -13,6 +13,11 @@ FileContentProvider = Literal[ ] +class FileContentStreamingResult(NamedTuple): + stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]] + headers: Dict[str, str] + + class FileContentStreamingResponse: """ Iterator wrapper for file content streaming that carries LiteLLM metadata diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index c42305ac6b2..086ee848521 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -32,6 +32,7 @@ import litellm from litellm import LlmProviders from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.files.streaming import FileContentStreamingResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -1756,12 +1757,21 @@ class OpenAIFilesAPI(BaseLLM): file_content_request: FileContentRequest, openai_client: AsyncOpenAI, chunk_size: int = 1024 * 1024, - ) -> AsyncIterator[bytes]: - async with openai_client.files.with_streaming_response.content( + ) -> FileContentStreamingResult: + response_cm = openai_client.files.with_streaming_response.content( **file_content_request - ) as response: - async for chunk in response.iter_bytes(chunk_size=chunk_size): - yield chunk + ) + response = await response_cm.__aenter__() + headers = dict(response.headers) + + async def _stream() -> AsyncIterator[bytes]: + try: + async for chunk in response.iter_bytes(chunk_size=chunk_size): + yield chunk + finally: + await response_cm.__aexit__(None, None, None) + + return FileContentStreamingResult(stream_iterator=_stream(), headers=headers) def file_content_streaming( self, @@ -1774,7 +1784,7 @@ class OpenAIFilesAPI(BaseLLM): organization: Optional[str], chunk_size: int = 1024 * 1024, client: Optional[Union[OpenAI, AsyncOpenAI]] = None, - ) -> Union[Iterator[bytes], AsyncIterator[bytes]]: + ) -> FileContentStreamingResult: openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1800,13 +1810,19 @@ class OpenAIFilesAPI(BaseLLM): chunk_size=chunk_size, ) - def _stream() -> Iterator[bytes]: - with cast(OpenAI, openai_client).files.with_streaming_response.content( - **file_content_request - ) as response: - yield from response.iter_bytes(chunk_size=chunk_size) + response_cm = cast(OpenAI, openai_client).files.with_streaming_response.content( + **file_content_request + ) + response = response_cm.__enter__() + headers = dict(response.headers) - return _stream() + def _stream() -> Iterator[bytes]: + try: + yield from response.iter_bytes(chunk_size=chunk_size) + finally: + response_cm.__exit__(None, None, None) + + return FileContentStreamingResult(stream_iterator=_stream(), headers=headers) async def aretrieve_file( self, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 05b11721587..acb2ffc3fe9 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -114,25 +114,30 @@ async def _get_streaming_file_content_response( file_id=original_file_id, ) + stream_result = await litellm.afile_content_streaming( + **{ + "custom_llm_provider": custom_llm_provider, + "file_id": file_id, + **data, + } # type: ignore + ) stream_iterator = cast( AsyncIterator[bytes], - await litellm.afile_content_streaming( - **{ - "custom_llm_provider": custom_llm_provider, - "file_id": file_id, - **data, - } # type: ignore - ), + stream_result.stream_iterator, ) hidden_params = getattr(stream_iterator, "_hidden_params", {}) or {} - response_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=hidden_params.get("model_id", "") or "", - cache_key=hidden_params.get("cache_key", "") or "", - api_base=hidden_params.get("api_base", "") or "", - version=version, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - ) + response_headers = { + **stream_result.headers, + **ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=hidden_params.get("model_id", "") or "", + cache_key=hidden_params.get("cache_key", "") or "", + api_base=hidden_params.get("api_base", "") or "", + version=version, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + ), + } + return StreamingResponse( _stream_file_content_with_logging( stream_iterator=stream_iterator, diff --git a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py b/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py index 3362af82a35..fcc0d9a97b4 100644 --- a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py +++ b/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py @@ -2,6 +2,7 @@ import pytest from typing import AsyncIterator, cast from litellm.files import main as files_main +from litellm.files.streaming import FileContentStreamingResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -17,7 +18,10 @@ async def test_afile_content_streaming_routes_to_openai_streaming_handler( def _mock_file_content_streaming(**kwargs): captured_kwargs.update(kwargs) - return _mock_stream() + return FileContentStreamingResult( + stream_iterator=_mock_stream(), + headers={"content-length": "11"}, + ) monkeypatch.setattr( files_main.openai_files_instance, @@ -25,7 +29,7 @@ async def test_afile_content_streaming_routes_to_openai_streaming_handler( _mock_file_content_streaming, ) - stream_iterator = await files_main.afile_content_streaming( + stream_result = await files_main.afile_content_streaming( file_id="file-abc123", custom_llm_provider="openai", api_key="sk-test", @@ -34,10 +38,11 @@ async def test_afile_content_streaming_routes_to_openai_streaming_handler( chunk_size=8, ) - async_stream_iterator = cast(AsyncIterator[bytes], stream_iterator) + async_stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator) chunks = [chunk async for chunk in async_stream_iterator] assert chunks == [b"hello ", b"world"] + assert stream_result.headers["content-length"] == "11" assert captured_kwargs["_is_async"] is True assert captured_kwargs["file_content_request"]["file_id"] == "file-abc123" assert captured_kwargs["api_key"] == "sk-test" @@ -56,7 +61,10 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet yield b"hello" def _mock_file_content_streaming(**kwargs): - return _mock_stream() + return FileContentStreamingResult( + stream_iterator=_mock_stream(), + headers={"content-length": "5"}, + ) async def _mock_async_success_handler( self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs @@ -81,17 +89,18 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet lambda self, result, start_time, end_time, cache_hit=None: None, ) - stream_iterator = await files_main.afile_content_streaming( + stream_result = await files_main.afile_content_streaming( file_id="file-abc123", custom_llm_provider="openai", api_key="sk-test", api_base="https://api.openai.com/v1", ) - async_stream_iterator = cast(AsyncIterator[bytes], stream_iterator) + async_stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator) chunks = [chunk async for chunk in async_stream_iterator] assert chunks == [b"hello"] + assert stream_result.headers["content-length"] == "5" assert captured_standard_logging_object is not None assert captured_standard_logging_object["call_type"] == "afile_content_streaming" assert captured_standard_logging_object["custom_llm_provider"] == "openai" diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 9e279920240..31a6b37437e 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -14,6 +14,7 @@ sys.path.insert( import litellm from litellm import Router +from litellm.files.streaming import FileContentStreamingResult from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users @@ -1575,7 +1576,10 @@ def test_get_file_content_streams_openai_direct_path( yield b"hello " yield b"world" - return _stream() + return FileContentStreamingResult( + stream_iterator=_stream(), + headers={"content-length": "11"}, + ) async def _fail_buffered_path(*args, **kwargs): raise AssertionError("buffered afile_content path should not be used") @@ -1604,6 +1608,7 @@ def test_get_file_content_streams_openai_direct_path( assert response.status_code == 200, response.text assert response.content == b"hello world" assert response.headers["content-type"].startswith("application/octet-stream") + assert response.headers["content-length"] == "11" assert captured_kwargs["custom_llm_provider"] == "openai" assert captured_kwargs["file_id"] == "file-abc123" proxy_logging_obj.update_request_status.assert_awaited_once()