mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-22 20:25:12 +00:00
Merge pull request #25569 from BerriAI/litellm_harish_april11
Litellm harish april11
This commit is contained in:
ishaan-berri
+134
-16
@@ -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,Coroutine, Dict, Literal, Optional, Union, cast
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -30,12 +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 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
|
||||
@@ -55,10 +53,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,
|
||||
@@ -69,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()
|
||||
@@ -772,8 +776,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 +793,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 +824,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 +874,23 @@ def file_content(
|
||||
|
||||
_is_async = kwargs.pop("afile_content", False) is True
|
||||
|
||||
if stream and _should_sdk_support_streaming(custom_llm_provider):
|
||||
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 +1014,89 @@ 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 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,
|
||||
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)
|
||||
@@ -0,0 +1,236 @@
|
||||
import datetime
|
||||
import traceback
|
||||
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 (
|
||||
Logging as LiteLLMLoggingObj,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingHiddenParams, StandardLoggingPayload
|
||||
|
||||
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: Dict[str, Any] = {}
|
||||
self._logging_completed = False
|
||||
self._close_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()
|
||||
)
|
||||
self._sync_hidden_params()
|
||||
|
||||
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
|
||||
|
||||
async def aclose(self) -> None:
|
||||
if self._close_completed:
|
||||
return
|
||||
|
||||
self._close_completed = True
|
||||
self._logging_completed = True
|
||||
stream_to_close = self.stream_iterator
|
||||
self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(()))
|
||||
|
||||
# Shield cleanup from request cancellation so upstream HTTP connections
|
||||
# are released promptly on client disconnects.
|
||||
with anyio.CancelScope(shield=True):
|
||||
if hasattr(stream_to_close, "aclose"):
|
||||
await cast(AsyncIterator[bytes], stream_to_close).aclose() # type: ignore[attr-defined]
|
||||
elif hasattr(stream_to_close, "close"):
|
||||
result = cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined]
|
||||
if result is not None:
|
||||
await result
|
||||
|
||||
def close(self) -> None:
|
||||
if self._close_completed:
|
||||
return
|
||||
|
||||
self._close_completed = True
|
||||
self._logging_completed = True
|
||||
stream_to_close = self.stream_iterator
|
||||
self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(()))
|
||||
|
||||
if hasattr(stream_to_close, "close"):
|
||||
cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined]
|
||||
|
||||
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: dict[str, Any] = {}
|
||||
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
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
|
||||
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(Dict[str, Any], 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
|
||||
)
|
||||
@@ -0,0 +1,11 @@
|
||||
from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Union
|
||||
|
||||
|
||||
FileContentProvider = Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
|
||||
]
|
||||
|
||||
|
||||
class FileContentStreamingResult(NamedTuple):
|
||||
stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]]
|
||||
headers: Dict[str, str]
|
||||
@@ -32,6 +32,7 @@ import litellm
|
||||
from litellm import LlmProviders
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_RETRIES
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
@@ -1751,6 +1752,92 @@ 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,
|
||||
) -> FileContentStreamingResult:
|
||||
response_cm = openai_client.files.with_streaming_response.content(
|
||||
**file_content_request
|
||||
)
|
||||
response = await response_cm.__aenter__()
|
||||
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:
|
||||
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)
|
||||
|
||||
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,
|
||||
) -> FileContentStreamingResult:
|
||||
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,
|
||||
)
|
||||
|
||||
response_cm = cast(OpenAI, openai_client).files.with_streaming_response.content(
|
||||
**file_content_request
|
||||
)
|
||||
response = response_cm.__enter__()
|
||||
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:
|
||||
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)
|
||||
|
||||
async def aretrieve_file(
|
||||
self,
|
||||
file_id: str,
|
||||
@@ -3045,4 +3132,4 @@ class OpenAIAssistantsAPI(BaseLLM):
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return response
|
||||
return response
|
||||
@@ -0,0 +1,149 @@
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Optional, Tuple, cast
|
||||
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
import litellm
|
||||
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
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
|
||||
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,
|
||||
) -> bool:
|
||||
return (
|
||||
custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS
|
||||
)
|
||||
|
||||
@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],
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
version: str,
|
||||
) -> StreamingResponse:
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
stream_result = cast(
|
||||
FileContentStreamingResult,
|
||||
await litellm.afile_content(
|
||||
**{
|
||||
"custom_llm_provider": cast(
|
||||
FileContentProvider, 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,
|
||||
)
|
||||
@@ -21,7 +21,6 @@ from fastapi import (
|
||||
UploadFile,
|
||||
status,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm import CreateFileRequest, get_secret_str
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
@@ -47,7 +46,7 @@ from litellm.types.llms.openai import (
|
||||
OpenAIFilesPurpose,
|
||||
)
|
||||
|
||||
from .common_utils import (
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
encode_file_id_with_model,
|
||||
extract_file_creation_params,
|
||||
@@ -55,7 +54,6 @@ from .common_utils import (
|
||||
handle_model_based_routing,
|
||||
prepare_data_with_credentials,
|
||||
)
|
||||
from .storage_backend_service import StorageBackendFileService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@@ -159,6 +157,9 @@ async def route_create_file(
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_data,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.storage_backend_service import (
|
||||
StorageBackendFileService,
|
||||
)
|
||||
|
||||
# Extract file data
|
||||
file_data = extract_file_data(cast(Any, _create_file_request.get("file")))
|
||||
@@ -633,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:
|
||||
@@ -731,6 +732,41 @@ async def get_file_content( # noqa: PLR0915
|
||||
check_file_id_encoding=True,
|
||||
)
|
||||
|
||||
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=resolved_custom_llm_provider,
|
||||
):
|
||||
verbose_proxy_logger.debug(
|
||||
"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=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,
|
||||
)
|
||||
|
||||
if should_route:
|
||||
# Use model-based routing with credentials from config
|
||||
prepare_data_with_credentials(
|
||||
@@ -738,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,
|
||||
|
||||
@@ -243,8 +243,6 @@ from typing import (
|
||||
get_args,
|
||||
)
|
||||
|
||||
from openai import OpenAIError as OriginalError
|
||||
|
||||
# These are lazy loaded via __getattr__
|
||||
from litellm.llms.base_llm.base_utils import (
|
||||
BaseLLMModelInfo,
|
||||
|
||||
@@ -27,8 +27,16 @@ async def test_file_operations():
|
||||
|
||||
get_file_content = await openai_client.files.content(file_id=uploaded_file.id)
|
||||
print("get_file_content=", get_file_content.content)
|
||||
response = get_file_content.response
|
||||
|
||||
assert get_file_content.content == file_content
|
||||
assert response.status_code == 200
|
||||
assert response.headers.get("content-type") == "application/octet-stream"
|
||||
assert response.headers.get("content-length") is not None
|
||||
assert int(response.headers["content-length"]) == len(get_file_content.content)
|
||||
assert response.headers.get("content-disposition") is not None
|
||||
assert uploaded_file.filename in response.headers["content-disposition"]
|
||||
assert response.headers.get("x-request-id") is not None
|
||||
# try get_file_content.write_to_file
|
||||
get_file_content.write_to_file("get_file_content.jsonl")
|
||||
|
||||
|
||||
@@ -0,0 +1,323 @@
|
||||
import pytest
|
||||
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_with_stream_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 FileContentStreamingResult(
|
||||
stream_iterator=_mock_stream(),
|
||||
headers={"content-length": "11"},
|
||||
)
|
||||
|
||||
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",
|
||||
organization="org-123",
|
||||
chunk_size=8,
|
||||
stream=True,
|
||||
),
|
||||
)
|
||||
|
||||
async_stream_iterator = cast(AsyncIterator[bytes], stream_result.stream_iterator)
|
||||
chunks = [chunk async for chunk in async_stream_iterator]
|
||||
|
||||
assert chunks == [b"hello ", b"world"]
|
||||
assert stream_result.headers["content-length"] == "11"
|
||||
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
|
||||
assert captured_kwargs["client"] is None
|
||||
|
||||
|
||||
@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 FileContentStreamingResult(
|
||||
stream_iterator=_mock_stream(),
|
||||
headers={"content-length": "5"},
|
||||
)
|
||||
|
||||
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_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)
|
||||
chunks = [chunk async for chunk in async_stream_iterator]
|
||||
|
||||
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"
|
||||
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"
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
async def _mock_stream():
|
||||
nonlocal close_called
|
||||
try:
|
||||
yield b"hello"
|
||||
yield b"world"
|
||||
finally:
|
||||
close_called = True
|
||||
|
||||
stream = FileContentStreamingResponse(
|
||||
stream_iterator=_mock_stream(),
|
||||
file_id="file-abc123",
|
||||
model="gpt-4o",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=None,
|
||||
)
|
||||
|
||||
assert await stream.__anext__() == b"hello"
|
||||
|
||||
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
|
||||
@@ -13,7 +13,6 @@ import pytest
|
||||
import yaml
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
def build_cache_config(enable_cache: bool = True) -> Optional[Dict]:
|
||||
"""
|
||||
Build Redis cache configuration from environment variables.
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import ANY
|
||||
from unittest.mock import ANY, AsyncMock
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
@@ -14,11 +15,15 @@ sys.path.insert(
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
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
|
||||
from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import (
|
||||
FileContentStreamingHandler,
|
||||
)
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIFileObject
|
||||
|
||||
client = TestClient(app)
|
||||
from litellm.caching.caching import DualCache
|
||||
@@ -77,6 +82,127 @@ def setup_proxy_logging_object(monkeypatch, llm_router: Router) -> ProxyLogging:
|
||||
return proxy_logging_object
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_file_content_with_logging_closes_inner_iterator_on_early_exit():
|
||||
class MockStreamIterator:
|
||||
def __init__(self) -> None:
|
||||
self._chunks = iter([b"hello", b"world"])
|
||||
self.aclose = AsyncMock()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self._chunks)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration
|
||||
|
||||
stream_iterator = MockStreamIterator()
|
||||
proxy_logging_obj = AsyncMock()
|
||||
|
||||
generator = FileContentStreamingHandler.stream_file_content_with_logging(
|
||||
stream_iterator=stream_iterator,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=AsyncMock(),
|
||||
data={"litellm_call_id": "call-123"},
|
||||
)
|
||||
|
||||
assert await generator.__anext__() == b"hello"
|
||||
|
||||
await generator.aclose()
|
||||
|
||||
stream_iterator.aclose.assert_awaited_once()
|
||||
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
|
||||
@@ -283,7 +409,7 @@ def test_target_storage_invokes_storage_backend(
|
||||
)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.StorageBackendFileService.upload_file_to_storage_backend",
|
||||
"litellm.proxy.openai_files_endpoints.storage_backend_service.StorageBackendFileService.upload_file_to_storage_backend",
|
||||
new=async_mock,
|
||||
)
|
||||
|
||||
@@ -342,7 +468,7 @@ def test_target_storage_with_target_models(
|
||||
)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.StorageBackendFileService.upload_file_to_storage_backend",
|
||||
"litellm.proxy.openai_files_endpoints.storage_backend_service.StorageBackendFileService.upload_file_to_storage_backend",
|
||||
new=async_mock,
|
||||
)
|
||||
|
||||
@@ -1552,3 +1678,198 @@ 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(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
async def _stream():
|
||||
yield b"hello "
|
||||
yield b"world"
|
||||
|
||||
return FileContentStreamingResult(
|
||||
stream_iterator=_stream(),
|
||||
headers={"content-length": "11"},
|
||||
)
|
||||
|
||||
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),
|
||||
)
|
||||
|
||||
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 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()
|
||||
|
||||
|
||||
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
|
||||
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(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
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: (
|
||||
True,
|
||||
"azure-gpt-3-5-turbo",
|
||||
"file-original-123",
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "azure-key",
|
||||
"api_base": "https://azure.example.com",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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"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 "stream" not in captured_kwargs
|
||||
mock_streaming_response.assert_not_awaited()
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
def test_get_file_content_non_openai_provider_skips_streaming_handler(
|
||||
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.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs = {}
|
||||
|
||||
async def _mock_afile_content(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
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: (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",
|
||||
"custom-llm-provider": "azure",
|
||||
},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.content == b"azure-bytes"
|
||||
assert captured_kwargs["custom_llm_provider"] == "azure"
|
||||
assert "stream" not in captured_kwargs
|
||||
mock_streaming_response.assert_not_awaited()
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user