From 065d38181ab80a3e26d1c0f8db2f322acdb6fde7 Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 8 Oct 2026 06:28:38 +0000 Subject: [PATCH] fix(proxy): safely refresh managed batch file details Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 126 +++++++----- .../proxy/hooks/test_managed_files.py | 194 +++++++++++++++++- 2 files changed, 265 insertions(+), 55 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index f1aa4e1e1c9..9437d70227c 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -196,13 +196,29 @@ def _proxy_llm_router() -> Router | None: return cast(Router | None, getattr(proxy_server_module, "llm_router", None)) -def _has_provider_file_retrieve_route( +def _provider_file_retrieve_credentials( *, llm_router: Router | None, model_id: str | None, +) -> Mapping[str, object] | None: + if llm_router is None or model_id is None: + return None + try: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id) + except Exception as error: + verbose_logger.warning( + f"Failed to retrieve credentials for provider file model_id={model_id}: {error}" + ) + return None + return cast(Mapping[str, object], credentials) if credentials else None + + +def _has_provider_file_retrieve_route( + *, + router_credentials: Mapping[str, object] | None, model_name: str | None, ) -> bool: - return (llm_router is not None and model_id is not None) or ( + return bool(router_credentials) or ( model_name is not None and "/" in model_name ) @@ -222,6 +238,8 @@ class _ManagedFileRow(Protocol): class _ManagedFileTableActions(Protocol): async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + async def find_many( self, where: Mapping[str, object], @@ -714,50 +732,55 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_id: str | None, model_name: str | None, llm_router: Router | None = None, - ) -> OpenAIFileObject | None: + raise_on_failure: bool = False, + ) -> tuple[OpenAIFileObject | None, bool]: route_llm_router: Final = llm_router if llm_router is not None else _proxy_llm_router() - if not _has_provider_file_retrieve_route( + router_credentials: Final = _provider_file_retrieve_credentials( llm_router=route_llm_router, model_id=model_id, + ) + model_name_provider: Final = _model_name_file_retrieve_provider(model_name) + fetch_route_available: Final = _has_provider_file_retrieve_route( + router_credentials=router_credentials, model_name=model_name, - ): - return None + ) + if not fetch_route_available: + return None, False try: - if route_llm_router is not None and model_id is not None: - credentials: Final = cast( - Mapping[str, object], - route_llm_router.get_deployment_credentials_with_provider(model_id) or {}, - ) + if router_credentials is not None: provider_file_object: Final = await asyncio.wait_for( self._afile_retrieve_with_retries( provider_file_id=provider_file_id, - call_options=credentials, - internal_model_credentials=credentials, + call_options=router_credentials, + internal_model_credentials=router_credentials, ), timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, ) - return provider_file_object.model_copy(update={"id": unified_file_id}) - if model_name is not None and "/" in model_name: - custom_llm_provider: Final = cast( - FileRetrieveProvider, model_name.split("/", 1)[0] - ) + return provider_file_object.model_copy(update={"id": unified_file_id}), True + if model_name_provider is not None: provider_file_object_by_model_name: Final = await asyncio.wait_for( self._afile_retrieve_with_retries( provider_file_id=provider_file_id, - call_options={"custom_llm_provider": custom_llm_provider}, + call_options={"custom_llm_provider": model_name_provider}, ), timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, ) - return provider_file_object_by_model_name.model_copy( - update={"id": unified_file_id} + return ( + provider_file_object_by_model_name.model_copy( + update={"id": unified_file_id} + ), + True, ) - return None - except Exception as e: + return None, False + except Exception as error: verbose_logger.warning( - f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: {type(e).__name__} {e}" + f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: " + f"{type(error).__name__} {error}" ) - return None + if raise_on_failure: + raise + return None, True async def _save_refreshed_file_object( self, @@ -765,20 +788,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object: OpenAIFileObject, ) -> None: file_object_json: Final = file_object.model_dump_json() - await _managed_file_table(self.prisma_client).upsert( + updated_row_count: Final = await _managed_file_table(self.prisma_client).update_many( 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}, - }, + data={"file_object": file_object_json}, ) + if updated_row_count == 0: + return refreshed_row: Final = stored.model_copy(update={"file_object": file_object}) await self.internal_usage_cache.async_set_cache( key=stored.unified_file_id, @@ -802,17 +817,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if stored_object is not None and not stored_object.litellm_details_fallback: return - fetch_route_available: Final = _has_provider_file_retrieve_route( - llm_router=_proxy_llm_router(), - 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 + provider_fetch_result: Final = ( + (None, True) if fallback_written_recently else await self._fetch_provider_file_object( unified_file_id=unified_file_id, @@ -821,6 +831,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name=model_name, ) ) + provider_object, fetch_route_available = provider_fetch_result if stored_object is not None and provider_object is None and ( size_bytes is None or stored_object.bytes == size_bytes ): @@ -1840,14 +1851,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # so callers see a consistent ID (matching Case 3 which does response.id = file_id). 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: + if file_object.litellm_details_fallback 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( + refreshed_file_object, _ = await self._fetch_provider_file_object( unified_file_id=file_id, provider_file_id=provider_file_id, model_id=model_id, - model_name=None, + model_name=model_id, llm_router=llm_router, ) if refreshed_file_object is None: @@ -1860,26 +1871,33 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # 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. - if not llm_router: + model_mapping: Final = next(iter(stored_file_object.model_mappings.items()), None) + if model_mapping is None: raise Exception( f"LiteLLM Managed File object with id={file_id} has no file_object " f"and llm_router is required to fetch from provider" ) + model_id, model_file_id = model_mapping try: - model_id, model_file_id = next(iter(stored_file_object.model_mappings.items())) - response: Final = await self._fetch_provider_file_object( + response, fetch_route_available = await self._fetch_provider_file_object( unified_file_id=file_id, provider_file_id=model_file_id, model_id=model_id, - model_name=None, + model_name=model_id, llm_router=llm_router, + raise_on_failure=True, ) - 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 + if not fetch_route_available: + raise Exception( + f"LiteLLM Managed File object with id={file_id} has no file_object " + f"and llm_router is required to fetch from provider" + ) + if response is None: + raise ValueError("Provider file details could not be retrieved") + return _public_file_object(response, file_id) async def afile_list( self, diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 536649765c6..9a00f602167 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -17,6 +17,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( is_base64_encoded_unified_file_id, encode_file_id_with_model, ) +from litellm.types.llms.openai import OpenAIFileObject class _InMemoryManagedFileTable: @@ -28,6 +29,9 @@ class _InMemoryManagedFileTable: self.upsert_calls: list[ tuple[Mapping[str, object], Mapping[str, Mapping[str, object]]] ] = [] + self.update_many_calls: list[ + tuple[Mapping[str, object], Mapping[str, object]] + ] = [] async def find_first( self, where: Mapping[str, object] @@ -74,6 +78,29 @@ class _InMemoryManagedFileTable: self.rows[unified_file_id] = row return row + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: + self.update_many_calls.append((where, data)) + unified_file_id: Final = cast(str, where["unified_file_id"]) + existing_row: Final = self.rows.get(unified_file_id) + if existing_row is None: + return 0 + file_object: Final = OpenAIFileObject.model_validate( + json.loads(cast(str, data["file_object"])) + ) + self.rows[unified_file_id] = existing_row.model_copy( + update={"file_object": file_object} + ) + return 1 + + async def delete( + self, where: Mapping[str, object] + ) -> LiteLLM_ManagedFileTable | None: + return self.rows.pop(cast(str, where["unified_file_id"]), None) + async def find_many( self, where: Mapping[str, object], @@ -2472,9 +2499,10 @@ async def test_afile_retrieve_refreshes_marked_fallback_and_preserves_ownership( 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"] == { + assert managed_file_table.update_many_calls[0][1] == { "file_object": response.model_dump_json() } + assert managed_file_table.upsert_calls == [] cached_row = await proxy_managed_files.internal_usage_cache.async_get_cache( key="unified-output", litellm_parent_otel_span=None, @@ -2482,6 +2510,170 @@ async def test_afile_retrieve_refreshes_marked_fallback_and_preserves_ownership( assert cached_row["file_object"]["bytes"] == 836 +@pytest.mark.asyncio +async def test_afile_retrieve_refreshes_marked_fallback_without_router_from_model_name(): + row: Final = _marked_fallback_file_row().model_copy( + update={"model_mappings": {"bedrock/model-x": "provider-output"}} + ) + provider_object: Final = 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) + + with patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + return_value=provider_object, + ) as retrieve: + response: Final = await proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=None, + ) + + assert response.bytes == 836 + assert response.id == "unified-output" + assert response.litellm_details_fallback is None + retrieve.assert_awaited_once() + assert retrieve.await_args.kwargs["custom_llm_provider"] == "bedrock" + assert retrieve.await_args.kwargs["max_retries"] == 0 + assert "_litellm_internal_model_credentials" not in retrieve.await_args.kwargs + updated_row: Final = managed_file_table.rows["unified-output"] + assert updated_row.file_object == response + assert updated_row.model_mappings == {"bedrock/model-x": "provider-output"} + assert updated_row.created_by == "user-123" + assert updated_row.team_id == "team-123" + + +@pytest.mark.asyncio +async def test_afile_retrieve_uses_model_name_when_router_does_not_know_deployment(): + row: Final = _marked_fallback_file_row().model_copy( + update={"model_mappings": {"bedrock/model-x": "provider-output"}} + ) + provider_object: Final = OpenAIFileObject( + id="provider-output", + object="file", + bytes=836, + created_at=456, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = None + proxy_managed_files, _ = _managed_files_with_fake_prisma(row) + + with patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + return_value=provider_object, + ) as retrieve: + response: Final = await proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert response.bytes == 836 + retrieve.assert_awaited_once() + router.get_deployment_credentials_with_provider.assert_called_once_with( + "bedrock/model-x" + ) + assert retrieve.await_args.kwargs["custom_llm_provider"] == "bedrock" + assert retrieve.await_args.kwargs["max_retries"] == 0 + assert "_litellm_internal_model_credentials" not in retrieve.await_args.kwargs + + +@pytest.mark.asyncio +async def test_afile_retrieve_does_not_recreate_deleted_row_after_refresh(): + row: Final = _marked_fallback_file_row() + provider_object: Final = OpenAIFileObject( + id="provider-output", + object="file", + bytes=836, + created_at=456, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row) + + async def delete_row_then_return_file( + *, file_id: str, **_options: object + ) -> OpenAIFileObject: + assert file_id == "provider-output" + deleted_row: Final = await managed_file_table.delete( + where={"unified_file_id": "unified-output"} + ) + assert deleted_row is not None + return provider_object + + with patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + side_effect=delete_row_then_return_file, + ): + response: Final = await proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert response.bytes == 836 + assert "unified-output" not in managed_file_table.rows + assert managed_file_table.upsert_calls == [] + cached_row: Final = await proxy_managed_files.internal_usage_cache.async_get_cache( + key="unified-output", + litellm_parent_otel_span=None, + ) + assert cached_row is None + + +@pytest.mark.asyncio +async def test_afile_retrieve_case3_includes_provider_error_text(): + stored_file: Final = LiteLLM_ManagedFileTable( + unified_file_id="unified-output", + file_object=None, + model_mappings={"model-123": "provider-output"}, + flat_model_file_ids=["provider-output"], + created_by="user-123", + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) + + with ( + patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + side_effect=HTTPException( + status_code=404, + detail="provider file is missing", + ), + ), + pytest.raises(Exception) as error, + ): + await proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert type(error.value) is Exception + assert str(error.value) == ( + "Failed to retrieve file unified-output from provider: " + "404: provider file is missing" + ) + + @pytest.mark.asyncio async def test_afile_retrieve_returns_marked_fallback_without_marker_when_refresh_fails(): router = MagicMock()