scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(proxy): keep batch listing provider-free

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-10-08 16:23:57 +00:00
parent 065d38181a
commit 9721d9f82b
4 changed files with 585 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

96
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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")

4
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -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)

497
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

@ -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()

5
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -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