test(proxy): cover model-name managed batch file retrieval

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

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -709,34 +709,40 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         provider_file_id: str,
         model_id: str | None,
         model_name: str | None,
+        llm_router: Router | None = None,
     ) -> OpenAIFileObject | None:
-        llm_router: Final = _proxy_llm_router()
+        route_llm_router: Final = llm_router if llm_router is not None else _proxy_llm_router()
         if not _has_provider_file_retrieve_route(
-            llm_router=llm_router,
+            llm_router=route_llm_router,
             model_id=model_id,
             model_name=model_name,
         ):
             return None
 
         try:
-            if llm_router is not None and model_id is not None:
+            if route_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 {}
+                    Mapping[str, object],
+                    route_llm_router.get_deployment_credentials_with_provider(model_id) or {},
                 )
                 provider_file_object: Final = await self._afile_retrieve_with_retries(
                     provider_file_id=provider_file_id,
                     call_options=credentials,
                     internal_model_credentials=credentials,
                 )
-            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(
+                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]
+                )
+                provider_file_object_by_model_name: 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})
+                return provider_file_object_by_model_name.model_copy(
+                    update={"id": unified_file_id}
+                )
+            return None
         except Exception as e:
             verbose_logger.warning(f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: {e}")
             return None
@@ -1822,6 +1828,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                         provider_file_id=provider_file_id,
                         model_id=model_id,
                         model_name=None,
+                        llm_router=llm_router,
                     )
                     if refreshed_file_object is None:
                         return _public_file_object(file_object, file_id)
@@ -1846,6 +1853,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 provider_file_id=model_file_id,
                 model_id=model_id,
                 model_name=None,
+                llm_router=llm_router,
             )
             if response is None:
                 raise ValueError("Provider file details could not be retrieved")

tests/integration/management/test_batch_output_file_listing.py

@@ -495,13 +495,13 @@ class _BedrockControlPlane:
             )
             return Reply(body=json.dumps({"jobArn": self.job_arn}).encode())
         if request.method == "GET" and request.target.endswith(self.job_id):
-            input_uri, output_uri, job_name = self.job_locations.get()
-            self.job_locations.put((input_uri, output_uri, job_name))
+            input_uri, output_uri, retrieved_job_name = self.job_locations.get()
+            self.job_locations.put((input_uri, output_uri, retrieved_job_name))
             return Reply(
                 body=json.dumps(
                     {
                         "jobArn": self.job_arn,
-                        "jobName": job_name,
+                        "jobName": retrieved_job_name,
                         "modelId": BEDROCK_MODEL_ID,
                         "status": "Completed",
                         "submitTime": 1700000000,

tests/unit/enterprise/proxy/hooks/test_managed_files.py

@@ -1794,6 +1794,101 @@ async def test_store_batch_output_file_stores_provider_object_with_unified_id():
     }
 
 
+@pytest.mark.asyncio
+async def test_store_batch_output_file_retries_model_name_provider_fetch():
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.llms.openai import OpenAIFileObject
+
+    provider_file_object = OpenAIFileObject(
+        id="provider-output",
+        object="file",
+        bytes=123,
+        created_at=456,
+        filename="output.jsonl",
+        purpose="batch_output",
+    )
+    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(
+        "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",
+            provider_file_id="provider-output",
+            model_id=None,
+            model_name="bedrock/model-x",
+            owner=UserAPIKeyAuth(user_id="user-123"),
+            litellm_parent_otel_span=None,
+        )
+
+    assert retrieve.await_count == 3
+    assert all(
+        call.kwargs["custom_llm_provider"] == "bedrock"
+        and call.kwargs["max_retries"] == 0
+        and "_litellm_internal_model_credentials" not in call.kwargs
+        for call in retrieve.call_args_list
+    )
+    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"
+    assert stored_object.bytes == 123
+    assert stored_object.litellm_details_fallback is None
+
+
+@pytest.mark.asyncio
+async def test_store_batch_output_file_marks_model_name_fallback_after_four_transient_failures():
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    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(
+        "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=None,
+            model_name="bedrock/model-x",
+            owner=UserAPIKeyAuth(user_id="user-123"),
+            litellm_parent_otel_span=None,
+            size_bytes=836,
+        )
+
+    assert retrieve.await_count == 4
+    assert all(
+        call.kwargs["custom_llm_provider"] == "bedrock"
+        and call.kwargs["max_retries"] == 0
+        and "_litellm_internal_model_credentials" not in call.kwargs
+        for call in retrieve.call_args_list
+    )
+    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.bytes == 836
+    assert stored_object.litellm_details_fallback is True
+
+
 @pytest.mark.asyncio
 async def test_store_batch_output_file_falls_back_when_provider_retrieve_raises():
     import litellm.proxy.proxy_server as proxy_server_module