scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Code Comments incorporated.

Browse source 
- Static Methods for Streaming Handler Function

 - Remove the afile_content_streaming wrapper function. Enabled with a stream boolean in afile_content

 - Cleaned up test cases after refactor
This commit is contained in:
harish876 2026-04-10 22:41:13 +00:00
parent 1c74e17bed
commit ccf3dc3161
8 changed files with 435 additions and 261 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

109
litellm/files/file_content_streaming_handler.py Normal file
View file

@ -0,0 +1,109 @@
from typing import Any, AsyncIterator, Dict, Optional, cast
from fastapi.responses import StreamingResponse
import litellm
from litellm.files.types import FileContentStreamingResult
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.openai_files_endpoints.common_utils import (
prepare_data_with_credentials,
)
from litellm.proxy.utils import ProxyLogging
class FileContentStreamingHandler:
@staticmethod
def should_stream_file_content(
*,
custom_llm_provider: str,
is_base64_unified_file_id: Any,
) -> bool:
return (
custom_llm_provider == "openai"
and bool(is_base64_unified_file_id) is False
)
@staticmethod
async def stream_file_content_with_logging(
stream_iterator: AsyncIterator[bytes],
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
data: Dict[str, Any],
):
try:
async for chunk in stream_iterator:
yield chunk
await proxy_logging_obj.update_request_status(
litellm_call_id=data.get("litellm_call_id", ""), status="success"
)
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=data,
)
raise
finally:
if hasattr(stream_iterator, "aclose"):
await stream_iterator.aclose() # type: ignore[attr-defined]
@staticmethod
async def get_streaming_file_content_response(
*,
custom_llm_provider: str,
file_id: str,
data: Dict[str, Any],
should_route: bool,
original_file_id: Optional[str],
credentials: Optional[Dict[str, Any]],
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
version: str,
) -> StreamingResponse:
if should_route:
prepare_data_with_credentials(
data=data,
credentials=credentials, # type: ignore[arg-type]
file_id=original_file_id,
)
stream_result = cast(
FileContentStreamingResult,
await litellm.afile_content(
**{
"custom_llm_provider": custom_llm_provider,
"file_id": file_id,
"stream": True,
**data,
} # type: ignore
),
)
stream_iterator = cast(
AsyncIterator[bytes],
stream_result.stream_iterator,
)
hidden_params = getattr(stream_iterator, "_hidden_params", {}) or {}
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(
FileContentStreamingHandler.stream_file_content_with_logging(
stream_iterator=stream_iterator,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
data=data,
),
media_type="application/octet-stream",
headers=response_headers,
)

255
litellm/files/main.py
View file

@ -771,8 +771,10 @@ async def afile_content(
custom_llm_provider: FileContentProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
chunk_size: int = 1024 * 1024,
stream: bool = False,
**kwargs,
) -> HttpxBinaryResponseContent:
) -> Union[HttpxBinaryResponseContent, FileContentStreamingResult]:
"""
Async: Get file contents
@ -786,11 +788,13 @@ async def afile_content(
# Use a partial function to pass your keyword arguments
func = partial(
file_content,
file_id,
model,
custom_llm_provider,
extra_headers,
extra_body,
file_id=file_id,
model=model,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
chunk_size=chunk_size,
stream=stream,
**kwargs,
)
@ -815,8 +819,15 @@ def file_content(
custom_llm_provider: Optional[Union[FileContentProvider, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
chunk_size: int = 1024 * 1024,
stream: bool = False,
**kwargs,
) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]:
) -> Union[
HttpxBinaryResponseContent,
FileContentStreamingResult,
Coroutine[Any, Any, HttpxBinaryResponseContent],
Coroutine[Any, Any, FileContentStreamingResult],
]:
"""
Returns the contents of the specified file.
@ -858,6 +869,23 @@ def file_content(
_is_async = kwargs.pop("afile_content", False) is True
if stream:
return file_content_streaming(
file_id=file_id,
model=model,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
chunk_size=chunk_size,
optional_params=optional_params,
timeout=timeout,
logging_obj=cast(
Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")
),
_is_async=_is_async,
client=client,
)
# Check if provider has a custom files config (e.g., Anthropic, Manus)
provider_config = ProviderConfigManager.get_provider_files_config(
model="",
@ -983,151 +1011,86 @@ def file_content(
raise e
@client
async def afile_content_streaming(
file_id: str,
custom_llm_provider: FileContentProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
chunk_size: int = 1024 * 1024,
**kwargs,
) -> FileContentStreamingResult:
"""
Async wrapper for file_content_streaming.
"""
try:
loop = asyncio.get_running_loop()
kwargs["afile_content_streaming"] = True
model = kwargs.pop("model", None)
# Use a partial function to pass your keyword arguments
func = partial(
file_content_streaming,
file_id,
model,
custom_llm_provider,
extra_headers,
extra_body,
chunk_size,
**kwargs,
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response # type: ignore
return response
except Exception as e:
raise e
@client
def file_content_streaming(
*,
file_id: str,
model: Optional[str] = None,
custom_llm_provider: Optional[Union[FileContentProvider, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
chunk_size: int = 1024 * 1024,
**kwargs,
model: Optional[str],
custom_llm_provider: Optional[Union[FileContentProvider, str]],
extra_headers: Optional[Dict[str, str]],
extra_body: Optional[Dict[str, str]],
chunk_size: int,
optional_params: GenericLiteLLMParams,
timeout: Union[float, httpx.Timeout],
logging_obj: Optional[LiteLLMLoggingObj],
_is_async: bool,
client: Optional[Any],
) -> Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]]:
"""
Prototype API: Returns a byte iterator for file contents.
if logging_obj is not None:
logging_obj.model = model or ""
logging_obj.model_call_details["model"] = model or ""
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
Supports OpenAI-compatible providers and Azure.
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
logging_obj = cast(
Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")
litellm_params = logging_obj.model_call_details.get("litellm_params", {}) or {}
if optional_params.api_base is not None:
litellm_params["api_base"] = optional_params.api_base
logging_obj.model_call_details["litellm_params"] = litellm_params
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,
)
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
if (
timeout is not None
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(cast(str, custom_llm_provider)) is False
):
timeout = timeout.read or 600
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
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,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
response = openai_files_instance.file_content_streaming(
_is_async=_is_async,
file_content_request=FileContentRequest(
file_id=file_id,
extra_headers=extra_headers,
extra_body=extra_body,
),
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=openai_creds.organization,
chunk_size=chunk_size,
client=client,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format(
custom_llm_provider
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
_is_async = kwargs.pop("afile_content_streaming", False) is True
if asyncio.iscoroutine(response):
async def _await_and_wrap() -> FileContentStreamingResult:
return _wrap_streaming_result(await response)
if logging_obj is not None:
logging_obj.model = model or ""
logging_obj.model_call_details["model"] = model or ""
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
return _await_and_wrap()
litellm_params = logging_obj.model_call_details.get("litellm_params", {}) or {}
if optional_params.api_base is not None:
litellm_params["api_base"] = optional_params.api_base
logging_obj.model_call_details["litellm_params"] = litellm_params
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,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
response = openai_files_instance.file_content_streaming(
_is_async=_is_async,
file_content_request=FileContentRequest(
file_id=file_id,
extra_headers=extra_headers,
extra_body=extra_body,
),
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=openai_creds.organization,
chunk_size=chunk_size,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format(
custom_llm_provider
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
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
return _wrap_streaming_result(response)

1
litellm/files/streaming.py
View file

@ -43,6 +43,7 @@ class FileContentStreamingResponse:
if logging_obj is not None and getattr(logging_obj, "start_time", None)
else datetime.datetime.now()
)
self._sync_hidden_params()
def __iter__(self) -> "FileContentStreamingResponse":
if not hasattr(self.stream_iterator, "__next__"):

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

@ -1765,11 +1765,18 @@ class OpenAIFilesAPI(BaseLLM):
headers = dict(response.headers)
async def _stream() -> AsyncIterator[bytes]:
exc: Optional[BaseException] = None
try:
async for chunk in response.iter_bytes(chunk_size=chunk_size):
yield chunk
except BaseException as e:
exc = e
raise
finally:
await response_cm.__aexit__(None, None, None)
if exc is None:
await response_cm.__aexit__(None, None, None)
else:
await response_cm.__aexit__(type(exc), exc, exc.__traceback__)
return FileContentStreamingResult(stream_iterator=_stream(), headers=headers)
@ -1817,10 +1824,17 @@ class OpenAIFilesAPI(BaseLLM):
headers = dict(response.headers)
def _stream() -> Iterator[bytes]:
exc: Optional[BaseException] = None
try:
yield from response.iter_bytes(chunk_size=chunk_size)
except BaseException as e:
exc = e
raise
finally:
response_cm.__exit__(None, None, None)
if exc is None:
response_cm.__exit__(None, None, None)
else:
response_cm.__exit__(type(exc), exc, exc.__traceback__)
return FileContentStreamingResult(stream_iterator=_stream(), headers=headers)

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

@ -7,7 +7,7 @@
import asyncio
import traceback
from typing import Any, AsyncIterator, Optional, cast, get_args
from typing import Any, Optional, cast, get_args
import httpx
from fastapi import (
@ -21,11 +21,12 @@ from fastapi import (
UploadFile,
status,
)
from fastapi.responses import StreamingResponse
import litellm
from litellm import CreateFileRequest, get_secret_str
from litellm._logging import verbose_proxy_logger
from litellm.files.file_content_streaming_handler import (
FileContentStreamingHandler,
)
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -54,7 +55,6 @@ from .common_utils import (
extract_file_creation_params,
get_credentials_for_model,
handle_model_based_routing,
prepare_data_with_credentials,
)
from .storage_backend_service import StorageBackendFileService
@ -63,96 +63,6 @@ router = APIRouter()
files_config = None
def _should_stream_file_content(
*,
custom_llm_provider: str,
is_base64_unified_file_id: Any,
) -> bool:
return (
custom_llm_provider == "openai"
and bool(is_base64_unified_file_id) is False
)
async def _stream_file_content_with_logging(
stream_iterator: AsyncIterator[bytes],
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
data: Dict[str, Any],
):
try:
async for chunk in stream_iterator:
yield chunk
await proxy_logging_obj.update_request_status(
litellm_call_id=data.get("litellm_call_id", ""), status="success"
)
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=data,
)
raise
finally:
if hasattr(stream_iterator, "aclose"):
await stream_iterator.aclose() # type: ignore[attr-defined]
async def _get_streaming_file_content_response(
*,
custom_llm_provider: str,
file_id: str,
data: Dict[str, Any],
should_route: bool,
original_file_id: Optional[str],
credentials: Optional[Dict[str, Any]],
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
version: str,
) -> StreamingResponse:
if should_route:
prepare_data_with_credentials(
data=data,
credentials=credentials, # type: ignore[arg-type]
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],
stream_result.stream_iterator,
)
hidden_params = getattr(stream_iterator, "_hidden_params", {}) or {}
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,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
data=data,
),
media_type="application/octet-stream",
headers=response_headers,
)
def set_files_config(config):
global files_config
if config is None:
@ -822,14 +732,14 @@ async def get_file_content( # noqa: PLR0915
check_file_id_encoding=True,
)
if _should_stream_file_content(
if FileContentStreamingHandler.should_stream_file_content(
custom_llm_provider=custom_llm_provider,
is_base64_unified_file_id=is_base64_unified_file_id,
):
verbose_proxy_logger.debug(
"Routing file content request to streaming response helper"
)
return await _get_streaming_file_content_response(
return await FileContentStreamingHandler.get_streaming_file_content_response(
custom_llm_provider=custom_llm_provider,
file_id=file_id,
data=data,

2
litellm/utils.py
View file

@ -2131,8 +2131,6 @@ def _is_async_request(
_STREAMING_CALL_TYPES = frozenset(
{
"afile_content_streaming",
"file_content_streaming",
CallTypes.generate_content_stream,
CallTypes.agenerate_content_stream,
CallTypes.generate_content_stream.value,

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

@ -1,14 +1,15 @@
import pytest
from typing import AsyncIterator, cast
from typing import AsyncIterator, Iterator, cast
from litellm.files import main as files_main
from litellm.files.streaming import FileContentStreamingResponse
from litellm.files.types import FileContentStreamingResult
from litellm.llms.openai.openai import OpenAIFilesAPI
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@pytest.mark.asyncio
async def test_afile_content_streaming_routes_to_openai_streaming_handler(
async def test_afile_content_with_stream_routes_to_openai_streaming_handler(
monkeypatch,
):
captured_kwargs = {}
@ -30,13 +31,17 @@ async def test_afile_content_streaming_routes_to_openai_streaming_handler(
_mock_file_content_streaming,
)
stream_result = await files_main.afile_content_streaming(
stream_result = cast(
FileContentStreamingResult,
await files_main.afile_content(
file_id="file-abc123",
custom_llm_provider="openai",
api_key="sk-test",
api_base="https://api.openai.com/v1",
organization="org-123",
chunk_size=8,
stream=True,
),
)
async_stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator)
@ -50,6 +55,7 @@ async def test_afile_content_streaming_routes_to_openai_streaming_handler(
assert captured_kwargs["api_base"] == "https://api.openai.com/v1"
assert captured_kwargs["organization"] == "org-123"
assert captured_kwargs["chunk_size"] == 8
assert captured_kwargs["client"] is None
@pytest.mark.asyncio
@ -90,11 +96,15 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet
lambda self, result, start_time, end_time, cache_hit=None: None,
)
stream_result = await files_main.afile_content_streaming(
stream_result = cast(
FileContentStreamingResult,
await files_main.afile_content(
file_id="file-abc123",
custom_llm_provider="openai",
api_key="sk-test",
api_base="https://api.openai.com/v1",
stream=True,
),
)
async_stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator)
@ -103,7 +113,7 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet
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["call_type"] == "afile_content"
assert captured_standard_logging_object["custom_llm_provider"] == "openai"
assert captured_standard_logging_object["response"]["id"] == "file-abc123"
assert (
@ -112,6 +122,35 @@ async def test_afile_content_streaming_builds_standard_logging_object_on_complet
)
@pytest.mark.asyncio
async def test_afile_content_streaming_shim_sets_stream_flag(
monkeypatch,
):
captured_kwargs = {}
def _mock_file_content_streaming(**kwargs):
captured_kwargs.update(kwargs)
return FileContentStreamingResult(
stream_iterator=iter(()),
headers={},
)
monkeypatch.setattr(
files_main.openai_files_instance,
"file_content_streaming",
_mock_file_content_streaming,
)
await files_main.afile_content(
file_id="file-abc123",
custom_llm_provider="openai",
api_key="sk-test",
stream=True,
)
assert captured_kwargs["_is_async"] is True
@pytest.mark.asyncio
async def test_file_content_streaming_response_aclose_closes_underlying_async_generator():
close_called = False
@ -137,3 +176,148 @@ async def test_file_content_streaming_response_aclose_closes_underlying_async_ge
await stream.aclose()
assert close_called is True
@pytest.mark.asyncio
async def test_afile_content_streaming_populates_hidden_params_before_iteration(
monkeypatch,
):
async def _mock_stream():
yield b"hello"
def _mock_file_content_streaming(**kwargs):
return FileContentStreamingResult(
stream_iterator=_mock_stream(),
headers={"content-length": "5"},
)
monkeypatch.setattr(
files_main.openai_files_instance,
"file_content_streaming",
_mock_file_content_streaming,
)
stream_result = cast(
FileContentStreamingResult,
await files_main.afile_content(
file_id="file-abc123",
custom_llm_provider="openai",
api_key="sk-test",
api_base="https://api.openai.com/v1",
stream=True,
),
)
stream_iterator = cast(FileContentStreamingResponse, stream_result.stream_iterator)
assert stream_iterator._hidden_params["api_base"] == "https://api.openai.com/v1"
assert stream_iterator._hidden_params["litellm_model_name"] is None
@pytest.mark.asyncio
async def test_afile_content_streaming_passes_exception_to_context_manager_exit():
class MockAsyncResponse:
headers = {"content-length": "1"}
async def iter_bytes(self, chunk_size: int):
yield b"a"
raise RuntimeError("stream failed")
class MockAsyncResponseContextManager:
def __init__(self):
self.exc_info = None
async def __aenter__(self):
return MockAsyncResponse()
async def __aexit__(self, exc_type, exc, tb):
self.exc_info = (exc_type, exc, tb)
class MockAsyncFiles:
def __init__(self, response_cm):
self.with_streaming_response = self
self._response_cm = response_cm
def content(self, **kwargs):
return self._response_cm
class MockAsyncOpenAIClient:
def __init__(self, response_cm):
self.files = MockAsyncFiles(response_cm)
response_cm = MockAsyncResponseContextManager()
api = OpenAIFilesAPI()
stream_result = await api.afile_content_streaming(
file_content_request={"file_id": "file-abc123"},
openai_client=MockAsyncOpenAIClient(response_cm), # type: ignore[arg-type]
chunk_size=1,
)
stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator)
assert await stream_iterator.__anext__() == b"a"
with pytest.raises(RuntimeError, match="stream failed") as exc_info:
await stream_iterator.__anext__()
assert response_cm.exc_info is not None
assert response_cm.exc_info[0] is RuntimeError
assert response_cm.exc_info[1] is exc_info.value
assert response_cm.exc_info[2] is not None
def test_file_content_streaming_passes_exception_to_context_manager_exit():
class MockSyncResponse:
headers = {"content-length": "1"}
def iter_bytes(self, chunk_size: int) -> Iterator[bytes]:
yield b"a"
raise RuntimeError("stream failed")
class MockSyncResponseContextManager:
def __init__(self):
self.exc_info = None
def __enter__(self):
return MockSyncResponse()
def __exit__(self, exc_type, exc, tb):
self.exc_info = (exc_type, exc, tb)
class MockSyncFiles:
def __init__(self, response_cm):
self.with_streaming_response = self
self._response_cm = response_cm
def content(self, **kwargs):
return self._response_cm
class MockSyncOpenAIClient:
def __init__(self, response_cm):
self.files = MockSyncFiles(response_cm)
response_cm = MockSyncResponseContextManager()
api = OpenAIFilesAPI()
stream_result = api.file_content_streaming(
_is_async=False,
file_content_request={"file_id": "file-abc123"},
api_base="https://api.openai.com/v1",
api_key="sk-test",
timeout=60,
max_retries=None,
organization=None,
chunk_size=1,
client=MockSyncOpenAIClient(response_cm), # type: ignore[arg-type]
)
stream_iterator = cast(Iterator[bytes], stream_result.stream_iterator)
assert next(stream_iterator) == b"a"
with pytest.raises(RuntimeError, match="stream failed") as exc_info:
next(stream_iterator)
assert response_cm.exc_info is not None
assert response_cm.exc_info[0] is RuntimeError
assert response_cm.exc_info[1] is exc_info.value
assert response_cm.exc_info[2] is not None

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

@ -14,10 +14,8 @@ sys.path.insert(
import litellm
from litellm import Router
from litellm.files.file_content_streaming_handler import FileContentStreamingHandler
from litellm.files.types import FileContentStreamingResult
from litellm.proxy.openai_files_endpoints.files_endpoints import (
_stream_file_content_with_logging,
)
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
@ -100,7 +98,7 @@ async def test_stream_file_content_with_logging_closes_inner_iterator_on_early_e
stream_iterator = MockStreamIterator()
proxy_logging_obj = AsyncMock()
generator = _stream_file_content_with_logging(
generator = FileContentStreamingHandler.stream_file_content_with_logging(
stream_iterator=stream_iterator,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=AsyncMock(),
@ -1606,7 +1604,7 @@ def test_get_file_content_streams_openai_direct_path(
captured_kwargs = {}
async def _mock_afile_content_streaming(**kwargs):
async def _mock_afile_content(**kwargs):
captured_kwargs.update(kwargs)
async def _stream():
@ -1618,11 +1616,7 @@ def test_get_file_content_streams_openai_direct_path(
headers={"content-length": "11"},
)
async def _fail_buffered_path(*args, **kwargs):
raise AssertionError("buffered afile_content path should not be used")
monkeypatch.setattr(litellm, "afile_content_streaming", _mock_afile_content_streaming)
monkeypatch.setattr(litellm, "afile_content", _fail_buffered_path)
monkeypatch.setattr(litellm, "afile_content", _mock_afile_content)
monkeypatch.setattr(
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
lambda **kwargs: (False, None, None, None),
@ -1648,5 +1642,6 @@ def test_get_file_content_streams_openai_direct_path(
assert response.headers["content-length"] == "11"
assert captured_kwargs["custom_llm_provider"] == "openai"
assert captured_kwargs["file_id"] == "file-abc123"
assert captured_kwargs["stream"] is True
proxy_logging_obj.update_request_status.assert_awaited_once()
proxy_logging_obj.post_call_failure_hook.assert_not_called()