From 9721d9f82b11bc8d29d29346adeb6dc3e4434e6a Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 8 Oct 2026 16:23:57 +0000 Subject: [PATCH] fix(proxy): keep batch listing provider-free Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 96 +++- .../openai_files_endpoints/common_utils.py | 4 + .../proxy/hooks/test_managed_files.py | 497 +++++++++++++++++- .../common_utils/test_check_batch_cost.py | 5 + 4 files changed, 585 insertions(+), 17 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 9437d70227c..ff91ee11c73 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -213,14 +213,38 @@ def _provider_file_retrieve_credentials( return cast(Mapping[str, object], credentials) if credentials else None +_PROVIDER_FILE_RETRIEVE_PROVIDERS: Final[frozenset[str]] = frozenset( + { + "openai", + "azure", + "gemini", + "vertex_ai", + "bedrock", + "hosted_vllm", + "litellm_proxy", + "manus", + "anthropic", + "mistral", + "xai", + } +) + + +def _model_name_file_retrieve_provider(model_name: str | None) -> FileRetrieveProvider | None: + if model_name is None: + return None + provider, separator, _ = model_name.partition("/") + if not separator or provider not in _PROVIDER_FILE_RETRIEVE_PROVIDERS: + return None + return cast(FileRetrieveProvider, provider) + + def _has_provider_file_retrieve_route( *, router_credentials: Mapping[str, object] | None, model_name: str | None, ) -> bool: - return bool(router_credentials) or ( - model_name is not None and "/" in model_name - ) + return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None class _ManagedFileRow(Protocol): @@ -733,6 +757,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name: str | None, llm_router: Router | None = None, raise_on_failure: bool = False, + allow_default_provider: bool = False, ) -> tuple[OpenAIFileObject | None, bool]: route_llm_router: Final = llm_router if llm_router is not None else _proxy_llm_router() router_credentials: Final = _provider_file_retrieve_credentials( @@ -744,7 +769,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): router_credentials=router_credentials, model_name=model_name, ) - if not fetch_route_available: + default_provider_route: Final = allow_default_provider and route_llm_router is not None + if not fetch_route_available and not default_provider_route: return None, False try: @@ -767,9 +793,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, ) return ( - provider_file_object_by_model_name.model_copy( - update={"id": unified_file_id} + provider_file_object_by_model_name.model_copy(update={"id": unified_file_id}), + True, + ) + if default_provider_route: + provider_file_object_by_default_route: Final = await asyncio.wait_for( + self._afile_retrieve_with_retries( + provider_file_id=provider_file_id, + call_options=router_credentials or {}, ), + timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, + ) + return ( + provider_file_object_by_default_route.model_copy(update={"id": unified_file_id}), True, ) return None, False @@ -779,6 +815,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"{type(error).__name__} {error}" ) if raise_on_failure: + if isinstance(error, TimeoutError) and not str(error): + raise TimeoutError( + "Provider file retrieve timed out " + f"after {BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS} seconds" + ) from error raise return None, True @@ -811,9 +852,34 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): owner: UserAPIKeyAuth, litellm_parent_otel_span: Span | None, size_bytes: int | None = None, + fetch_provider_details: bool = True, ) -> None: stored_file: Final = await self.get_unified_file_id(unified_file_id, litellm_parent_otel_span) stored_object: Final = stored_file.file_object if stored_file is not None else None + if not fetch_provider_details: + if stored_file is not None: + return + router_credentials: Final = _provider_file_retrieve_credentials( + llm_router=_proxy_llm_router(), + model_id=model_id, + ) + file_object_without_provider_details: Final = _batch_output_file_object( + unified_file_id, + provider_file_id, + size_bytes or 0, + fallback=_has_provider_file_retrieve_route( + router_credentials=router_credentials, + model_name=model_name, + ), + ) + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=file_object_without_provider_details, + litellm_parent_otel_span=litellm_parent_otel_span, + model_mappings={model_id: provider_file_id} if model_id else {}, + user_api_key_dict=owner, + ) + return if stored_object is not None and not stored_object.litellm_details_fallback: return @@ -832,8 +898,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) ) 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 + if ( + stored_file is not None + and provider_object is None + and (size_bytes is None or (stored_object is not None and stored_object.bytes == size_bytes)) ): return @@ -852,6 +920,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) ) + if stored_file is not None: + await self._save_refreshed_file_object(stored_file, file_object) + return + await self.store_unified_file_id( file_id=unified_file_id, file_object=file_object, @@ -1002,6 +1074,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict=user_api_key_dict, db_batch_object=row, unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id), + fetch_provider_details=False, ) except Exception as e: verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}") @@ -1874,8 +1947,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): 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" + f"LiteLLM Managed File object with id={file_id} has no file_object and no provider route to fetch it" ) model_id, model_file_id = model_mapping @@ -1887,13 +1959,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name=model_id, llm_router=llm_router, raise_on_failure=True, + allow_default_provider=True, ) 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" + f"LiteLLM Managed File object with id={file_id} has no file_object and no provider route to fetch it" ) if response is None: raise ValueError("Provider file details could not be retrieved") diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 284063a5ed5..bd4f8ca89e7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -125,6 +125,7 @@ class ManagedBatchOutputFileWriter(Protocol): owner: "UserAPIKeyAuth", litellm_parent_otel_span: "Span | None", size_bytes: int | None = None, + fetch_provider_details: bool = True, ) -> None: ... @@ -1313,6 +1314,8 @@ async def ensure_batch_response_managed_file_ids( user_api_key_dict: "UserAPIKeyAuth | None" = None, db_batch_object: object | None = None, unified_batch_id: str | Literal[False] | None = None, + *, + fetch_provider_details: bool = True, ) -> None: """Normalize batch file IDs to managed unified IDs before DB persistence.""" await resolve_input_file_id_to_unified(response, prisma_client) @@ -1352,6 +1355,7 @@ async def ensure_batch_response_managed_file_ids( owner=effective_auth, litellm_parent_otel_span=effective_auth.parent_otel_span, size_bytes=None, + fetch_provider_details=fetch_provider_details, ) setattr(response, file_attr, new_unified_file_id) verbose_proxy_logger.debug("Converted batch %s %r to managed ID before DB write", file_attr, raw_file_id) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 9a00f602167..10998a11d42 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -3,13 +3,16 @@ import json import time import asyncio from collections.abc import Awaitable, Callable, Mapping +from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException -from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles +from openai import APIConnectionError +from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles from litellm.caching import DualCache from litellm.proxy._types import CallTypes, LiteLLM_ManagedFileTable from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -37,7 +40,18 @@ class _InMemoryManagedFileTable: self, where: Mapping[str, object] ) -> LiteLLM_ManagedFileTable | None: self.find_first_calls.append(where) - return self.rows.get(cast(str, where["unified_file_id"])) + unified_file_id: Final = where.get("unified_file_id") + if isinstance(unified_file_id, str): + return self.rows.get(unified_file_id) + raw_file_filter: Final = where.get("flat_model_file_ids") + if isinstance(raw_file_filter, Mapping): + raw_file_id: Final = raw_file_filter.get("has") + if isinstance(raw_file_id, str): + return next( + (row for row in self.rows.values() if raw_file_id in row.flat_model_file_ids), + None, + ) + return None async def upsert( self, @@ -176,6 +190,46 @@ def _marked_fallback_file_row( ) +def _batch_response_for_output_listing() -> "LiteLLMBatch": + from openai.types.batch import BatchRequestCounts + from litellm.types.utils import LiteLLMBatch + + response: Final = LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="input-file", + object="batch", + status="completed", + output_file_id="provider-output", + request_counts=BatchRequestCounts(completed=1, failed=0, total=1), + ) + response.hidden_params = { + "model_id": "model-123", + "model_name": "bedrock/model-x", + } + return response + + +async def _resolve_batch_for_output_listing( + managed_files: PROXY_LiteLLMManagedFiles, + response: "LiteLLMBatch", +) -> "LiteLLMBatch | None": + from litellm.proxy._types import UserAPIKeyAuth + + return await managed_files._resolve_listed_batch( + row=SimpleNamespace( + unified_object_id="batch-row", + created_by="user-123", + team_id="team-123", + ), + batch_obj=response, + unified_id_by_raw_id={}, + user_api_key_dict=UserAPIKeyAuth(user_id="user-123", team_id="team-123"), + ) + + def test_get_file_ids_from_messages(): proxy_managed_files = PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=MagicMock() @@ -2059,6 +2113,272 @@ async def test_store_batch_output_file_falls_back_without_model_id(): assert "litellm_details_fallback" not in stored_object.model_dump() +@pytest.mark.asyncio +async def test_listed_batch_saves_basic_output_without_provider_fetch(): + import litellm.proxy.proxy_server as proxy_server_module + + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma() + response: Final = _batch_response_for_output_listing() + with ( + patch.object(proxy_server_module, "llm_router", None), + patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, + ): + resolved: Final = await _resolve_batch_for_output_listing( + proxy_managed_files, response + ) + + assert resolved is response + retrieve.assert_not_called() + stored_file: Final = next(iter(managed_file_table.rows.values())) + assert stored_file.file_object is not None + assert stored_file.file_object.bytes == 0 + assert stored_file.file_object.litellm_details_fallback is True + assert stored_file.model_mappings == {"model-123": "provider-output"} + + +@pytest.mark.asyncio +async def test_listed_batch_leaves_existing_marked_output_untouched(): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + + proxy_managed_files, _ = _managed_files_with_fake_prisma() + unified_file_id: Final = proxy_managed_files.get_unified_output_file_id( + output_file_id="provider-output", + model_id="model-123", + model_name="bedrock/model-x", + ) + row: Final = _marked_fallback_file_row().model_copy( + update={"unified_file_id": unified_file_id} + ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row) + response: Final = _batch_response_for_output_listing() + with ( + patch.object(proxy_server_module, "llm_router", None), + patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, + ): + resolved: Final = await _resolve_batch_for_output_listing( + proxy_managed_files, response + ) + await proxy_managed_files.store_batch_output_file( + unified_file_id=unified_file_id, + provider_file_id="provider-output", + model_id="model-123", + model_name="bedrock/model-x", + owner=UserAPIKeyAuth(user_id="user-123", team_id="team-123"), + litellm_parent_otel_span=None, + fetch_provider_details=False, + ) + + assert resolved is response + retrieve.assert_not_called() + assert managed_file_table.rows[unified_file_id] is row + assert managed_file_table.upsert_calls == [] + + +@pytest.mark.asyncio +async def test_batch_retrieve_registration_fetches_details_by_default(): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.openai_files_endpoints.common_utils import ( + ensure_batch_response_managed_file_ids, + ) + + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma() + response: Final = _batch_response_for_output_listing() + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = { + "api_key": "key", + "custom_llm_provider": "bedrock", + } + provider_file_object: Final = OpenAIFileObject( + id="provider-output", + object="file", + bytes=836, + created_at=456, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + with ( + patch.object(proxy_server_module, "llm_router", router), + patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + return_value=provider_file_object, + ) as retrieve, + ): + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=proxy_managed_files, + prisma_client=proxy_managed_files.prisma_client, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-123"), + db_batch_object=SimpleNamespace(created_by="user-123", team_id=None), + ) + + retrieve.assert_awaited_once() + output_file_id: Final = response.output_file_id + assert output_file_id is not None + stored_file: Final = managed_file_table.rows[output_file_id] + assert stored_file.file_object is not None + assert stored_file.file_object.bytes == 836 + + +@pytest.mark.parametrize( + ("error_factory", "expected_attempts"), + [ + pytest.param(lambda: HTTPException(status_code=429), 2, id="429"), + pytest.param(lambda: HTTPException(status_code=408), 2, id="408"), + pytest.param(lambda: HTTPException(status_code=503), 2, id="503"), + pytest.param( + lambda: httpx.ConnectError( + "connection failed", + request=httpx.Request("GET", "https://api.openai.com/v1/files/file-1"), + ), + 2, + id="connect-error", + ), + pytest.param( + lambda: APIConnectionError( + message="connection failed", + request=httpx.Request("GET", "https://api.openai.com/v1/files/file-1"), + ), + 2, + id="openai-api-connection-error", + ), + pytest.param(lambda: asyncio.TimeoutError(), 2, id="async-timeout"), + pytest.param(lambda: HTTPException(status_code=400), 1, id="400"), + pytest.param(lambda: HTTPException(status_code=403), 1, id="403"), + pytest.param(lambda: HTTPException(status_code=404), 1, id="404"), + pytest.param(lambda: ValueError("invalid file"), 1, id="value-error"), + ], +) +@pytest.mark.asyncio +async def test_batch_file_detail_retry_classification( + error_factory: Callable[[], Exception], + expected_attempts: int, +) -> None: + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + + sleep_calls: Final[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 + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + provider_file_object: Final = OpenAIFileObject( + id="provider-output", + object="file", + bytes=836, + created_at=456, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + retrieve_side_effect: Final = ( + [error_factory(), provider_file_object] + if expected_attempts == 2 + else error_factory() + ) + with ( + patch.object(proxy_server_module, "llm_router", router), + patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + side_effect=retrieve_side_effect, + ) as retrieve, + ): + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-output", + provider_file_id="provider-output", + model_id="model-123", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + assert retrieve.await_count == expected_attempts + assert [call.kwargs["max_retries"] for call in retrieve.call_args_list] == [ + 0 + ] * expected_attempts + assert sleep_calls == ([0.5] if expected_attempts == 2 else []) + stored_object: Final = managed_file_table.rows["unified-output"].file_object + assert stored_object is not None + assert stored_object.bytes == (836 if expected_attempts == 2 else 0) + + +@pytest.mark.asyncio +async def test_store_batch_output_file_ignores_unmarked_row_with_provider_route(): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import OpenAIFileObject + + existing_file_object: Final = OpenAIFileObject( + id="unified-output", + object="file", + bytes=1, + created_at=1, + filename="output.jsonl", + purpose="batch_output", + ) + existing_file_row: Final = LiteLLM_ManagedFileTable( + unified_file_id="unified-output", + file_object=existing_file_object, + model_mappings={"model-123": "provider-output"}, + flat_model_file_ids=["provider-output"], + created_by="user-123", + ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( + existing_file_row + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + with ( + patch.object(proxy_server_module, "llm_router", router), + patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, + ): + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-output", + provider_file_id="provider-output", + model_id="model-123", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + retrieve.assert_not_called() + assert managed_file_table.rows["unified-output"].file_object is existing_file_object + assert managed_file_table.upsert_calls == [] + + +@pytest.mark.asyncio +async def test_store_batch_output_file_rejects_non_provider_model_prefix(): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma() + with ( + patch.object(proxy_server_module, "llm_router", None), + patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, + ): + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-output", + provider_file_id="provider-output", + model_id="model-123", + model_name="my-team/gpt-4o", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + retrieve.assert_not_called() + stored_object: Final = managed_file_table.rows["unified-output"].file_object + assert stored_object is not None + assert stored_object.litellm_details_fallback is None + + @pytest.mark.asyncio async def test_store_batch_output_file_marks_fallback_after_four_transient_failures(): import litellm.proxy.proxy_server as proxy_server_module @@ -2207,7 +2527,67 @@ async def test_store_batch_output_file_refreshes_marked_fallback_on_later_write( assert refreshed.id == "unified-output" assert refreshed.bytes == 836 assert refreshed.litellm_details_fallback is None - assert managed_file_table.upsert_calls[0][1]["update"]["file_object"] == refreshed.model_dump_json() + assert managed_file_table.update_many_calls == [ + ({"unified_file_id": "unified-output"}, {"file_object": refreshed.model_dump_json()}) + ] + assert managed_file_table.upsert_calls == [] + + +@pytest.mark.asyncio +async def test_store_batch_output_file_does_not_recreate_deleted_row_after_refresh(): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + + 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( + _marked_fallback_file_row() + ) + router: Final = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + + async def delete_then_return_file(**_kwargs: object) -> OpenAIFileObject: + await managed_file_table.delete(where={"unified_file_id": "unified-output"}) + return provider_object + + with ( + patch.object(proxy_server_module, "llm_router", router), + patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + side_effect=delete_then_return_file, + ), + ): + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-output", + provider_file_id="provider-output", + model_id="model-123", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + assert "unified-output" not in managed_file_table.rows + assert managed_file_table.upsert_calls == [] + assert len(managed_file_table.update_many_calls) == 1 + saved_file_object_json: Final = cast( + str, managed_file_table.update_many_calls[0][1]["file_object"] + ) + saved_file_object: Final = json.loads(saved_file_object_json) + assert saved_file_object["id"] == "unified-output" + assert ( + await proxy_managed_files.internal_usage_cache.async_get_cache( + key="unified-output", + litellm_parent_otel_span=None, + ) + is None + ) @pytest.mark.asyncio @@ -2427,7 +2807,7 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none() llm_router=None, ) - assert "llm_router is required" in str(exc_info.value) + assert "no provider route to fetch it" in str(exc_info.value) @pytest.mark.asyncio @@ -2647,7 +3027,7 @@ async def test_afile_retrieve_case3_includes_provider_error_text(): created_by="user-123", ) router: Final = MagicMock() - router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + router.get_deployment_credentials_with_provider.return_value = None proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) with ( @@ -2674,6 +3054,113 @@ async def test_afile_retrieve_case3_includes_provider_error_text(): ) +@pytest.mark.asyncio +async def test_afile_retrieve_case3_uses_default_provider_for_unknown_router_deployment(): + 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 = None + proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) + provider_file_object: Final = OpenAIFileObject( + id="provider-output", + object="file", + bytes=836, + created_at=456, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + + with patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + return_value=provider_file_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() + assert retrieve.await_args.kwargs["file_id"] == "provider-output" + assert retrieve.await_args.kwargs["max_retries"] == 0 + assert "custom_llm_provider" not in retrieve.await_args.kwargs + assert "_litellm_internal_model_credentials" not in retrieve.await_args.kwargs + + +@pytest.mark.asyncio +async def test_afile_retrieve_case3_times_out_on_default_provider_route(): + 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 = None + proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) + + async def never_answers(**_options: object) -> None: + await asyncio.Event().wait() + + with ( + patch( + "litellm.afile_retrieve", + new_callable=AsyncMock, + side_effect=never_answers, + ) as retrieve, + patch( + "litellm_enterprise.proxy.hooks.managed_files.BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS", + 0.01, + ), + pytest.raises(Exception, match="Provider file retrieve timed out"), + ): + await asyncio.wait_for( + proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=router, + ), + timeout=1, + ) + + retrieve.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_afile_retrieve_case3_without_route_raises_accurate_error(): + import litellm.proxy.proxy_server as proxy_server_module + + stored_file: Final = LiteLLM_ManagedFileTable( + unified_file_id="unified-output", + file_object=None, + model_mappings={"my-team/gpt-4o": "provider-output"}, + flat_model_file_ids=["provider-output"], + created_by="user-123", + ) + proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) + with ( + patch.object(proxy_server_module, "llm_router", None), + patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, + pytest.raises(Exception, match="no provider route to fetch it"), + ): + await proxy_managed_files.afile_retrieve( + file_id="unified-output", + litellm_parent_otel_span=None, + llm_router=None, + ) + + retrieve.assert_not_called() + + @pytest.mark.asyncio async def test_afile_retrieve_returns_marked_fallback_without_marker_when_refresh_fails(): router = MagicMock() diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 094356108e0..58d6a44e4a4 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1063,6 +1063,7 @@ class TestCheckBatchCost: owner=mock_hook.store_batch_output_file.await_args.kwargs["owner"], litellm_parent_otel_span=None, size_bytes=None, + fetch_provider_details=True, ) owner = mock_hook.store_batch_output_file.await_args.kwargs["owner"] assert owner.user_id == "user-1" @@ -1643,6 +1644,10 @@ class TestCheckBatchCost: model_name="gpt-5-batch", ) assert mock_hook.store_batch_output_file.await_count == 2 + assert { + call.kwargs["model_name"] + for call in mock_hook.store_batch_output_file.await_args_list + } == {"gpt-5-batch"} stored = { c.kwargs["provider_file_id"]: c.kwargs["unified_file_id"] for c in mock_hook.store_batch_output_file.call_args_list