From 646d4c2e46a47227690e995a1bd8b0cd5bd9f02c Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 8 Oct 2026 02:38:55 +0000 Subject: [PATCH] fix(proxy): retry provider file details and refresh fallback batch output entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 228 +++++++-- litellm/types/llms/openai.py | 18 +- .../test_batch_output_file_listing.py | 1 + .../proxy/hooks/test_managed_files.py | 461 +++++++++++++++--- .../common_utils/test_check_batch_cost.py | 2 + .../unit/types/llms/test_types_llms_openai.py | 11 + 6 files changed, 593 insertions(+), 128 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d9d2929a8d9..364b6a56ce4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index f7e3d644f34..cd77d6088f6 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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 diff --git a/tests/integration/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index 606c7b9718e..291587eab76 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/tests/integration/management/test_batch_output_file_listing.py @@ -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: diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index c24215d947b..0e4b32ca25d 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 3d844e8b9e8..094356108e0 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -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" diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py index e59643f6509..871e4cc8cfb 100644 --- a/tests/unit/types/llms/test_types_llms_openai.py +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -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