Add file content streaming support for OpenAI and related utilities

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,
+        )

litellm/files/main.py

@@ -10,7 +10,7 @@
 import time
 import uuid as uuid_module
 from functools import partial
-from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
+from typing import Any,Coroutine, Dict, Literal, Optional, Union, cast
 
 import httpx
 
@@ -36,6 +36,8 @@
 
 import litellm
 from litellm import get_secret_str
+from litellm.files.streaming import FileContentStreamingResponse
+from litellm.files.types import 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
@@ -55,10 +57,7 @@
     OpenAIFileObject,
 )
 from litellm.types.router import *
-from litellm.types.utils import (
-    OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS,
-    LlmProviders,
-)
+from litellm.types.utils import OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, LlmProviders
 from litellm.utils import (
     ProviderConfigManager,
     client,
@@ -772,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
 
@@ -787,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,
         )
 
@@ -816,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.
 
@@ -859,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="",
@@ -982,3 +1009,88 @@ def file_content(
         return response
     except Exception as e:
         raise e
+
+
+def file_content_streaming(
+    *,
+    file_id: str,
+    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]]:
+    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
+
+        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,
+            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
+            ),
+        )
+
+    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)