Repository navigation
Add file content streaming support for OpenAI and related utilities #25450
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 5 commits
7ebc144
1310803
af4d4ab
044d434
baba3eb
1c74e17
ccf3dc3
4e5e739
d1dda3d
c5d93e6
68bc6de
f523ccb
69eb345
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
|
|
@@ -982,3 +981,153 @@ def file_content( | |
| return response | ||
| except Exception as e: | ||
| 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( | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. this feels like a lot of duplicate code. Why can't we just add a stream=True/False on def file_content ? That way you don't need this new function
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Counterpoint here. I think keeping file_content_streaming() separate is the cleaner choice because this is not just a stream=True transport toggle on the existing API. The streaming path returns a different shape, carries headers alongside an iterator, and has iterator-specific logging and cleanup behavior like aclose() on disconnect. Keeping it separate preserves the existing file_content() contract, makes the rollout to other providers incremental, and keeps the streaming-specific behavior isolated and easier to test. The original function code can be removed once we migrate all paths to a streaming one. |
||
| 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, | ||
| ) -> Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]]: | ||
| """ | ||
| Prototype API: Returns a byte iterator for file contents. | ||
|
|
||
| Supports OpenAI-compatible providers and Azure. | ||
| """ | ||
| try: | ||
| optional_params = GenericLiteLLMParams(**kwargs) | ||
| logging_obj = cast( | ||
| Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj") | ||
| ) | ||
|
|
||
| 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 | ||
|
|
||
| _is_async = kwargs.pop("afile_content_streaming", False) is True | ||
|
|
||
| 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, | ||
| ) | ||
| 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 | ||
Uh oh!
There was an error while loading. Please reload this page.