scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(proxy): retry provider file details and refresh fallback batch output entries

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-10-08 02:38:55 +00:00
parent 4f8ba125df
commit 646d4c2e46
6 changed files with 593 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,10 +1,11 @@
# What is this?
## This hook is used to check for LiteLLM managed files in the request body, and replace them with model-specific file id
import asyncio
import base64
import json
import time
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -38,7 +39,7 @@ from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_metadata,
)
from openai import AsyncOpenAI
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
@ -158,7 +159,9 @@ 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:
def _batch_output_file_object(
unified_file_id: str, raw_file_id: str, size_bytes: int, *, fallback: bool
) -> OpenAIFileObject:
filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id
return OpenAIFileObject(
id=unified_file_id,
@ -168,6 +171,35 @@ def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes
created_at=int(time.time()),
bytes=size_bytes,
status="processed",
litellm_details_fallback=True if fallback else None,
)
def _public_file_object(file_object: OpenAIFileObject, unified_file_id: str) -> OpenAIFileObject:
return file_object.model_copy(update={"id": unified_file_id, "litellm_details_fallback": None})
def _is_transient_file_retrieve_error(error: Exception) -> bool:
status_code: Final = getattr(error, "status_code", None)
if isinstance(status_code, int):
return status_code in {408, 429} or status_code >= 500
return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
def _proxy_llm_router() -> Router | None:
import litellm.proxy.proxy_server as proxy_server_module
return cast(Router | None, getattr(proxy_server_module, "llm_router", None))
def _has_provider_file_retrieve_route(
*,
llm_router: Router | None,
model_id: str | None,
model_name: str | None,
) -> bool:
return (llm_router is not None and model_id is not None) or (
model_name is not None and "/" in model_name
)
@ -278,13 +310,21 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
def __init__(
self,
internal_usage_cache: InternalUsageCache,
prisma_client: PrismaClient,
*,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
):
self.internal_usage_cache = internal_usage_cache
self.prisma_client = prisma_client
self._sleep = sleep
@staticmethod
def _get_prometheus_logger():
@ -629,47 +669,104 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
owner=owner_identity,
)
async def _retrieve_batch_output_file_object(
async def _afile_retrieve_with_retries(
self,
*,
provider_file_id: str,
call_options: Mapping[str, object],
internal_model_credentials: Mapping[str, object] | None = None,
) -> OpenAIFileObject:
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
retrieve_options: Final = cast(
FileRetrieveCallOptions,
{**call_options, "max_retries": 0},
)
for attempt in range(len(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS) + 1):
if attempt > 0:
await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1])
try:
if internal_model_credentials is None:
return await retrieve_file(
file_id=provider_file_id,
**retrieve_options,
)
return await retrieve_file(
file_id=provider_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)),
**retrieve_options,
)
except Exception as error:
if not _is_transient_file_retrieve_error(error) or attempt == len(
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS
):
raise
raise RuntimeError("Provider file retrieve retry loop ended without a result")
async def _fetch_provider_file_object(
self,
*,
unified_file_id: str,
provider_file_id: str,
model_id: str | None,
model_name: str | None,
size_bytes: int | None,
) -> OpenAIFileObject:
if model_id is None and model_name is None:
return _batch_output_file_object(unified_file_id, provider_file_id, size_bytes or 0)
) -> OpenAIFileObject | None:
llm_router: Final = _proxy_llm_router()
if not _has_provider_file_retrieve_route(
llm_router=llm_router,
model_id=model_id,
model_name=model_name,
):
return None
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 not None and model_id is not None:
credentials: Final = cast(
Mapping[str, object], llm_router.get_deployment_credentials_with_provider(model_id) or {}
)
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
return (
await retrieve_file(
file_id=provider_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
**cast(FileRetrieveCallOptions, credentials),
)
).model_copy(update={"id": unified_file_id})
if model_name is None or "/" not in model_name:
return _batch_output_file_object(unified_file_id, provider_file_id, size_bytes or 0)
custom_llm_provider: Final = cast(FileRetrieveProvider, model_name.split("/", 1)[0])
provider_retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
return (
await provider_retrieve_file(
custom_llm_provider=custom_llm_provider,
file_id=provider_file_id,
provider_file_object: Final = await self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options=credentials,
internal_model_credentials=credentials,
)
).model_copy(update={"id": unified_file_id})
elif model_name is not None and "/" in model_name:
custom_llm_provider: Final = cast(FileRetrieveProvider, model_name.split("/", 1)[0])
provider_file_object: Final = await self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options={"custom_llm_provider": custom_llm_provider},
)
else:
return None
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)
return None
async def _save_refreshed_file_object(
self,
stored: LiteLLM_ManagedFileTable,
file_object: OpenAIFileObject,
) -> None:
file_object_json: Final = file_object.model_dump_json()
await _managed_file_table(self.prisma_client).upsert(
where={"unified_file_id": stored.unified_file_id},
data={
"create": {
"unified_file_id": stored.unified_file_id,
"model_mappings": json.dumps(stored.model_mappings),
"flat_model_file_ids": stored.flat_model_file_ids,
"created_by": stored.created_by,
"team_id": stored.team_id,
"file_object": file_object_json,
},
"update": {"file_object": file_object_json},
},
)
refreshed_row: Final = stored.model_copy(update={"file_object": file_object})
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id,
value=refreshed_row.model_dump(),
litellm_parent_otel_span=None,
)
async def store_batch_output_file(
self,
@ -683,15 +780,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
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:
stored_object: Final = stored_file.file_object if stored_file is not None else None
if stored_object is not None and not stored_object.litellm_details_fallback:
return
file_object: Final = await self._retrieve_batch_output_file_object(
fetch_route_available: Final = _has_provider_file_retrieve_route(
llm_router=_proxy_llm_router(),
model_id=model_id,
model_name=model_name,
)
provider_object: Final = await self._fetch_provider_file_object(
unified_file_id=unified_file_id,
provider_file_id=provider_file_id,
model_id=model_id,
model_name=model_name,
size_bytes=size_bytes,
)
if stored_object is not None and provider_object is None and (
size_bytes is None or stored_object.bytes == size_bytes
):
return
file_object: Final = (
provider_object
if provider_object is not None
else (
stored_object.model_copy(update={"bytes": size_bytes})
if stored_object is not None and size_bytes is not None
else _batch_output_file_object(
unified_file_id,
provider_file_id,
size_bytes or 0,
fallback=fetch_route_available,
)
)
)
await self.store_unified_file_id(
@ -905,7 +1026,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
}
)
return [
parsed_file_object.model_copy(update={"id": row.unified_file_id})
_public_file_object(parsed_file_object, row.unified_file_id)
for row in file_ids
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
]
@ -1691,10 +1812,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Case 2: Managed file and the file object exists in the database
# The stored file_object has the raw provider ID. Replace with the unified ID
# so callers see a consistent ID (matching Case 3 which does response.id = file_id).
if stored_file_object and stored_file_object.file_object:
# Use model_copy to ensure the ID update persists (Pydantic v2 compatibility)
response = stored_file_object.file_object.model_copy(update={"id": file_id})
return response
if stored_file_object.file_object is not None:
file_object: Final = stored_file_object.file_object
if file_object.litellm_details_fallback and llm_router is not None and stored_file_object.model_mappings:
try:
model_id, provider_file_id = next(iter(stored_file_object.model_mappings.items()))
refreshed_file_object: Final = await self._fetch_provider_file_object(
unified_file_id=file_id,
provider_file_id=provider_file_id,
model_id=model_id,
model_name=None,
)
if refreshed_file_object is None:
return _public_file_object(file_object, file_id)
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
return _public_file_object(refreshed_file_object, file_id)
except Exception as error:
verbose_logger.warning(f"Failed to refresh batch file object for file_id={file_id}: {error}")
return _public_file_object(file_object, file_id)
# Case 3: Managed file exists in the database but not the file object (for. e.g the batch task might not have run)
# So we fetch the file object from the provider. We deliberately do not store the result to avoid interfering with batch cost tracking code.
@ -1706,14 +1841,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
try:
model_id, model_file_id = next(iter(stored_file_object.model_mappings.items()))
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: Final = await self._fetch_provider_file_object(
unified_file_id=file_id,
provider_file_id=model_file_id,
model_id=model_id,
model_name=None,
)
response.id = file_id # Replace with unified ID
return response
if response is None:
raise ValueError("Provider file details could not be retrieved")
return _public_file_object(response, file_id)
except Exception as e:
raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e
@ -1778,7 +1914,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
**cursor_args,
)
matches.extend(
parsed_file_object.model_copy(update={"id": row.unified_file_id})
_public_file_object(parsed_file_object, row.unified_file_id)
for row in chunk
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
and (purpose is None or parsed_file_object.purpose == purpose)

18
litellm/types/llms/openai.py
View file

@ -365,6 +365,7 @@ _JsonValue: TypeAlias = object
BATCH_GUARDRAIL_RESPONSE_FIELD: Final = "litellm_batch_guardrail"
LITELLM_DETAILS_FALLBACK_RESPONSE_FIELD: Final = "litellm_details_fallback"
class OpenAIFileObject(LiteLLMBaseModel):
@ -413,6 +414,9 @@ class OpenAIFileObject(LiteLLMBaseModel):
Absent on every other upload, so OpenAI-shaped clients see an unchanged response.
"""
litellm_details_fallback: bool | None = None
"""Set by the LiteLLM proxy on a saved batch output file entry built without provider metadata; stripped from API responses."""
_hidden_params: dict = PrivateAttr(default={"response_cost": 0.0}) # no cost for writing a file
@property
@ -424,13 +428,19 @@ class OpenAIFileObject(LiteLLMBaseModel):
self._hidden_params = hidden_params
@model_serializer(mode="wrap")
def _omit_absent_batch_guardrail( # noqa: ANN202 # annotating it replaces the model's serialization schema
def _omit_absent_proxy_only_fields( # noqa: ANN202 # annotating it replaces the model's serialization schema
self, handler: SerializerFunctionWrapHandler
):
serialized: Final[Mapping[str, object]] = handler(self)
if self.litellm_batch_guardrail is not None:
return serialized
return {key: value for key, value in serialized.items() if key != BATCH_GUARDRAIL_RESPONSE_FIELD}
fields_to_omit: Final = tuple(
field_name
for field_name, value in (
(BATCH_GUARDRAIL_RESPONSE_FIELD, self.litellm_batch_guardrail),
(LITELLM_DETAILS_FALLBACK_RESPONSE_FIELD, self.litellm_details_fallback),
)
if value is None
)
return {key: value for key, value in serialized.items() if key not in fields_to_omit}
def __contains__(self, key) -> bool:
# Define custom behavior for the 'in' operator

1
tests/integration/management/test_batch_output_file_listing.py
View file

@ -346,6 +346,7 @@ def test_output_file_lists_with_basic_details_when_provider_file_lookup_fails(ga
assert (output["purpose"], output["filename"]) == ("batch_output", f"file-out-{batch.scenario.scenario_id}"), (
output
)
assert "litellm_details_fallback" not in output, output
def _tls_context(directory: Path) -> ssl.SSLContext:

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

@ -1,6 +1,7 @@
import base64
import json
from collections.abc import Mapping
import asyncio
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -72,6 +73,13 @@ class _InMemoryManagedFileTable:
self.rows[unified_file_id] = row
return row
async def find_many(
self,
where: Mapping[str, object],
**_query: object,
) -> list[LiteLLM_ManagedFileTable]:
return list(self.rows.values())
class _InMemoryManagedObjectTable:
async def find_first(self, where: Mapping[str, object]) -> None:
@ -103,15 +111,43 @@ class _InMemoryManagedFilesPrismaClient:
def _managed_files_with_fake_prisma(
*rows: LiteLLM_ManagedFileTable,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
) -> tuple[PROXY_LiteLLMManagedFiles, _InMemoryManagedFileTable]:
managed_file_table: Final = _InMemoryManagedFileTable(*rows)
prisma_client: Final = _InMemoryManagedFilesPrismaClient(managed_file_table)
managed_files: Final = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
DualCache(), prisma_client=prisma_client, sleep=sleep
)
return managed_files, managed_file_table
def _marked_fallback_file_row(
*,
size_bytes: int = 0,
created_at: int = 123,
) -> LiteLLM_ManagedFileTable:
from litellm.types.llms.openai import OpenAIFileObject
file_object: Final = OpenAIFileObject(
id="unified-output",
object="file",
bytes=size_bytes,
created_at=created_at,
filename="output.jsonl",
purpose="batch_output",
status="processed",
litellm_details_fallback=True,
)
return LiteLLM_ManagedFileTable(
unified_file_id="unified-output",
file_object=file_object,
model_mappings={"model-123": "provider-output"},
flat_model_file_ids=["provider-output"],
created_by="user-123",
team_id="team-123",
)
def test_get_file_ids_from_messages():
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=MagicMock()
@ -1713,11 +1749,26 @@ async def test_store_batch_output_file_stores_provider_object_with_unified_id():
)
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma()
sleep_calls: list[float] = []
async def record_sleep(delay: float) -> None:
sleep_calls.append(delay)
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
sleep=record_sleep
)
with (
patch.object(proxy_server_module, "llm_router", router),
patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_file_object) as retrieve,
patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=[
HTTPException(status_code=500),
HTTPException(status_code=500),
provider_file_object,
],
) as retrieve,
):
await proxy_managed_files.store_batch_output_file(
unified_file_id="unified-output",
@ -1727,7 +1778,9 @@ async def test_store_batch_output_file_stores_provider_object_with_unified_id():
litellm_parent_otel_span=None,
)
retrieve.assert_awaited_once()
assert retrieve.await_count == 3
assert [call.kwargs["max_retries"] for call in retrieve.call_args_list] == [0, 0, 0]
assert sleep_calls == [0.5, 1.0]
stored_object = managed_file_table.rows["unified-output"].file_object
assert stored_object is not None
assert stored_object.id == "unified-output"
@ -1810,6 +1863,194 @@ async def test_store_batch_output_file_falls_back_without_model_id():
== stored_object.model_dump_json()
)
assert managed_file_table.rows["unified-error"].model_mappings == {}
assert stored_object.litellm_details_fallback is None
assert "litellm_details_fallback" not in stored_object.model_dump()
@pytest.mark.asyncio
async def test_store_batch_output_file_marks_fallback_after_four_transient_failures():
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"}
sleep_calls: list[float] = []
async def record_sleep(delay: float) -> None:
sleep_calls.append(delay)
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
sleep=record_sleep
)
with (
patch.object(proxy_server_module, "llm_router", router),
patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=[HTTPException(status_code=500) for _ in range(4)],
) 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", team_id="team-123"),
litellm_parent_otel_span=None,
size_bytes=836,
)
assert retrieve.await_count == 4
assert [call.kwargs["max_retries"] for call in retrieve.call_args_list] == [0, 0, 0, 0]
assert sleep_calls == [0.5, 1.0, 2.0]
stored_object = managed_file_table.rows["unified-output"].file_object
assert stored_object is not None
assert stored_object.litellm_details_fallback is True
assert json.loads(managed_file_table.upsert_calls[0][1]["create"]["file_object"])[
"litellm_details_fallback"
] is True
@pytest.mark.asyncio
async def test_store_batch_output_file_does_not_retry_not_found():
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"}
sleep_calls: list[float] = []
async def record_sleep(delay: float) -> None:
sleep_calls.append(delay)
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
sleep=record_sleep
)
with (
patch.object(proxy_server_module, "llm_router", router),
patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=404),
) 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()
assert sleep_calls == []
stored_object = managed_file_table.rows["unified-output"].file_object
assert stored_object is not None
assert stored_object.litellm_details_fallback is True
@pytest.mark.asyncio
async def test_store_batch_output_file_refreshes_marked_fallback_on_later_write():
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
row = _marked_fallback_file_row()
provider_object = OpenAIFileObject(
id="provider-output",
object="file",
bytes=836,
created_at=456,
filename="output.jsonl",
purpose="batch_output",
status="processed",
)
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row)
with (
patch.object(proxy_server_module, "llm_router", router),
patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_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", team_id="team-123"),
litellm_parent_otel_span=None,
)
retrieve.assert_awaited_once()
refreshed = managed_file_table.rows["unified-output"].file_object
assert refreshed is not None
assert refreshed.id == "unified-output"
assert refreshed.bytes == 836
assert refreshed.litellm_details_fallback is None
assert managed_file_table.upsert_calls[0][1]["update"]["file_object"] == refreshed.model_dump_json()
@pytest.mark.asyncio
async def test_store_batch_output_file_keeps_marked_fallback_when_retryable_details_still_fail():
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, managed_file_table = _managed_files_with_fake_prisma(
_marked_fallback_file_row()
)
with (
patch.object(proxy_server_module, "llm_router", router),
patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=404),
) 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", team_id="team-123"),
litellm_parent_otel_span=None,
)
retrieve.assert_awaited_once()
assert managed_file_table.upsert_calls == []
assert managed_file_table.rows["unified-output"].file_object.created_at == 123
@pytest.mark.asyncio
async def test_store_batch_output_file_updates_only_size_for_failed_marked_fallback():
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, managed_file_table = _managed_files_with_fake_prisma(
_marked_fallback_file_row()
)
with (
patch.object(proxy_server_module, "llm_router", router),
patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=404),
),
):
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", team_id="team-123"),
litellm_parent_otel_span=None,
size_bytes=836,
)
updated = managed_file_table.rows["unified-output"].file_object
assert updated is not None
assert updated.bytes == 836
assert updated.created_at == 123
assert updated.litellm_details_fallback is True
@pytest.mark.asyncio
@ -1883,19 +2124,21 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none():
"""
from litellm.types.llms.openai import OpenAIFileObject
prisma_client = AsyncMock()
internal_usage_cache = MagicMock()
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=internal_usage_cache,
prisma_client=prisma_client,
stored_file = LiteLLM_ManagedFileTable(
unified_file_id="test-unified-file-id",
file_object=None,
model_mappings={"model-123": "file-provider-xyz"},
flat_model_file_ids=["file-provider-xyz"],
created_by="user-123",
)
sleep_calls: list[float] = []
# Mock get_unified_file_id to return a stored object with file_object=None
stored_file = MagicMock()
stored_file.file_object = None
stored_file.model_mappings = {"model-123": "file-provider-xyz"}
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file)
async def record_sleep(delay: float) -> None:
sleep_calls.append(delay)
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
stored_file, sleep=record_sleep
)
# Mock the router and provider response
provider_file_response = OpenAIFileObject(
@ -1915,7 +2158,11 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none():
}
)
with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_afile_retrieve:
with patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=[HTTPException(status_code=500), provider_file_response],
) as mock_afile_retrieve:
mock_afile_retrieve.return_value = provider_file_response
unified_file_id = "test-unified-file-id"
@ -1928,7 +2175,10 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none():
# Should return the provider response with the unified file ID
assert result is not None
assert result.id == unified_file_id
mock_afile_retrieve.assert_called_once()
assert mock_afile_retrieve.await_count == 2
assert [call.kwargs["max_retries"] for call in mock_afile_retrieve.call_args_list] == [0, 0]
assert sleep_calls == [0.5]
assert managed_file_table.upsert_calls == []
@pytest.mark.asyncio
@ -1937,19 +2187,14 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none()
Test that afile_retrieve raises an appropriate error when file_object is None
and no llm_router is provided to fetch from the provider.
"""
prisma_client = AsyncMock()
internal_usage_cache = MagicMock()
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=internal_usage_cache,
prisma_client=prisma_client,
stored_file = LiteLLM_ManagedFileTable(
unified_file_id="test-unified-file-id",
file_object=None,
model_mappings={"model-123": "file-provider-xyz"},
flat_model_file_ids=["file-provider-xyz"],
created_by="user-123",
)
# Mock get_unified_file_id to return a stored object with file_object=None
stored_file = MagicMock()
stored_file.file_object = None
stored_file.model_mappings = {"model-123": "file-provider-xyz"}
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file)
proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file)
unified_file_id = "test-unified-file-id"
@ -1971,15 +2216,6 @@ async def test_afile_retrieve_returns_stored_file_object_when_exists():
"""
from litellm.types.llms.openai import OpenAIFileObject
prisma_client = AsyncMock()
internal_usage_cache = MagicMock()
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=internal_usage_cache,
prisma_client=prisma_client,
)
# Mock get_unified_file_id to return a stored object WITH file_object
stored_file_object = OpenAIFileObject(
id="test-unified-file-id",
object="file",
@ -1988,18 +2224,93 @@ async def test_afile_retrieve_returns_stored_file_object_when_exists():
filename="input.jsonl",
purpose="batch",
)
stored_file = MagicMock()
stored_file.file_object = stored_file_object
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file)
stored_file = LiteLLM_ManagedFileTable(
unified_file_id="test-unified-file-id",
file_object=stored_file_object,
model_mappings={"model-123": "provider-file-id"},
flat_model_file_ids=["provider-file-id"],
created_by="user-123",
)
proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file)
router = MagicMock()
result = await proxy_managed_files.afile_retrieve(
file_id="test-unified-file-id",
with patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve:
result = await proxy_managed_files.afile_retrieve(
file_id="test-unified-file-id",
litellm_parent_otel_span=None,
llm_router=router,
)
assert result == stored_file_object
retrieve.assert_not_called()
@pytest.mark.asyncio
async def test_afile_retrieve_refreshes_marked_fallback_and_preserves_ownership():
from litellm.types.llms.openai import OpenAIFileObject
row = _marked_fallback_file_row()
provider_object = OpenAIFileObject(
id="provider-output",
object="file",
bytes=836,
created_at=456,
filename="output.jsonl",
purpose="batch_output",
status="processed",
)
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row)
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
with patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_object) as retrieve:
response = await proxy_managed_files.afile_retrieve(
file_id="unified-output",
litellm_parent_otel_span=None,
llm_router=router,
)
assert response == provider_object.model_copy(update={"id": "unified-output"})
assert "litellm_details_fallback" not in response.model_dump()
retrieve.assert_awaited_once()
updated_row = managed_file_table.rows["unified-output"]
assert updated_row.file_object == response
assert updated_row.created_by == "user-123"
assert updated_row.team_id == "team-123"
assert managed_file_table.upsert_calls[0][1]["update"] == {
"file_object": response.model_dump_json()
}
cached_row = await proxy_managed_files.internal_usage_cache.async_get_cache(
key="unified-output",
litellm_parent_otel_span=None,
llm_router=None,
)
assert cached_row["file_object"]["bytes"] == 836
@pytest.mark.asyncio
async def test_afile_retrieve_returns_marked_fallback_without_marker_when_refresh_fails():
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
_marked_fallback_file_row()
)
# Should return the stored file object directly
assert result == stored_file_object
with patch(
"litellm.afile_retrieve",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=404),
) as retrieve:
response = await proxy_managed_files.afile_retrieve(
file_id="unified-output",
litellm_parent_otel_span=None,
llm_router=router,
)
assert response.bytes == 0
assert response.id == "unified-output"
assert "litellm_details_fallback" not in response.model_dump()
retrieve.assert_awaited_once()
assert managed_file_table.upsert_calls == []
@pytest.mark.asyncio
@ -2008,16 +2319,7 @@ async def test_afile_retrieve_raises_error_for_non_managed_file():
Test that afile_retrieve raises an error when the file_id is not found
in the managed files table.
"""
prisma_client = AsyncMock()
internal_usage_cache = MagicMock()
proxy_managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=internal_usage_cache,
prisma_client=prisma_client,
)
# Mock get_unified_file_id to return None (file not found)
proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None)
proxy_managed_files, _ = _managed_files_with_fake_prisma()
with pytest.raises(Exception, match='LiteLLM Managed File object with id=non-existent-file-id') as exc_info:
await proxy_managed_files.afile_retrieve(
@ -2237,25 +2539,9 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids():
@pytest.mark.asyncio
async def test_afile_list_returns_persisted_batch_output_file():
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
unified_file_id = "managed-output-file-id"
row = MagicMock()
row.unified_file_id = unified_file_id
row.file_object = OpenAIFileObject(
id=unified_file_id,
object="file",
purpose="batch_output",
filename="output.jsonl.out",
created_at=1,
bytes=42,
status="processed",
).model_dump()
prisma_client = MagicMock()
prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row])
proxy_managed_files = PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
proxy_managed_files, _ = _managed_files_with_fake_prisma(
_marked_fallback_file_row()
)
result = await proxy_managed_files.afile_list(
@ -2266,10 +2552,29 @@ async def test_afile_list_returns_persisted_batch_output_file():
)
listed = result.data[0]
assert listed.id == unified_file_id
assert listed.filename == "output.jsonl.out"
assert listed.bytes == 42
assert listed.id == "unified-output"
assert listed.filename == "output.jsonl"
assert listed.bytes == 0
assert listed.purpose == "batch_output"
assert "litellm_details_fallback" not in listed.model_dump()
@pytest.mark.asyncio
async def test_get_user_created_file_ids_strips_fallback_marker():
from litellm.proxy._types import UserAPIKeyAuth
proxy_managed_files, _ = _managed_files_with_fake_prisma(
_marked_fallback_file_row()
)
files = await proxy_managed_files.get_user_created_file_ids(
UserAPIKeyAuth(user_id="user-123"),
["provider-output"],
)
assert len(files) == 1
assert files[0].id == "unified-output"
assert "litellm_details_fallback" not in files[0].model_dump()
@pytest.mark.asyncio

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

@ -30,11 +30,13 @@ def test_batch_output_file_object_derives_metadata():
unified_file_id="unified-output",
raw_file_id="s3://bucket/path/to/output.jsonl.out",
size_bytes=4321,
fallback=False,
)
provider_file = _batch_output_file_object(
unified_file_id="unified-provider",
raw_file_id="file-abc",
size_bytes=0,
fallback=False,
)
assert output_file.id == "unified-output"

11
tests/unit/types/llms/test_types_llms_openai.py
View file

@ -552,6 +552,17 @@ class TestOpenAIFileObjectBatchGuardrailSerialization:
original = self._file_object(litellm_batch_guardrail=self._report())
assert OpenAIFileObject(**original.model_dump()) == original
def test_details_fallback_marker_is_omitted_when_unset_and_round_trips_when_set(self):
from litellm.types.llms.openai import OpenAIFileObject
without_marker = self._file_object()
assert "litellm_details_fallback" not in without_marker.model_dump()
assert "litellm_details_fallback" not in without_marker.model_dump_json()
with_marker = self._file_object(litellm_details_fallback=True)
assert with_marker.model_dump()["litellm_details_fallback"] is True
assert OpenAIFileObject.model_validate_json(with_marker.model_dump_json()) == with_marker
def test_serialization_json_schema_still_describes_the_model(self):
"""A return annotation on the wrap serializer would collapse this to a bare object."""
from litellm.types.llms.openai import OpenAIFileObject