scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Refactor file content streaming handling to improve routing and support

Browse source 
- Introduced a new method in `FileContentStreamingHandler` to resolve streaming request parameters, enhancing the routing logic based on credentials.
- Updated the `should_stream_file_content` method to check against supported providers.
- Cleaned up type hints and imports across multiple files for better organization and clarity.
- Added comprehensive tests to validate the new routing behavior and ensure original data integrity during streaming requests.
This commit is contained in:
harish876 2026-04-11 18:56:15 +00:00
parent f523ccb2fe
commit 69eb34597c
6 changed files with 206 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/files/main.py
View file

@ -30,14 +30,10 @@ FileRetrieveProvider = Literal[
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "manus", "anthropic"]
FileContentProvider = Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
]
import litellm
from litellm import get_secret_str
from litellm.files.streaming import FileContentStreamingResponse
from litellm.files.types import FileContentStreamingResult
from litellm.files.types import FileContentProvider, 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
@ -68,6 +64,15 @@ from litellm.utils import (
base_llm_http_handler = BaseLLMHTTPHandler()
####### ENVIRONMENT VARIABLES ###################
def _should_sdk_support_streaming(
custom_llm_provider: Optional[Union[FileContentProvider, str]],
) -> bool:
"""
Return whether file content streaming is supported for the provider.
"""
return custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS
openai_files_instance = OpenAIFilesAPI()
azure_files_instance = AzureOpenAIFilesAPI()
vertex_ai_files_instance = VertexAIFilesHandler()
@ -869,7 +874,7 @@ def file_content(
_is_async = kwargs.pop("afile_content", False) is True
if stream:
if stream and _should_sdk_support_streaming(custom_llm_provider):
return file_content_streaming(
file_id=file_id,
model=model,
@ -1075,8 +1080,9 @@ def file_content_streaming(
)
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
message="LiteLLM doesn't support {} for streaming 'file_content'. Supported providers are {}.".format(
custom_llm_provider,
sorted(OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS),
),
model="n/a",
llm_provider=custom_llm_provider,

8
litellm/files/streaming.py
View file

@ -1,8 +1,9 @@
import datetime
import traceback
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional, Union, cast
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Optional, Union, cast
import anyio
from litellm.files.types import FileContentProvider
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
@ -10,11 +11,6 @@ if TYPE_CHECKING:
)
from litellm.types.utils import StandardLoggingHiddenParams, StandardLoggingPayload
FileContentProvider = Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
]
class FileContentStreamingResponse:
"""
Iterator wrapper for file content streaming that carries LiteLLM metadata

7
litellm/files/types.py
View file

@ -1,4 +1,9 @@
from typing import AsyncIterator, Dict, Iterator, NamedTuple, Union
from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Union
FileContentProvider = Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
]
class FileContentStreamingResult(NamedTuple):

81
litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py
View file

@ -1,9 +1,10 @@
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Optional, cast
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Optional, Tuple, cast
from fastapi.responses import StreamingResponse
import litellm
from litellm.files.types import FileContentStreamingResult
from litellm.files.types import FileContentProvider, FileContentStreamingResult
from litellm.types.utils import OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
@ -11,15 +12,61 @@ if TYPE_CHECKING:
class FileContentStreamingHandler:
@staticmethod
def resolve_streaming_request_params(
*,
custom_llm_provider: str,
file_id: str,
data: Dict[str, Any],
should_route: bool,
original_file_id: Optional[str],
credentials: Optional[Dict[str, Any]],
) -> Tuple[str, str, Dict[str, Any]]:
"""
Resolve the provider, file ID, and request payload to use for streaming.
For model-routed requests, this derives the effective provider from
credentials, applies `prepare_data_with_credentials()` to a copied
payload, swaps in the decoded/original file ID, and removes `model`
so `afile_content()` does not re-resolve the provider. This helper
does not mutate the passed-in `data` dictionary. Non-routed requests
return the original provider, file ID, and data unchanged.
"""
if should_route and credentials is not None:
from litellm.proxy.openai_files_endpoints.common_utils import (
prepare_data_with_credentials,
)
resolved_streaming_data = dict(data)
prepare_data_with_credentials(
data=resolved_streaming_data,
credentials=credentials,
file_id=original_file_id,
)
resolved_streaming_data.pop("model", None)
resolved_streaming_provider = cast(
str, credentials["custom_llm_provider"]
)
resolved_custom_llm_provider = resolved_streaming_provider
resolved_file_id = cast(str, resolved_streaming_data["file_id"])
else:
resolved_streaming_data = data
resolved_custom_llm_provider = custom_llm_provider
resolved_file_id = file_id
return (
resolved_custom_llm_provider,
resolved_file_id,
resolved_streaming_data,
)
@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
custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS
)
@staticmethod
@ -52,9 +99,6 @@ class FileContentStreamingHandler:
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,
@ -62,28 +106,13 @@ class FileContentStreamingHandler:
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
prepare_data_with_credentials,
)
effective_custom_llm_provider = custom_llm_provider
if should_route:
if credentials is None or credentials.get("custom_llm_provider") is None:
raise ValueError(
"Model-based file routing requires credentials with custom_llm_provider"
)
prepare_data_with_credentials(
data=data,
credentials=credentials,
file_id=original_file_id,
)
effective_custom_llm_provider = cast(str, credentials["custom_llm_provider"])
stream_result = cast(
FileContentStreamingResult,
await litellm.afile_content(
**{
"custom_llm_provider": effective_custom_llm_provider,
"custom_llm_provider": cast(
FileContentProvider, custom_llm_provider
),
"file_id": file_id,
"stream": True,
**data,

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

@ -634,7 +634,7 @@ async def get_file_content( # noqa: PLR0915
or await get_custom_llm_provider_from_request_body(request=request)
or "openai"
)
## check if file_id is a litellm managed file
is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
if is_base64_unified_file_id:
@ -735,21 +735,33 @@ async def get_file_content( # noqa: PLR0915
from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import (
FileContentStreamingHandler,
)
(
resolved_custom_llm_provider,
resolved_file_id,
resolved_streaming_data,
) = FileContentStreamingHandler.resolve_streaming_request_params(
custom_llm_provider=custom_llm_provider,
file_id=file_id,
data=data,
should_route=should_route,
original_file_id=original_file_id,
credentials=credentials,
)
if FileContentStreamingHandler.should_stream_file_content(
custom_llm_provider=custom_llm_provider,
is_base64_unified_file_id=is_base64_unified_file_id,
custom_llm_provider=resolved_custom_llm_provider,
):
verbose_proxy_logger.debug(
"Routing file content request to streaming response helper"
"Using streaming file content helper for custom_llm_provider=%s, original_file_id=%s, file_id=%s, model_used=%s",
resolved_custom_llm_provider,
original_file_id,
resolved_file_id,
model_used,
)
return await FileContentStreamingHandler.get_streaming_file_content_response(
custom_llm_provider=custom_llm_provider,
file_id=file_id,
data=data,
should_route=should_route,
original_file_id=original_file_id,
credentials=credentials,
custom_llm_provider=resolved_custom_llm_provider,
file_id=resolved_file_id,
data=resolved_streaming_data,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=version,
@ -762,7 +774,6 @@ async def get_file_content( # noqa: PLR0915
credentials=credentials, # type: ignore
file_id=original_file_id, # Use decoded file ID if from encoded ID
)
response = await litellm.afile_content(
custom_llm_provider=credentials["custom_llm_provider"], # type: ignore
**data,

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

@ -116,6 +116,93 @@ async def test_stream_file_content_with_logging_closes_inner_iterator_on_early_e
proxy_logging_obj.update_request_status.assert_not_called()
def test_resolve_streaming_request_params_non_routed_returns_original_values():
data = {"file_id": "file-abc123", "metadata": {"k": "v"}}
(
resolved_custom_llm_provider,
resolved_file_id,
resolved_streaming_data,
) = FileContentStreamingHandler.resolve_streaming_request_params(
custom_llm_provider="openai",
file_id="file-abc123",
data=data,
should_route=False,
original_file_id=None,
credentials=None,
)
assert resolved_custom_llm_provider == "openai"
assert resolved_file_id == "file-abc123"
assert resolved_streaming_data is data
def test_resolve_streaming_request_params_routed_uses_credentials_and_original_file_id():
data = {
"file_id": "file-encoded-123",
"model": "azure-gpt-3-5-turbo",
"metadata": {"k": "v"},
}
credentials = {
"custom_llm_provider": "azure",
"api_key": "azure-key",
"api_base": "https://azure.example.com",
}
(
resolved_custom_llm_provider,
resolved_file_id,
resolved_streaming_data,
) = FileContentStreamingHandler.resolve_streaming_request_params(
custom_llm_provider="openai",
file_id="file-encoded-123",
data=data,
should_route=True,
original_file_id="file-original-123",
credentials=credentials,
)
assert resolved_custom_llm_provider == "azure"
assert resolved_file_id == "file-original-123"
assert resolved_streaming_data["file_id"] == "file-original-123"
assert resolved_streaming_data["api_key"] == "azure-key"
assert resolved_streaming_data["api_base"] == "https://azure.example.com"
assert "custom_llm_provider" not in resolved_streaming_data
assert "model" not in resolved_streaming_data
assert data["file_id"] == "file-encoded-123"
assert data["model"] == "azure-gpt-3-5-turbo"
def test_resolve_streaming_request_params_routed_preserves_input_data_object():
data = {
"file_id": "file-encoded-123",
"model": "openai/gpt-4o",
}
credentials = {
"custom_llm_provider": "openai",
"api_key": "sk-test",
}
(
_resolved_custom_llm_provider,
_resolved_file_id,
resolved_streaming_data,
) = FileContentStreamingHandler.resolve_streaming_request_params(
custom_llm_provider="openai",
file_id="file-encoded-123",
data=data,
should_route=True,
original_file_id=None,
credentials=credentials,
)
assert resolved_streaming_data is not data
assert data == {
"file_id": "file-encoded-123",
"model": "openai/gpt-4o",
}
def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router):
"""
Asserts 'create_file' is called with the correct arguments
@ -1650,7 +1737,7 @@ def test_get_file_content_streams_openai_direct_path(
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def test_get_file_content_streams_with_routed_provider(
def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider_is_not_supported(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
import litellm.proxy.proxy_server as ps
@ -1666,17 +1753,25 @@ def test_get_file_content_streams_with_routed_provider(
async def _mock_afile_content(**kwargs):
captured_kwargs.update(kwargs)
async def _stream():
yield b"hello "
yield b"world"
return FileContentStreamingResult(
stream_iterator=_stream(),
headers={"content-length": "11"},
return HttpxBinaryResponseContent(
response=httpx.Response(
status_code=200,
content=b"azure-bytes",
headers={
"content-type": "application/octet-stream",
"content-length": "11",
},
)
)
mock_streaming_response = mocker.AsyncMock()
monkeypatch.setattr(litellm, "afile_content", _mock_afile_content)
monkeypatch.setattr(
FileContentStreamingHandler,
"get_streaming_file_content_response",
mock_streaming_response,
)
monkeypatch.setattr(
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
lambda **kwargs: (
@ -1706,13 +1801,13 @@ def test_get_file_content_streams_with_routed_provider(
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
assert response.content == b"hello world"
assert response.content == b"azure-bytes"
assert captured_kwargs["custom_llm_provider"] == "azure"
assert captured_kwargs["file_id"] == "file-original-123"
assert captured_kwargs["api_key"] == "azure-key"
assert captured_kwargs["api_base"] == "https://azure.example.com"
assert captured_kwargs["stream"] is True
proxy_logging_obj.update_request_status.assert_awaited_once()
assert "stream" not in captured_kwargs
mock_streaming_response.assert_not_awaited()
proxy_logging_obj.post_call_failure_hook.assert_not_called()