diff --git a/litellm/files/file_content_streaming_handler.py b/litellm/files/file_content_streaming_handler.py new file mode 100644 index 00000000000..dd39e1b0312 --- /dev/null +++ b/litellm/files/file_content_streaming_handler.py @@ -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, + ) diff --git a/litellm/files/main.py b/litellm/files/main.py index 13abdfd92fd..9bcb3976f0f 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -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 \ No newline at end of file + return _wrap_streaming_result(response) \ No newline at end of file diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index dab190b9b32..f49b7054881 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -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__"): diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 76f300ec9c7..b48edf53d5e 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -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) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 468e0baaa4f..eefcd12ddbb 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -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, diff --git a/litellm/utils.py b/litellm/utils.py index d8e14de6372..939bbd91ce9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, 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 66c69640216..ac5fb91f6f3 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 @@ -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 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 37363801240..47bcaa2dfc5 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,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()