scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Introduced Content-Length response headers into the streaming response. This provides a 1:1 behaviour mapping similar to the non streaming behaviour.

Browse source
This commit is contained in:
harish876 2026-04-10 06:54:08 +00:00
parent 13108039c8
commit af4d4ab2ee
6 changed files with 105 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/files/main.py
View file

@ -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

7
litellm/files/streaming.py
View file

@ -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

40
litellm/llms/openai/openai.py
View file

@ -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,

35
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -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,

21
tests/test_litellm/llms/openai/test_openai_file_content_streaming.py
View file

@ -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"

7
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -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()