scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138

fix(proxy): persist metadata for batch output files

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-10-07 20:45:56 +00:00
parent 54d7c65493
commit 72392916c0
7 changed files with 554 additions and 139 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -2,7 +2,6 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
import time
from collections.abc import Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -20,7 +19,6 @@ from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.types.llms.openai import OpenAIFileObject
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -49,19 +47,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
)
def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes: int) -> OpenAIFileObject:
filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id
return OpenAIFileObject(
id=unified_file_id,
object="file",
purpose="batch_output",
filename=filename,
created_at=int(time.time()),
bytes=size_bytes,
status="processed",
)
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@ -830,15 +815,14 @@ class CheckBatchCost:
custom_llm_provider=custom_llm_provider,
)
# CheckBatchCost bypasses async_post_call_success_hook, so convert raw
# output/error file IDs to managed base64 IDs before the DB write here.
managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files")
if managed_files_hook is not None:
from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter
managed_files_hook: Final = self.proxy_logging_obj.get_proxy_hook("managed_files")
if isinstance(managed_files_hook, ManagedBatchOutputFileWriter):
managed_file_writer: Final = managed_files_hook
from litellm.proxy._types import UserAPIKeyAuth
managed_file_model_name = self._get_managed_file_model_name(
job=job, deployment_info=deployment_info
)
managed_file_model_name = self._get_managed_file_model_name(job=job, deployment_info=deployment_info)
_minimal_auth = UserAPIKeyAuth(
user_id=job.created_by or "default-user-id",
team_id=getattr(job, "team_id", None),
@ -847,21 +831,18 @@ class CheckBatchCost:
_raw_file_id = cast(str | None, getattr(response, _file_attr, None))
if _raw_file_id and not _is_base64_encoded_unified_file_id(_raw_file_id):
try:
_unified_file_id = managed_files_hook.get_unified_output_file_id(
_unified_file_id = managed_file_writer.get_unified_output_file_id(
output_file_id=_raw_file_id,
model_id=model_id,
model_name=managed_file_model_name,
)
await managed_files_hook.store_unified_file_id(
file_id=_unified_file_id,
file_object=_batch_output_file_object(
unified_file_id=_unified_file_id,
raw_file_id=_raw_file_id,
size_bytes=len(content_bytes or b"") if _file_attr == "output_file_id" else 0,
),
await managed_file_writer.store_batch_output_file(
unified_file_id=_unified_file_id,
provider_file_id=_raw_file_id,
model_id=model_id,
owner=_minimal_auth,
litellm_parent_otel_span=None,
model_mappings={model_id: _raw_file_id},
user_api_key_dict=_minimal_auth,
size_bytes=len(content_bytes) if _file_attr == "output_file_id" else None,
)
setattr(response, _file_attr, _unified_file_id)
verbose_proxy_logger.info(

139
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -3,6 +3,7 @@
import base64
import json
import time
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import (
@ -15,6 +16,7 @@ from typing import (
Optional,
Protocol,
TypedDict,
Unpack,
Union,
cast,
)
@ -31,7 +33,7 @@ from litellm._internal_context import with_service_target
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.constants import MAX_FILE_LIST_LIMIT
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.files.types import FileRetrieveCallOptions
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -96,6 +98,17 @@ from litellm.types.utils import (
SpecialEnums,
)
class _ManagedFileRetrieve(Protocol):
async def __call__(
self,
*,
file_id: str,
_litellm_internal_model_credentials: Mapping[str, object],
**kwargs: Unpack[FileRetrieveCallOptions],
) -> OpenAIFileObject: ...
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prisma.models import (
@ -146,6 +159,19 @@ def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) ->
return None
def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes: int) -> OpenAIFileObject:
filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id
return OpenAIFileObject(
id=unified_file_id,
object="file",
purpose="batch_output",
filename=filename,
created_at=int(time.time()),
bytes=size_bytes,
status="processed",
)
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
@ -595,14 +621,74 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
for file_id in provider_file_ids:
model_name = decode_model_from_file_id(file_id)
raw_file_id = get_original_file_id(file_id)
await self.store_unified_file_id(
file_id=file_id,
file_object=None,
await self.store_batch_output_file(
unified_file_id=file_id,
provider_file_id=raw_file_id,
model_id=model_name or None,
litellm_parent_otel_span=litellm_parent_otel_span,
model_mappings={model_name: raw_file_id} if model_name else {},
user_api_key_dict=owner_identity,
owner=owner_identity,
)
async def _retrieve_batch_output_file_object(
self,
*,
unified_file_id: str,
provider_file_id: str,
model_id: str | None,
size_bytes: int | None,
) -> OpenAIFileObject:
if model_id is None:
return _batch_output_file_object(unified_file_id, provider_file_id, size_bytes or 0)
try:
import litellm.proxy.proxy_server as proxy_server_module
llm_router: Final = cast(Router | None, getattr(proxy_server_module, "llm_router", None))
if llm_router is None:
raise ValueError("Managed file retrieval requires a configured LLM router")
credentials: Final = cast(
Mapping[str, object], llm_router.get_deployment_credentials_with_provider(model_id) or {}
)
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
provider_file_object: Final = await retrieve_file(
file_id=provider_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
**cast(FileRetrieveCallOptions, credentials),
)
return provider_file_object.model_copy(update={"id": unified_file_id})
except Exception as e:
verbose_logger.warning(f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: {e}")
return _batch_output_file_object(unified_file_id, provider_file_id, size_bytes or 0)
async def store_batch_output_file(
self,
*,
unified_file_id: str,
provider_file_id: str,
model_id: str | None,
owner: UserAPIKeyAuth,
litellm_parent_otel_span: Span | None,
size_bytes: int | None = None,
) -> None:
stored_file: Final = await self.get_unified_file_id(unified_file_id, litellm_parent_otel_span)
if stored_file is not None and stored_file.file_object is not None:
return
file_object: Final = await self._retrieve_batch_output_file_object(
unified_file_id=unified_file_id,
provider_file_id=provider_file_id,
model_id=model_id,
size_bytes=size_bytes,
)
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
litellm_parent_otel_span=litellm_parent_otel_span,
model_mappings={model_id: provider_file_id} if model_id else {},
user_api_key_dict=owner,
)
async def list_user_batches(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -1474,42 +1560,12 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
setattr(response, file_attr, unified_file_id)
# Use llm_router credentials when available. Without credentials,
# Azure and other auth-required providers return 500/401.
file_object = None
try:
# Import module and use getattr for better testability with mocks
import litellm.proxy.proxy_server as proxy_server_module
_llm_router = getattr(proxy_server_module, "llm_router", None)
if _llm_router is not None and model_id:
_creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {}
file_object = await litellm.afile_retrieve(
file_id=provider_file_id,
**cast(FileRetrieveCallOptions, _creds),
)
else:
file_object = await litellm.afile_retrieve(
custom_llm_provider=cast(
FileRetrieveProvider,
model_name.split("/")[0] if model_name and "/" in model_name else "openai",
),
file_id=provider_file_id,
)
verbose_logger.debug(
f"Successfully retrieved file object for {file_attr}={provider_file_id}"
)
except Exception as e:
verbose_logger.warning(
f"Failed to retrieve file object for {file_attr}={provider_file_id}: {str(e)}. Storing with None and will fetch on-demand."
)
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
await self.store_batch_output_file(
unified_file_id=unified_file_id,
provider_file_id=provider_file_id,
model_id=model_id,
owner=user_api_key_dict,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
model_mappings={model_id: provider_file_id},
user_api_key_dict=user_api_key_dict,
)
request_metadata: Final = data.get("litellm_metadata")
await self.store_unified_object_id(
@ -1639,6 +1695,7 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id) or {}
response = await litellm.afile_retrieve(
file_id=model_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
**cast(FileRetrieveCallOptions, credentials),
)
response.id = file_id # Replace with unified ID

89
litellm/llms/bedrock/files/transformation.py
View file

@ -8,9 +8,10 @@ from collections.abc import Iterable, Mapping, MutableMapping, Sequence
from contextlib import suppress
from dataclasses import dataclass
from datetime import datetime
from email.utils import parsedate_to_datetime
from itertools import chain
from types import MappingProxyType
from typing import Any, Final, Literal, TypeAlias, TypedDict
from typing import Any, Final, Literal, TypeAlias, TypedDict, cast
from urllib.parse import quote, unquote, urlencode
import httpx
@ -70,12 +71,52 @@ from ..common_utils import (
)
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
S3_RETRIEVE_FILE_ID_PARAM: Final = "_s3_retrieve_file_id"
S3_RETRIEVE_FILE_KEY_PARAM: Final = "_s3_retrieve_file_key"
LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose"
LIST_FILES_LOCATION_PARAM: Final = "_s3_list_files_location"
def _retrieved_s3_file_size(raw_response: Response) -> int:
status_code: Final = raw_response.status_code
if status_code == 206:
content_range: Final = cast(str, raw_response.headers.get("Content-Range", ""))
range_parts: Final = content_range.removeprefix("bytes 0-0/")
if content_range.startswith("bytes 0-0/") and range_parts.isdigit():
return int(range_parts)
raise BedrockError(
status_code=status_code,
message=f"Invalid S3 Content-Range header: {content_range}",
headers=raw_response.headers,
response=raw_response,
)
if status_code == 200:
content_length: Final = cast(str, raw_response.headers.get("Content-Length", ""))
if content_length.isdigit():
return int(content_length)
raise BedrockError(
status_code=status_code,
message=f"Invalid S3 Content-Length header: {content_length}",
headers=raw_response.headers,
response=raw_response,
)
if status_code >= 400:
raise BedrockError(
status_code=status_code,
message=raw_response.text,
headers=raw_response.headers,
response=raw_response,
)
raise BedrockError(
status_code=status_code,
message=f"S3 file retrieval returned HTTP {status_code}",
headers=raw_response.headers,
response=raw_response,
)
class _S3DeleteContext(LiteLLMBaseModel):
file_id: str = Field(min_length=1)
@ -1276,18 +1317,52 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def transform_retrieve_file_request(
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
raise NotImplementedError("BedrockFilesConfig does not support file retrieval")
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
_, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params)
url, params = self._transform_s3_file_request(
file_id=file_id,
method="GET",
optional_params=optional_params,
litellm_params=litellm_params,
)
signed_headers_object: Final = litellm_params.get(S3_SIGNED_REQUEST_HEADERS_PARAM)
if not isinstance(signed_headers_object, Mapping):
raise ValueError("S3 request signing did not produce request headers")
signed_headers: Final = cast(Mapping[str, str], signed_headers_object)
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = MappingProxyType({**signed_headers, "Range": "bytes=0-0"})
litellm_params[S3_RETRIEVE_FILE_ID_PARAM] = file_id
litellm_params[S3_RETRIEVE_FILE_KEY_PARAM] = object_key
return url, params
def transform_retrieve_file_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> OpenAIFileObject:
raise NotImplementedError("BedrockFilesConfig does not support file retrieval")
file_id: Final = litellm_params.get(S3_RETRIEVE_FILE_ID_PARAM)
object_key: Final = litellm_params.get(S3_RETRIEVE_FILE_KEY_PARAM)
if not isinstance(file_id, str) or not isinstance(object_key, str):
raise ValueError("S3 retrieve response is missing request context")
file_size: Final = _retrieved_s3_file_size(raw_response)
configured_bucket_name: Final = _listing_bucket_name(litellm_params, "batch_output")
_, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name)
relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key
last_modified: Final = cast(str, raw_response.headers.get("Last-Modified", ""))
created_at: Final = int(parsedate_to_datetime(last_modified).timestamp()) if last_modified else 0
return OpenAIFileObject(
id=file_id,
bytes=file_size,
created_at=created_at,
filename=posixpath.basename(object_key),
object="file",
purpose="batch_output" if relative_key.startswith(BEDROCK_MANAGED_S3_OUTPUT_PREFIX) else "batch",
status="processed",
)
def transform_delete_file_request(
self,

54
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -1,4 +1,5 @@
import base64
import logging
import mimetypes
import re
from collections.abc import Mapping, Sequence
@ -27,6 +28,7 @@ from litellm.types.utils import SpecialEnums
if TYPE_CHECKING:
from fastapi import Request
from opentelemetry.trace import Span
from prisma.models import LiteLLM_ManagedObjectTable
from litellm.proxy._types import UserAPIKeyAuth
@ -104,6 +106,27 @@ class ManagedFileIdResolver(Protocol):
) -> Mapping[str, str]: ...
@runtime_checkable
class ManagedBatchOutputFileWriter(Protocol):
def get_unified_output_file_id(
self,
output_file_id: str,
model_id: str,
model_name: str | None = None,
) -> str: ...
async def store_batch_output_file(
self,
*,
unified_file_id: str,
provider_file_id: str,
model_id: str | None,
owner: "UserAPIKeyAuth",
litellm_parent_otel_span: "Span | None",
size_bytes: int | None = None,
) -> None: ...
def _is_base64_encoded_unified_file_id(b64_uid: object) -> str | Literal[False]:
# Ensure b64_uid is a string and not a mock object
if not isinstance(b64_uid, str):
@ -1279,11 +1302,11 @@ def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mappi
async def ensure_batch_response_managed_file_ids(
response,
managed_files_obj,
prisma_client,
verbose_proxy_logger,
user_api_key_dict=None,
response: "LiteLLMBatch",
managed_files_obj: object | None,
prisma_client: "PrismaClient | None",
verbose_proxy_logger: logging.Logger,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
db_batch_object: object | None = None,
unified_batch_id: str | Literal[False] | None = None,
) -> None:
@ -1291,7 +1314,7 @@ async def ensure_batch_response_managed_file_ids(
await resolve_input_file_id_to_unified(response, prisma_client)
await resolve_output_file_ids_to_unified(response, prisma_client)
if managed_files_obj is None:
if not isinstance(managed_files_obj, ManagedBatchOutputFileWriter):
return
model_id: Final = _model_id_for_batch_response(response, unified_batch_id)
@ -1305,8 +1328,10 @@ async def ensure_batch_response_managed_file_ids(
if effective_auth is None:
return
for file_attr in ("output_file_id", "error_file_id"):
raw_file_id = getattr(response, file_attr, None)
for file_attr, raw_file_id in (
("output_file_id", response.output_file_id),
("error_file_id", response.error_file_id),
):
if not raw_file_id or _is_base64_encoded_unified_file_id(raw_file_id):
continue
try:
@ -1315,12 +1340,13 @@ async def ensure_batch_response_managed_file_ids(
model_id=model_id,
model_name=model_name,
)
await managed_files_obj.store_unified_file_id(
file_id=new_unified_file_id,
file_object=None,
litellm_parent_otel_span=getattr(effective_auth, "parent_otel_span", None),
model_mappings={model_id: raw_file_id},
user_api_key_dict=effective_auth,
await managed_files_obj.store_batch_output_file(
unified_file_id=new_unified_file_id,
provider_file_id=raw_file_id,
model_id=model_id,
owner=effective_auth,
litellm_parent_otel_span=effective_auth.parent_otel_span,
size_bytes=None,
)
setattr(response, file_attr, new_unified_file_id)
verbose_proxy_logger.debug("Converted batch %s %r to managed ID before DB write", file_attr, raw_file_id)

167
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

@ -525,6 +525,8 @@ async def test_output_file_id_for_batch_retrieve():
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=AsyncMock()
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None)
proxy_managed_files.store_unified_file_id = AsyncMock()
response = await proxy_managed_files.async_post_call_success_hook(
data={},
@ -585,6 +587,8 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=AsyncMock()
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None)
proxy_managed_files.store_unified_file_id = AsyncMock()
provider_output_file = OpenAIFileObject(
id="file-provider-output-id",
@ -661,6 +665,8 @@ async def test_error_file_id_for_failed_batch():
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=AsyncMock()
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None)
proxy_managed_files.store_unified_file_id = AsyncMock()
# Create a proper OpenAIFileObject for the error file
error_file_object = OpenAIFileObject(
@ -1514,6 +1520,145 @@ async def test_store_unified_file_id_updates_file_metadata_on_existing_row():
assert second_update["storage_url"] == "s3://bucket/output.jsonl"
@pytest.mark.asyncio
async def test_store_batch_output_file_skips_existing_file_object():
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=AsyncMock(),
)
proxy_managed_files.get_unified_file_id = AsyncMock(
return_value=MagicMock(
file_object=OpenAIFileObject(
id="existing",
object="file",
bytes=1,
created_at=1,
filename="output.jsonl",
purpose="batch_output",
)
)
)
proxy_managed_files.store_unified_file_id = AsyncMock()
with patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve:
await proxy_managed_files.store_batch_output_file(
unified_file_id="unified-output",
provider_file_id="provider-output",
model_id="model-123",
owner=UserAPIKeyAuth(user_id="user-123"),
litellm_parent_otel_span=None,
)
retrieve.assert_not_called()
proxy_managed_files.store_unified_file_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_store_batch_output_file_stores_provider_object_with_unified_id():
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
provider_file_object = OpenAIFileObject(
id="provider-output",
object="file",
bytes=123,
created_at=456,
filename="output.jsonl",
purpose="batch_output",
)
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=AsyncMock(),
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None))
proxy_managed_files.store_unified_file_id = AsyncMock()
with (
patch.object(proxy_server_module, "llm_router", router),
patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_file_object) as retrieve,
):
await proxy_managed_files.store_batch_output_file(
unified_file_id="unified-output",
provider_file_id="provider-output",
model_id="model-123",
owner=UserAPIKeyAuth(user_id="user-123"),
litellm_parent_otel_span=None,
)
retrieve.assert_awaited_once()
stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"]
assert stored_object.id == "unified-output"
assert stored_object.bytes == 123
@pytest.mark.asyncio
async def test_store_batch_output_file_falls_back_when_provider_retrieve_raises():
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import UserAPIKeyAuth
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=AsyncMock(),
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None))
proxy_managed_files.store_unified_file_id = AsyncMock()
with (
patch.object(proxy_server_module, "llm_router", router),
patch("litellm.afile_retrieve", new_callable=AsyncMock, side_effect=RuntimeError("provider unavailable")),
):
await proxy_managed_files.store_batch_output_file(
unified_file_id="unified-output",
provider_file_id="s3://bucket/output.jsonl",
model_id="model-123",
owner=UserAPIKeyAuth(user_id="user-123"),
litellm_parent_otel_span=None,
size_bytes=321,
)
stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"]
assert stored_object.id == "unified-output"
assert stored_object.filename == "output.jsonl"
assert stored_object.bytes == 321
assert stored_object.purpose == "batch_output"
@pytest.mark.asyncio
async def test_store_batch_output_file_falls_back_without_model_id():
from litellm.proxy._types import UserAPIKeyAuth
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=AsyncMock(),
)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None))
proxy_managed_files.store_unified_file_id = AsyncMock()
with patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve:
await proxy_managed_files.store_batch_output_file(
unified_file_id="unified-error",
provider_file_id="error.jsonl",
model_id=None,
owner=UserAPIKeyAuth(user_id="user-123"),
litellm_parent_otel_span=None,
)
retrieve.assert_not_called()
stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"]
assert stored_object.id == "unified-error"
assert stored_object.filename == "error.jsonl"
assert stored_object.bytes == 0
assert stored_object.purpose == "batch_output"
@pytest.mark.asyncio
async def test_afile_delete_returns_provider_response_when_stored_file_object_none():
"""
@ -1956,7 +2101,7 @@ async def test_afile_list_returns_persisted_batch_output_file():
prisma_client = MagicMock()
prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row])
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
@ -3491,8 +3636,10 @@ async def test_post_call_batch_sync_does_not_claim_ownership():
prisma_client = AsyncMock()
prisma_client.db.litellm_managedobjecttable.update_many.return_value = 0
proxy_managed_files = PROXY_LiteLLMManagedFiles(
MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client
MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()),
prisma_client=prisma_client,
)
prisma_client.db.litellm_managedfiletable.find_first.return_value = None
await proxy_managed_files.async_post_call_success_hook(
data={"batch_id": MODEL_ENCODED_BATCH_ID},
@ -3515,22 +3662,16 @@ async def test_post_call_batch_sync_updates_existing_row():
prisma_client.db.litellm_managedobjecttable.find_first.return_value = (
_owned_record(created_by="user_a", team_id="team_a")
)
proxy_managed_files = PROXY_LiteLLMManagedFiles(
MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client
)
proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client)
await proxy_managed_files.async_post_call_success_hook(
data={"batch_id": MODEL_ENCODED_BATCH_ID},
user_api_key_dict=UserAPIKeyAuth(
user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()
),
user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()),
response=_batch_response(MODEL_ENCODED_BATCH_ID),
)
update_call = prisma_client.db.litellm_managedobjecttable.update_many.await_args
assert update_call.kwargs["where"] == {
"unified_object_id": MODEL_ENCODED_BATCH_ID
}
assert update_call.kwargs["where"] == {"unified_object_id": MODEL_ENCODED_BATCH_ID}
assert update_call.kwargs["data"]["status"] == "completed"
prisma_client.db.litellm_managedobjecttable.upsert.assert_not_awaited()
@ -3550,8 +3691,10 @@ async def test_post_call_batch_sync_stores_output_file_ownership_from_batch_row(
_owned_record(created_by="user_a", team_id="team_a")
)
proxy_managed_files = PROXY_LiteLLMManagedFiles(
MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client
MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()),
prisma_client=prisma_client,
)
prisma_client.db.litellm_managedfiletable.find_first.return_value = None
await proxy_managed_files.async_post_call_success_hook(
data={"batch_id": MODEL_ENCODED_BATCH_ID},

133
tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py
View file

@ -2216,6 +2216,139 @@ class TestBedrockFileContentTransformation:
assert "x-amz-content-sha256" in authorization
assert "X-Amz-Date" in signed_headers
def test_transform_retrieve_file_request_adds_unsigned_range(self, monkeypatch):
from litellm.llms.bedrock.files.transformation import (
S3_SIGNED_REQUEST_HEADERS_PARAM,
BedrockFilesConfig,
)
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
url, params = BedrockFilesConfig().transform_retrieve_file_request(
file_id=self.S3_URI,
optional_params={},
litellm_params=litellm_params,
)
assert url == self.EXPECTED_URL
assert params == {}
assert litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Range"] == "bytes=0-0"
@pytest.mark.parametrize(
("file_id", "purpose"),
[
("s3://my-bucket/litellm-batch-outputs/job-123/output.jsonl", "batch_output"),
("s3://my-bucket/litellm-bedrock-files-job-123/input.jsonl", "batch"),
],
)
def test_transform_retrieve_file_response_parses_metadata(
self, file_id: str, purpose: str, monkeypatch: pytest.MonkeyPatch
) -> None:
import httpx
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
BedrockFilesConfig().transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
response = BedrockFilesConfig().transform_retrieve_file_response(
raw_response=httpx.Response(
206,
headers={
"Content-Range": "bytes 0-0/4321",
"Last-Modified": "Wed, 21 Oct 2015 07:28:00 GMT",
},
request=httpx.Request("GET", file_id),
),
logging_obj=MagicMock(),
litellm_params=litellm_params,
)
assert response.id == file_id
assert response.bytes == 4321
assert response.created_at == 1445412480
assert response.filename == file_id.rsplit("/", 1)[-1]
assert response.purpose == purpose
assert response.status == "processed"
assert response.object == "file"
def test_transform_retrieve_file_response_raises_for_empty_object(self, monkeypatch):
import httpx
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
BedrockFilesConfig().transform_retrieve_file_request(
file_id=self.S3_URI,
optional_params={},
litellm_params=litellm_params,
)
with pytest.raises(BedrockError):
BedrockFilesConfig().transform_retrieve_file_response(
raw_response=httpx.Response(
416,
headers={"Content-Range": "bytes */0"},
request=httpx.Request("GET", self.S3_URI),
),
logging_obj=MagicMock(),
litellm_params=litellm_params,
)
def test_transform_retrieve_file_response_uses_content_length_when_range_is_ignored(self, monkeypatch):
import httpx
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
BedrockFilesConfig().transform_retrieve_file_request(
file_id=self.S3_URI,
optional_params={},
litellm_params=litellm_params,
)
response = BedrockFilesConfig().transform_retrieve_file_response(
raw_response=httpx.Response(
200,
headers={"Content-Length": "4321"},
request=httpx.Request("GET", self.S3_URI),
),
logging_obj=MagicMock(),
litellm_params=litellm_params,
)
assert response.bytes == 4321
def test_transform_retrieve_file_response_raises_on_s3_error(self, monkeypatch):
import httpx
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
BedrockFilesConfig().transform_retrieve_file_request(
file_id=self.S3_URI,
optional_params={},
litellm_params=litellm_params,
)
with pytest.raises(BedrockError, match="AccessDenied"):
BedrockFilesConfig().transform_retrieve_file_response(
raw_response=httpx.Response(
403,
content=b"<Error><Code>AccessDenied</Code></Error>",
request=httpx.Request("GET", self.S3_URI),
),
logging_obj=MagicMock(),
litellm_params=litellm_params,
)
def test_transform_file_content_request_decodes_unified_file_id(self, monkeypatch):
"""Base64 unified ids carrying llm_output_file_id must resolve to their S3 object."""
import base64

66
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -24,7 +24,7 @@ _CLAIM_OUTPUT_FILE_ID = "file-output-123"
def test_batch_output_file_object_derives_metadata():
from litellm_enterprise.proxy.common_utils.check_batch_cost import _batch_output_file_object
from litellm_enterprise.proxy.hooks.managed_files import _batch_output_file_object
output_file = _batch_output_file_object(
unified_file_id="unified-output",
@ -1039,9 +1039,11 @@ class TestCheckBatchCost:
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=response)
mock_hook = MagicMock()
from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter
mock_hook = MagicMock(spec=ManagedBatchOutputFileWriter)
mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id]
mock_hook.store_unified_file_id = AsyncMock()
mock_hook.store_batch_output_file = AsyncMock()
check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook
await check_batch_cost_instance.check_batch_cost()
@ -1051,14 +1053,17 @@ class TestCheckBatchCost:
model_id="model-123",
model_name="gpt-5-batch",
)
stored = {
next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"]
for c in mock_hook.store_unified_file_id.call_args_list
}
assert stored == {raw_error_file_id: unified_error_file_id}
for store_call in mock_hook.store_unified_file_id.call_args_list:
assert store_call.kwargs["user_api_key_dict"].user_id == "user-1"
assert store_call.kwargs["user_api_key_dict"].team_id == "team-1"
mock_hook.store_batch_output_file.assert_awaited_once_with(
unified_file_id=unified_error_file_id,
provider_file_id=raw_error_file_id,
model_id="model-123",
owner=mock_hook.store_batch_output_file.await_args.kwargs["owner"],
litellm_parent_otel_span=None,
size_bytes=None,
)
owner = mock_hook.store_batch_output_file.await_args.kwargs["owner"]
assert owner.user_id == "user-1"
assert owner.team_id == "team-1"
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
@ -1529,6 +1534,8 @@ class TestCheckBatchCost:
Without this, GET /batches/{id} returns a raw file ID that cannot be routed
through the proxy, causing API_KEY errors when clients call GET /files/{id}/content.
"""
from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
@ -1563,12 +1570,12 @@ class TestCheckBatchCost:
mock_deployment.model_info.model_dump.return_value = {}
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
mock_hook = MagicMock()
mock_hook = MagicMock(spec=ManagedBatchOutputFileWriter)
mock_hook.get_unified_output_file_id.side_effect = [
fake_managed_output_id,
fake_managed_error_id,
]
mock_hook.store_unified_file_id = AsyncMock()
mock_hook.store_batch_output_file = AsyncMock()
check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook
mock_file_content = MagicMock()
@ -1631,28 +1638,21 @@ class TestCheckBatchCost:
model_id="model-123",
model_name="gpt-5-batch",
)
assert mock_hook.store_unified_file_id.await_count == 2
# {raw_file_id: managed_file_id} for each store call
assert mock_hook.store_batch_output_file.await_count == 2
stored = {
next(iter(c[1]["model_mappings"].values())): c[1]["file_id"]
for c in mock_hook.store_unified_file_id.call_args_list
c.kwargs["provider_file_id"]: c.kwargs["unified_file_id"]
for c in mock_hook.store_batch_output_file.call_args_list
}
assert stored == {
raw_output_file_id: fake_managed_output_id,
raw_error_file_id: fake_managed_error_id,
}
for store_call in mock_hook.store_unified_file_id.call_args_list:
file_object = store_call.kwargs["file_object"]
assert file_object is not None
assert file_object.object == "file"
assert file_object.purpose == "batch_output"
assert file_object.id == store_call.kwargs["file_id"]
stored_file_objects = {
next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_object"]
for c in mock_hook.store_unified_file_id.call_args_list
stored_sizes = {
c.kwargs["provider_file_id"]: c.kwargs["size_bytes"]
for c in mock_hook.store_batch_output_file.call_args_list
}
assert stored_file_objects[raw_output_file_id].bytes == len(mock_file_content.content)
assert stored_file_objects[raw_error_file_id].bytes == 0
assert stored_sizes[raw_output_file_id] == len(mock_file_content.content)
assert stored_sizes[raw_error_file_id] is None
assert mock_response.output_file_id == fake_managed_output_id
assert mock_response.error_file_id == fake_managed_error_id
@ -2227,10 +2227,10 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
from litellm_enterprise.proxy.common_utils.check_batch_cost import (
CheckBatchCost,
)
from enterprise.litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter
from litellm.types.utils import LiteLLMBatch
from enterprise.litellm_enterprise.proxy.hooks.managed_files import (
PROXY_LiteLLMManagedFiles,
)
router = MagicMock()
router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
@ -2241,13 +2241,13 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
deployment.model_info.model_dump.return_value = {}
router.get_deployment = MagicMock(return_value=deployment)
hook = MagicMock()
hook = MagicMock(spec=ManagedBatchOutputFileWriter)
hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: (
PROXY_LiteLLMManagedFiles.get_unified_output_file_id(
None, output_file_id=output_file_id, model_id=model_id, model_name=model_name
)
)
hook.store_unified_file_id = AsyncMock()
hook.store_batch_output_file = AsyncMock()
proxy_logging_obj = MagicMock()
proxy_logging_obj.get_proxy_hook.return_value = hook