From a61bceb0cf00dd05be46387e8d56a3dfb2daf013 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 14:00:51 -0700 Subject: [PATCH] fix(files): read storage-backed managed files from their storage backend The managed files hook's content read looped the file's model mappings and asked each deployment for the file. A file LiteLLM stored itself maps every model to its storage url, so the read sent that internal id to the upstream server, failed, and the batch rate limiter failed open: a key's TPM limit did not apply to a LiteLLM-executed batch. The hook now returns the stored bytes from the file's storage backend before it consults any deployment --- .../proxy/hooks/managed_files.py | 17 +++++-- .../proxy/test_managed_files_hook.py | 46 +++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 8eef8a5f1ce..09cd0ed192f 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -20,6 +20,7 @@ from typing import ( ) from uuid import NAMESPACE_URL, uuid5 +import httpx from fastapi import HTTPException from pydantic import ValidationError @@ -77,6 +78,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess CreateFileRequest, FileListPage, FileObject, + HttpxBinaryResponseContent, OpenAIFileObject, ResponsesAPIResponse, ) @@ -88,10 +90,6 @@ from litellm.types.utils import ( SpecialEnums, ) -if TYPE_CHECKING: - from litellm.types.llms.openai import HttpxBinaryResponseContent - - if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from prisma.models import ( @@ -1867,10 +1865,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span: Optional[Span], llm_router: Router, **data: Dict, - ) -> "HttpxBinaryResponseContent": + ) -> HttpxBinaryResponseContent: """ Get the content of a file from first model that has it """ + managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span) + if managed_file is not None and managed_file.storage_backend and managed_file.storage_url: + return await self._storage_backend_content(managed_file.storage_backend, managed_file.storage_url) + model_file_id_mapping = data.pop("model_file_id_mapping", None) model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping( [file_id], litellm_parent_otel_span @@ -1900,6 +1902,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): else: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") + async def _storage_backend_content(self, storage_backend_name: str, storage_url: str) -> HttpxBinaryResponseContent: + storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client) + content: Final = await storage_backend.download_file(storage_url) + return HttpxBinaryResponseContent(response=httpx.Response(status_code=httpx.codes.OK, content=content)) + async def _convert_storage_files_to_base64( self, messages: List[AllMessageValues], diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 419a460d098..74bd67efaf2 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -1067,6 +1067,7 @@ async def test_afile_content_passes_trusted_model_credentials_to_router(): managed_files = _make_managed_files_instance() unified_file_id = "unified-file-id" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_unified_file_id = AsyncMock(return_value=None) managed_files.get_model_file_id_mapping = AsyncMock( return_value={unified_file_id: {"model-123": s3_uri}} ) @@ -1238,6 +1239,7 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): managed_files = _make_managed_files_instance() unified_file_id = "unified-file-id" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_unified_file_id = AsyncMock(return_value=None) managed_files.get_model_file_id_mapping = AsyncMock( return_value={unified_file_id: {"model-123": s3_uri}} ) @@ -1268,6 +1270,7 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri(): managed_files = _make_managed_files_instance() unified_file_id = "litellm_proxy_unified_id_abc" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_unified_file_id = AsyncMock(return_value=None) managed_files.get_model_file_id_mapping = AsyncMock( return_value={unified_file_id: {"model-123": s3_uri}} ) @@ -1908,6 +1911,49 @@ async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provid assert response == FileDeleted(id=unified_file_id, object="file", deleted=True) +@pytest.mark.asyncio +async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provider_content(): + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + from prisma import Base64 + + from litellm.caching import DualCache + from litellm.models.managed_files import LiteLLM_ManagedFileTable + + storage_url = "litellm_db://content-row-1" + unified_file_id = _managed_deletion_file_id(storage_url) + stored_bytes = b'{"custom_id": "line-1", "method": "POST", "url": "/v1/chat/completions", "body": {}}\n' + row = LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + model_mappings={"vllm-batch": storage_url}, + flat_model_file_ids=[storage_url], + file_object=_make_file_object(unified_file_id), + storage_backend="litellm_db", + storage_url=storage_url, + ) + file_table = MagicMock(find_first=AsyncMock(return_value=row)) + content_table = MagicMock(find_unique=AsyncMock(return_value=MagicMock(content=Base64.encode(stored_bytes)))) + managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=MagicMock( + db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table) + ), + ) + router = MagicMock( + get_deployment_credentials_with_provider=MagicMock(return_value=None), + afile_content=AsyncMock(), + ) + + response = await managed_files.afile_content( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert response.content == stored_bytes + content_table.find_unique.assert_awaited_once_with(where={"id": "content-row-1"}) + router.afile_content.assert_not_awaited() + + @pytest.mark.asyncio async def test_store_unified_object_id_batch_processed_is_written_only_when_asked(): managed_files, mock_prisma = _make_object_store_instance()