fix(proxy): preserve provider metadata for managed batch files

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -840,6 +840,7 @@ class CheckBatchCost:
                             unified_file_id=_unified_file_id,
                             provider_file_id=_raw_file_id,
                             model_id=model_id,
+                            model_name=managed_file_model_name,
                             owner=_minimal_auth,
                             litellm_parent_otel_span=None,
                             size_bytes=len(content_bytes) if _file_attr == "output_file_id" else None,

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -16,7 +16,6 @@ from typing import (
     Optional,
     Protocol,
     TypedDict,
-    Unpack,
     Union,
     cast,
 )
@@ -25,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
 import httpx
 from fastapi import HTTPException
 from pydantic import ValidationError
-from typing_extensions import ReadOnly
+from typing_extensions import ReadOnly, Unpack
 
 import litellm
 from litellm import Router, verbose_logger
@@ -33,7 +32,7 @@ from litellm._internal_context import with_service_target
 from litellm._uuid import uuid
 from litellm.caching.caching import DualCache
 from litellm.constants import MAX_FILE_LIST_LIMIT
-from litellm.files.types import FileRetrieveCallOptions
+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
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -104,7 +103,7 @@ class _ManagedFileRetrieve(Protocol):
         self,
         *,
         file_id: str,
-        _litellm_internal_model_credentials: Mapping[str, object],
+        _litellm_internal_model_credentials: Mapping[str, object] | None = None,
         **kwargs: Unpack[FileRetrieveCallOptions],
     ) -> OpenAIFileObject: ...
 
@@ -625,6 +624,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 unified_file_id=file_id,
                 provider_file_id=raw_file_id,
                 model_id=model_name or None,
+                model_name=model_name,
                 litellm_parent_otel_span=litellm_parent_otel_span,
                 owner=owner_identity,
             )
@@ -635,27 +635,38 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         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:
+        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)
 
         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 None:
-                raise ValueError("Managed file retrieval requires a configured LLM router")
-            credentials: Final = cast(
-                Mapping[str, object], llm_router.get_deployment_credentials_with_provider(model_id) or {}
-            )
-            retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
-            provider_file_object: Final = await retrieve_file(
-                file_id=provider_file_id,
-                _litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
-                **cast(FileRetrieveCallOptions, credentials),
-            )
-            return provider_file_object.model_copy(update={"id": unified_file_id})
+            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,
+                )
+            ).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)
@@ -666,6 +677,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         unified_file_id: str,
         provider_file_id: str,
         model_id: str | None,
+        model_name: str | None = None,
         owner: UserAPIKeyAuth,
         litellm_parent_otel_span: Span | None,
         size_bytes: int | None = None,
@@ -678,6 +690,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             unified_file_id=unified_file_id,
             provider_file_id=provider_file_id,
             model_id=model_id,
+            model_name=model_name,
             size_bytes=size_bytes,
         )
 
@@ -1564,6 +1577,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                             unified_file_id=unified_file_id,
                             provider_file_id=provider_file_id,
                             model_id=model_id,
+                            model_name=resolved_model_name,
                             owner=user_api_key_dict,
                             litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
                         )

litellm/llms/bedrock/files/transformation.py

@@ -82,7 +82,7 @@ LIST_FILES_LOCATION_PARAM: Final = "_s3_list_files_location"
 def _retrieved_s3_file_size(raw_response: Response) -> int:
     status_code: Final = raw_response.status_code
     if status_code == 206:
-        content_range: Final = cast(str, raw_response.headers.get("Content-Range", ""))
+        content_range: Final = raw_response.headers.get("Content-Range", "")
         range_parts: Final = content_range.removeprefix("bytes 0-0/")
         if content_range.startswith("bytes 0-0/") and range_parts.isdigit():
             return int(range_parts)
@@ -93,7 +93,7 @@ def _retrieved_s3_file_size(raw_response: Response) -> int:
             response=raw_response,
         )
     if status_code == 200:
-        content_length: Final = cast(str, raw_response.headers.get("Content-Length", ""))
+        content_length: Final = raw_response.headers.get("Content-Length", "")
         if content_length.isdigit():
             return int(content_length)
         raise BedrockError(
@@ -1352,7 +1352,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         configured_bucket_name: Final = _listing_bucket_name(litellm_params, "batch_output")
         _, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name)
         relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key
-        last_modified: Final = cast(str, raw_response.headers.get("Last-Modified", ""))
+        last_modified: Final = raw_response.headers.get("Last-Modified", "")
         created_at: Final = int(parsedate_to_datetime(last_modified).timestamp()) if last_modified else 0
         return OpenAIFileObject(
             id=file_id,

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -121,6 +121,7 @@ class ManagedBatchOutputFileWriter(Protocol):
         unified_file_id: str,
         provider_file_id: str,
         model_id: str | None,
+        model_name: str | None = None,
         owner: "UserAPIKeyAuth",
         litellm_parent_otel_span: "Span | None",
         size_bytes: int | None = None,
@@ -1347,6 +1348,7 @@ async def ensure_batch_response_managed_file_ids(
                 unified_file_id=new_unified_file_id,
                 provider_file_id=raw_file_id,
                 model_id=model_id,
+                model_name=model_name,
                 owner=effective_auth,
                 litellm_parent_otel_span=effective_auth.parent_otel_span,
                 size_bytes=None,

tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py

@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
 
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.openai_files_endpoints.common_utils import (
+    ManagedBatchOutputFileWriter,
     ensure_batch_response_managed_file_ids,
     update_batch_in_database,
 )
@@ -39,9 +40,9 @@ def _build_batch_response(
 
 
 def _build_managed_files_mock(unified_id: str = "file-bWFuYWdlZF9vdXRwdXRfaWQ="):
-    mock = MagicMock()
+    mock = MagicMock(spec=ManagedBatchOutputFileWriter)
     mock.get_unified_output_file_id = MagicMock(return_value=unified_id)
-    mock.store_unified_file_id = AsyncMock()
+    mock.store_batch_output_file = AsyncMock()
     return mock
 
 
@@ -114,9 +115,7 @@ async def test_cancel_path_registers_output_file_under_batch_owner():
         operation="cancel",
     )
 
-    forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
-        "user_api_key_dict"
-    ]
+    forwarded_auth = mock_managed_files.store_batch_output_file.await_args.kwargs["owner"]
     assert forwarded_auth.user_id == "batch-owner"
     assert forwarded_auth.team_id == "batch-team"
     stored = json.loads(
@@ -156,9 +155,7 @@ async def test_update_batch_skips_lookup_when_db_batch_object_supplied():
     )
 
     mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_called()
-    forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
-        "user_api_key_dict"
-    ]
+    forwarded_auth = mock_managed_files.store_batch_output_file.await_args.kwargs["owner"]
     assert forwarded_auth.user_id == "caller-owner"
     assert forwarded_auth.team_id == "caller-team"
 
@@ -222,11 +219,11 @@ async def test_ensure_batch_response_swallows_conversion_errors():
         hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
     )
 
-    mock_managed_files = MagicMock()
+    mock_managed_files = MagicMock(spec=ManagedBatchOutputFileWriter)
     mock_managed_files.get_unified_output_file_id = MagicMock(
         side_effect=RuntimeError("boom")
     )
-    mock_managed_files.store_unified_file_id = AsyncMock()
+    mock_managed_files.store_batch_output_file = AsyncMock()
 
     mock_logger = MagicMock()
     await ensure_batch_response_managed_file_ids(
@@ -263,9 +260,7 @@ async def test_ensure_batch_response_builds_auth_from_db_batch_object():
         db_batch_object=db_batch_object,
     )
 
-    forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
-        "user_api_key_dict"
-    ]
+    forwarded_auth = mock_managed_files.store_batch_output_file.await_args.kwargs["owner"]
     assert forwarded_auth.user_id == "user-from-db"
     assert forwarded_auth.team_id == "team-from-db"
 

tests/unit/enterprise/proxy/test_managed_files_hook.py

@@ -183,7 +183,9 @@ def _make_managed_files_instance():
     )
 
     mock_cache = MagicMock()
+    mock_cache.async_get_cache = AsyncMock(return_value=None)
     mock_prisma = MagicMock()
+    mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
 
     instance = PROXY_LiteLLMManagedFiles(
         internal_usage_cache=mock_cache,
@@ -1384,9 +1386,11 @@ def _make_real_managed_files_instance():
     )
 
     mock_cache = MagicMock()
+    mock_cache.async_get_cache = AsyncMock(return_value=None)
     mock_cache.async_set_cache = AsyncMock()
 
     mock_prisma = MagicMock()
+    mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
     mock_prisma.db.litellm_managedfiletable.upsert = AsyncMock()
     mock_prisma.db.litellm_managedfiletable.create = AsyncMock(
         side_effect=AssertionError(

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -1057,6 +1057,7 @@ class TestCheckBatchCost:
             unified_file_id=unified_error_file_id,
             provider_file_id=raw_error_file_id,
             model_id="model-123",
+            model_name="gpt-5-batch",
             owner=mock_hook.store_batch_output_file.await_args.kwargs["owner"],
             litellm_parent_otel_span=None,
             size_bytes=None,