From 79b4705986683bf250d8a5deab76538c2d9be807 Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 8 Oct 2026 05:13:37 +0000 Subject: [PATCH] test(proxy): cover model-name managed batch file retrieval Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 28 ++++-- .../test_batch_output_file_listing.py | 6 +- .../proxy/hooks/test_managed_files.py | 95 +++++++++++++++++++ 3 files changed, 116 insertions(+), 13 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 364b6a56ce4..b768ef72e1d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/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") diff --git a/tests/integration/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index 9756f050701..26d1ae8de46 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/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, diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 0e4b32ca25d..b32523657b8 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/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