mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 10:24:03 +00:00
Add file content streaming support for OpenAI and related utilities
- Introduced `afile_content_streaming` and `file_content_streaming` functions in `litellm/files/main.py` to handle asynchronous and synchronous file content streaming. - Added `FileContentStreamingResponse` class in `litellm/files/streaming.py` to manage streaming responses with logging capabilities. - Updated OpenAI API integration in `litellm/llms/openai/openai.py` to support new streaming methods. - Enhanced file content retrieval in `litellm/proxy/openai_files_endpoints/files_endpoints.py` to route requests for streaming. - Added unit tests for the new streaming functionality in `tests/test_litellm/llms/openai/test_openai_file_content_streaming.py` and `tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py`. - Refactored type hints and imports for better clarity and organization across modified files.
This commit is contained in:
harish876
+147
-5
@@ -10,7 +10,7 @@ import contextvars
|
||||
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, AsyncIterator, Coroutine, Dict, Iterator, Literal, Optional, Union, cast
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -36,6 +36,7 @@ FileContentProvider = Literal[
|
||||
|
||||
import litellm
|
||||
from litellm import get_secret_str
|
||||
from litellm.files.streaming import FileContentStreamingResponse
|
||||
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 +56,7 @@ from litellm.types.llms.openai import (
|
||||
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 +980,147 @@ 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,
|
||||
) -> Union[Iterator[bytes], AsyncIterator[bytes]]:
|
||||
"""
|
||||
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,
|
||||
) -> Union[Iterator[bytes], AsyncIterator[bytes]]:
|
||||
"""
|
||||
Prototype API: Returns a byte iterator for file contents.
|
||||
|
||||
Supports OpenAI-compatible providers and Azure.
|
||||
"""
|
||||
try:
|
||||
optional_params = GenericLiteLLMParams(**kwargs)
|
||||
litellm_params_dict = get_litellm_params(**kwargs)
|
||||
client = kwargs.get("client")
|
||||
logging_obj = cast(
|
||||
Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")
|
||||
)
|
||||
|
||||
try:
|
||||
if model is not None:
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model, custom_llm_provider
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
response = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(()))
|
||||
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
|
||||
),
|
||||
)
|
||||
|
||||
return FileContentStreamingResponse(
|
||||
stream_iterator=response,
|
||||
file_id=file_id,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
@@ -0,0 +1,205 @@
|
||||
import datetime
|
||||
import traceback
|
||||
from typing import AsyncIterator, Dict, Iterator, Literal, Optional, Union, cast
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
Logging as LiteLLMLoggingObj,
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
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
|
||||
and emits success/failure callbacks once the stream finishes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]],
|
||||
file_id: str,
|
||||
model: Optional[str],
|
||||
custom_llm_provider: Optional[Union[FileContentProvider, str]],
|
||||
logging_obj: Optional[LiteLLMLoggingObj],
|
||||
) -> None:
|
||||
self.stream_iterator = stream_iterator
|
||||
self.file_id = file_id
|
||||
self.model = model
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
self.logging_obj = logging_obj
|
||||
self.standard_logging_object: Optional[StandardLoggingPayload] = None
|
||||
self._hidden_params: StandardLoggingHiddenParams = cast(
|
||||
StandardLoggingHiddenParams, {}
|
||||
)
|
||||
self._logging_completed = False
|
||||
self._start_time = (
|
||||
logging_obj.start_time
|
||||
if logging_obj is not None and getattr(logging_obj, "start_time", None)
|
||||
else datetime.datetime.now()
|
||||
)
|
||||
|
||||
def __iter__(self) -> "FileContentStreamingResponse":
|
||||
if not hasattr(self.stream_iterator, "__next__"):
|
||||
raise TypeError("File content stream does not support sync iteration")
|
||||
return self
|
||||
|
||||
def __next__(self) -> bytes:
|
||||
if not hasattr(self.stream_iterator, "__next__"):
|
||||
raise TypeError("File content stream does not support sync iteration")
|
||||
|
||||
try:
|
||||
return next(cast(Iterator[bytes], self.stream_iterator))
|
||||
except StopIteration:
|
||||
self._log_success_sync()
|
||||
raise
|
||||
except Exception as e:
|
||||
self._log_failure_sync(e)
|
||||
raise
|
||||
|
||||
def __aiter__(self) -> "FileContentStreamingResponse":
|
||||
if not hasattr(self.stream_iterator, "__anext__"):
|
||||
raise TypeError("File content stream does not support async iteration")
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
if not hasattr(self.stream_iterator, "__anext__"):
|
||||
raise TypeError("File content stream does not support async iteration")
|
||||
|
||||
try:
|
||||
return await cast(AsyncIterator[bytes], self.stream_iterator).__anext__()
|
||||
except StopAsyncIteration:
|
||||
await self._log_success_async()
|
||||
raise
|
||||
except Exception as e:
|
||||
await self._log_failure_async(e)
|
||||
raise
|
||||
|
||||
def _build_logging_response(self) -> Dict[str, str]:
|
||||
response = {
|
||||
"id": self.file_id,
|
||||
"object": "file.content",
|
||||
}
|
||||
if self.model:
|
||||
response["model"] = self.model
|
||||
return response
|
||||
|
||||
def _sync_hidden_params(self) -> None:
|
||||
litellm_params = {}
|
||||
if self.logging_obj is not None:
|
||||
litellm_params = (
|
||||
self.logging_obj.model_call_details.get("litellm_params", {}) or {}
|
||||
)
|
||||
|
||||
if "api_base" not in self._hidden_params and litellm_params.get("api_base"):
|
||||
self._hidden_params["api_base"] = litellm_params["api_base"]
|
||||
|
||||
# The generic client decorator infers `model` from the first positional arg,
|
||||
# which is `file_id` for this API. Correct it before logging callbacks run.
|
||||
self._hidden_params["litellm_model_name"] = self.model
|
||||
if "response_cost" not in self._hidden_params:
|
||||
self._hidden_params["response_cost"] = None
|
||||
|
||||
def _build_standard_logging_object(
|
||||
self,
|
||||
end_time: datetime.datetime,
|
||||
) -> Optional[StandardLoggingPayload]:
|
||||
if self.standard_logging_object is not None:
|
||||
return self.standard_logging_object
|
||||
|
||||
if self.logging_obj is None:
|
||||
return None
|
||||
|
||||
self._sync_hidden_params()
|
||||
payload = get_standard_logging_object_payload(
|
||||
kwargs=self.logging_obj.model_call_details,
|
||||
init_response_obj=self._build_logging_response(),
|
||||
start_time=self._start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self.logging_obj,
|
||||
status="success",
|
||||
)
|
||||
if payload is None:
|
||||
return None
|
||||
|
||||
merged_hidden_params = cast(
|
||||
StandardLoggingHiddenParams,
|
||||
{
|
||||
**cast(
|
||||
StandardLoggingHiddenParams, payload.get("hidden_params") or {}
|
||||
),
|
||||
**self._hidden_params,
|
||||
},
|
||||
)
|
||||
payload["hidden_params"] = merged_hidden_params
|
||||
payload["response"] = self._build_logging_response()
|
||||
if self.custom_llm_provider is not None:
|
||||
payload["custom_llm_provider"] = self.custom_llm_provider
|
||||
if self.model is not None:
|
||||
payload["model"] = self.model
|
||||
if self._hidden_params.get("api_base"):
|
||||
payload["api_base"] = cast(str, self._hidden_params["api_base"])
|
||||
|
||||
self.standard_logging_object = payload
|
||||
return payload
|
||||
|
||||
async def _log_success_async(self) -> None:
|
||||
if self._logging_completed or self.logging_obj is None:
|
||||
return
|
||||
|
||||
self._logging_completed = True
|
||||
end_time = datetime.datetime.now()
|
||||
standard_logging_object = self._build_standard_logging_object(end_time=end_time)
|
||||
await self.logging_obj.async_success_handler(
|
||||
result=self._build_logging_response(),
|
||||
start_time=self._start_time,
|
||||
end_time=end_time,
|
||||
standard_logging_object=standard_logging_object,
|
||||
)
|
||||
self.logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=self._build_logging_response(),
|
||||
start_time=self._start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
def _log_success_sync(self) -> None:
|
||||
if self._logging_completed or self.logging_obj is None:
|
||||
return
|
||||
|
||||
self._logging_completed = True
|
||||
end_time = datetime.datetime.now()
|
||||
standard_logging_object = self._build_standard_logging_object(end_time=end_time)
|
||||
self.logging_obj.success_handler(
|
||||
result=self._build_logging_response(),
|
||||
start_time=self._start_time,
|
||||
end_time=end_time,
|
||||
standard_logging_object=standard_logging_object,
|
||||
)
|
||||
|
||||
async def _log_failure_async(self, error: Exception) -> None:
|
||||
if self._logging_completed or self.logging_obj is None:
|
||||
return
|
||||
|
||||
self._logging_completed = True
|
||||
end_time = datetime.datetime.now()
|
||||
traceback_str = traceback.format_exc()
|
||||
self.logging_obj.failure_handler(
|
||||
error, traceback_str, self._start_time, end_time
|
||||
)
|
||||
await self.logging_obj.async_failure_handler(
|
||||
error, traceback_str, self._start_time, end_time
|
||||
)
|
||||
|
||||
def _log_failure_sync(self, error: Exception) -> None:
|
||||
if self._logging_completed or self.logging_obj is None:
|
||||
return
|
||||
|
||||
self._logging_completed = True
|
||||
end_time = datetime.datetime.now()
|
||||
self.logging_obj.failure_handler(
|
||||
error, traceback.format_exc(), self._start_time, end_time
|
||||
)
|
||||
@@ -1751,6 +1751,63 @@ class OpenAIFilesAPI(BaseLLM):
|
||||
|
||||
return HttpxBinaryResponseContent(response=response.response)
|
||||
|
||||
async def afile_content_streaming(
|
||||
self,
|
||||
file_content_request: FileContentRequest,
|
||||
openai_client: AsyncOpenAI,
|
||||
chunk_size: int = 1024 * 1024,
|
||||
) -> AsyncIterator[bytes]:
|
||||
async with openai_client.files.with_streaming_response.content(
|
||||
**file_content_request
|
||||
) as response:
|
||||
async for chunk in response.iter_bytes(chunk_size=chunk_size):
|
||||
yield chunk
|
||||
|
||||
def file_content_streaming(
|
||||
self,
|
||||
_is_async: bool,
|
||||
file_content_request: FileContentRequest,
|
||||
api_base: str,
|
||||
api_key: Optional[str],
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
max_retries: Optional[int],
|
||||
organization: Optional[str],
|
||||
chunk_size: int = 1024 * 1024,
|
||||
client: Optional[Union[OpenAI, AsyncOpenAI]] = None,
|
||||
) -> Union[Iterator[bytes], AsyncIterator[bytes]]:
|
||||
openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client(
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
organization=organization,
|
||||
client=client,
|
||||
_is_async=_is_async,
|
||||
)
|
||||
if openai_client is None:
|
||||
raise ValueError(
|
||||
"OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
|
||||
)
|
||||
|
||||
if _is_async is True:
|
||||
if not isinstance(openai_client, AsyncOpenAI):
|
||||
raise ValueError(
|
||||
"OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client."
|
||||
)
|
||||
return self.afile_content_streaming( # type: ignore
|
||||
file_content_request=file_content_request,
|
||||
openai_client=openai_client,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
|
||||
def _stream() -> Iterator[bytes]:
|
||||
with cast(OpenAI, openai_client).files.with_streaming_response.content(
|
||||
**file_content_request
|
||||
) as response:
|
||||
yield from response.iter_bytes(chunk_size=chunk_size)
|
||||
|
||||
return _stream()
|
||||
|
||||
async def aretrieve_file(
|
||||
self,
|
||||
file_id: str,
|
||||
@@ -3045,4 +3102,4 @@ class OpenAIAssistantsAPI(BaseLLM):
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return response
|
||||
return response
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
import asyncio
|
||||
import traceback
|
||||
from typing import Any, Optional, cast, get_args
|
||||
from typing import Any, AsyncIterator, Optional, cast, get_args
|
||||
|
||||
import httpx
|
||||
from fastapi import (
|
||||
@@ -21,6 +21,7 @@ from fastapi import (
|
||||
UploadFile,
|
||||
status,
|
||||
)
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
import litellm
|
||||
from litellm import CreateFileRequest, get_secret_str
|
||||
@@ -62,6 +63,88 @@ 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
|
||||
|
||||
|
||||
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_iterator = cast(
|
||||
AsyncIterator[bytes],
|
||||
await litellm.afile_content_streaming(
|
||||
**{
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"file_id": file_id,
|
||||
**data,
|
||||
} # type: ignore
|
||||
),
|
||||
)
|
||||
hidden_params = getattr(stream_iterator, "_hidden_params", {}) or {}
|
||||
response_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:
|
||||
@@ -633,7 +716,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:
|
||||
@@ -731,6 +814,25 @@ async def get_file_content( # noqa: PLR0915
|
||||
check_file_id_encoding=True,
|
||||
)
|
||||
|
||||
if _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(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
file_id=file_id,
|
||||
data=data,
|
||||
should_route=should_route,
|
||||
original_file_id=original_file_id,
|
||||
credentials=credentials,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
version=version,
|
||||
)
|
||||
|
||||
if should_route:
|
||||
# Use model-based routing with credentials from config
|
||||
prepare_data_with_credentials(
|
||||
@@ -738,7 +840,7 @@ 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,
|
||||
|
||||
@@ -2133,6 +2133,8 @@ 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,
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
import pytest
|
||||
from typing import AsyncIterator, cast
|
||||
|
||||
from litellm.files import main as files_main
|
||||
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(
|
||||
monkeypatch,
|
||||
):
|
||||
captured_kwargs = {}
|
||||
|
||||
async def _mock_stream():
|
||||
yield b"hello "
|
||||
yield b"world"
|
||||
|
||||
def _mock_file_content_streaming(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return _mock_stream()
|
||||
|
||||
monkeypatch.setattr(
|
||||
files_main.openai_files_instance,
|
||||
"file_content_streaming",
|
||||
_mock_file_content_streaming,
|
||||
)
|
||||
|
||||
stream_iterator = await files_main.afile_content_streaming(
|
||||
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,
|
||||
)
|
||||
|
||||
async_stream_iterator = cast(AsyncIterator[bytes], stream_iterator)
|
||||
chunks = [chunk async for chunk in async_stream_iterator]
|
||||
|
||||
assert chunks == [b"hello ", b"world"]
|
||||
assert captured_kwargs["_is_async"] is True
|
||||
assert captured_kwargs["file_content_request"]["file_id"] == "file-abc123"
|
||||
assert captured_kwargs["api_key"] == "sk-test"
|
||||
assert captured_kwargs["api_base"] == "https://api.openai.com/v1"
|
||||
assert captured_kwargs["organization"] == "org-123"
|
||||
assert captured_kwargs["chunk_size"] == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_streaming_builds_standard_logging_object_on_completion(
|
||||
monkeypatch,
|
||||
):
|
||||
captured_standard_logging_object = None
|
||||
|
||||
async def _mock_stream():
|
||||
yield b"hello"
|
||||
|
||||
def _mock_file_content_streaming(**kwargs):
|
||||
return _mock_stream()
|
||||
|
||||
async def _mock_async_success_handler(
|
||||
self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
|
||||
):
|
||||
nonlocal captured_standard_logging_object
|
||||
captured_standard_logging_object = kwargs.get("standard_logging_object")
|
||||
self.model_call_details["standard_logging_object"] = captured_standard_logging_object
|
||||
|
||||
monkeypatch.setattr(
|
||||
files_main.openai_files_instance,
|
||||
"file_content_streaming",
|
||||
_mock_file_content_streaming,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLMLoggingObj,
|
||||
"async_success_handler",
|
||||
_mock_async_success_handler,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLMLoggingObj,
|
||||
"handle_sync_success_callbacks_for_async_calls",
|
||||
lambda self, result, start_time, end_time, cache_hit=None: None,
|
||||
)
|
||||
|
||||
stream_iterator = await files_main.afile_content_streaming(
|
||||
file_id="file-abc123",
|
||||
custom_llm_provider="openai",
|
||||
api_key="sk-test",
|
||||
api_base="https://api.openai.com/v1",
|
||||
)
|
||||
|
||||
async_stream_iterator = cast(AsyncIterator[bytes], stream_iterator)
|
||||
chunks = [chunk async for chunk in async_stream_iterator]
|
||||
|
||||
assert chunks == [b"hello"]
|
||||
assert captured_standard_logging_object is not None
|
||||
assert captured_standard_logging_object["call_type"] == "afile_content_streaming"
|
||||
assert captured_standard_logging_object["custom_llm_provider"] == "openai"
|
||||
assert captured_standard_logging_object["response"]["id"] == "file-abc123"
|
||||
assert (
|
||||
captured_standard_logging_object["hidden_params"]["api_base"]
|
||||
== "https://api.openai.com/v1"
|
||||
)
|
||||
@@ -13,6 +13,30 @@ import pytest
|
||||
import yaml
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
def _patch_missing_responses_activate() -> None:
|
||||
"""
|
||||
Ensure tests using @responses.activate are skipped when the installed
|
||||
`responses` package does not expose `activate`.
|
||||
"""
|
||||
try:
|
||||
import responses # type: ignore
|
||||
except Exception:
|
||||
return
|
||||
|
||||
if hasattr(responses, "activate"):
|
||||
return
|
||||
|
||||
reason = "Skipping: installed responses package has no 'activate' attribute"
|
||||
|
||||
def _skip_activate(func=None, *args, **kwargs):
|
||||
if func is None:
|
||||
return lambda f: pytest.mark.skip(reason=reason)(f)
|
||||
return pytest.mark.skip(reason=reason)(func)
|
||||
|
||||
setattr(responses, "activate", _skip_activate)
|
||||
|
||||
|
||||
_patch_missing_responses_activate()
|
||||
|
||||
def build_cache_config(enable_cache: bool = True) -> Optional[Dict]:
|
||||
"""
|
||||
|
||||
@@ -1552,3 +1552,59 @@ def test_file_invalid_anchor_returns_500(
|
||||
)
|
||||
assert response.status_code == 500
|
||||
assert "created_at" in response.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_get_file_content_streams_openai_direct_path(
|
||||
mocker: MockerFixture, monkeypatch, llm_router: Router
|
||||
):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs = {}
|
||||
|
||||
async def _mock_afile_content_streaming(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
async def _stream():
|
||||
yield b"hello "
|
||||
yield b"world"
|
||||
|
||||
return _stream()
|
||||
|
||||
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.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
|
||||
lambda **kwargs: (False, None, None, None),
|
||||
)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/v1/files/file-abc123/content",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
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.headers["content-type"].startswith("application/octet-stream")
|
||||
assert captured_kwargs["custom_llm_provider"] == "openai"
|
||||
assert captured_kwargs["file_id"] == "file-abc123"
|
||||
proxy_logging_obj.update_request_status.assert_awaited_once()
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user