diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 354d816416d..f1aa4e1e1c9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -32,7 +32,11 @@ from litellm import Router, verbose_logger from litellm._internal_context import with_service_target from litellm._uuid import uuid from litellm.caching.caching import DualCache -from litellm.constants import BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, MAX_FILE_LIST_LIMIT +from litellm.constants import ( + BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS, + BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, + MAX_FILE_LIST_LIMIT, +) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR @@ -803,11 +807,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): 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, + fallback_written_recently: Final = ( + stored_object is not None + and time.time() - stored_object.created_at < BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS + ) + provider_object: Final = ( + None + if fallback_written_recently + else 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, + ) ) if stored_object is not None and provider_object is None and ( size_bytes is None or stored_object.bytes == size_bytes diff --git a/litellm/constants.py b/litellm/constants.py index 1001a4ad91d..7905212da87 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1834,6 +1834,7 @@ PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVA MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0 +BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60 STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))) # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 3dc896688a6..536649765c6 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -1,5 +1,6 @@ import base64 import json +import time import asyncio from collections.abc import Awaitable, Callable, Mapping from typing import Final, cast @@ -2074,6 +2075,36 @@ async def test_store_batch_output_file_marks_fallback_after_four_transient_failu ] is True +@pytest.mark.asyncio +async def test_store_batch_output_file_skips_lookup_for_fallback_written_moments_ago(): + 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(created_at=int(time.time())) + ) + + with ( + patch.object(proxy_server_module, "llm_router", router), + 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", team_id="team-123"), + litellm_parent_otel_span=None, + size_bytes=836, + ) + + retrieve.assert_not_awaited() + stored_object = managed_file_table.rows["unified-output"].file_object + assert stored_object is not None + assert (stored_object.bytes, stored_object.litellm_details_fallback) == (836, 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