diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 10075b73d1d..7d1b1966cea 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/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, diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d40c0cb6d8b..d9d2929a8d9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/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, ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 40e543c6903..c4aa84004dc 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/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, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 9d10d37d924..284063a5ed5 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/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, diff --git a/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 2049ff58f95..cf305c660cd 100644 --- a/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/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" diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 3ea3b97e1fe..7bf53b11cdb 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/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( 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 5b1a5052343..3d844e8b9e8 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/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,