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 4e14b56faff..10075b73d1d 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,7 +2,6 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ -import time from collections.abc import Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -20,7 +19,6 @@ from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository -from litellm.types.llms.openai import OpenAIFileObject if TYPE_CHECKING: from prisma import models as prisma_models @@ -49,19 +47,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( ) -def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes: int) -> OpenAIFileObject: - filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id - return OpenAIFileObject( - id=unified_file_id, - object="file", - purpose="batch_output", - filename=filename, - created_at=int(time.time()), - bytes=size_bytes, - status="processed", - ) - - class _ManagedObjectRow(Protocol): @property def id(self) -> str: ... @@ -830,15 +815,14 @@ class CheckBatchCost: custom_llm_provider=custom_llm_provider, ) - # CheckBatchCost bypasses async_post_call_success_hook, so convert raw - # output/error file IDs to managed base64 IDs before the DB write here. - managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None: + from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter + + managed_files_hook: Final = self.proxy_logging_obj.get_proxy_hook("managed_files") + if isinstance(managed_files_hook, ManagedBatchOutputFileWriter): + managed_file_writer: Final = managed_files_hook from litellm.proxy._types import UserAPIKeyAuth - managed_file_model_name = self._get_managed_file_model_name( - job=job, deployment_info=deployment_info - ) + managed_file_model_name = self._get_managed_file_model_name(job=job, deployment_info=deployment_info) _minimal_auth = UserAPIKeyAuth( user_id=job.created_by or "default-user-id", team_id=getattr(job, "team_id", None), @@ -847,21 +831,18 @@ class CheckBatchCost: _raw_file_id = cast(str | None, getattr(response, _file_attr, None)) if _raw_file_id and not _is_base64_encoded_unified_file_id(_raw_file_id): try: - _unified_file_id = managed_files_hook.get_unified_output_file_id( + _unified_file_id = managed_file_writer.get_unified_output_file_id( output_file_id=_raw_file_id, model_id=model_id, model_name=managed_file_model_name, ) - await managed_files_hook.store_unified_file_id( - file_id=_unified_file_id, - file_object=_batch_output_file_object( - unified_file_id=_unified_file_id, - raw_file_id=_raw_file_id, - size_bytes=len(content_bytes or b"") if _file_attr == "output_file_id" else 0, - ), + await managed_file_writer.store_batch_output_file( + unified_file_id=_unified_file_id, + provider_file_id=_raw_file_id, + model_id=model_id, + owner=_minimal_auth, litellm_parent_otel_span=None, - model_mappings={model_id: _raw_file_id}, - user_api_key_dict=_minimal_auth, + size_bytes=len(content_bytes) if _file_attr == "output_file_id" else None, ) setattr(response, _file_attr, _unified_file_id) verbose_proxy_logger.info( diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index c07bb216602..ee6f8156c09 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,6 +3,7 @@ import base64 import json +import time from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType from typing import ( @@ -15,6 +16,7 @@ from typing import ( Optional, Protocol, TypedDict, + Unpack, Union, cast, ) @@ -31,7 +33,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, FileRetrieveProvider +from litellm.files.types import FileRetrieveCallOptions 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 ( @@ -96,6 +98,17 @@ from litellm.types.utils import ( SpecialEnums, ) + +class _ManagedFileRetrieve(Protocol): + async def __call__( + self, + *, + file_id: str, + _litellm_internal_model_credentials: Mapping[str, object], + **kwargs: Unpack[FileRetrieveCallOptions], + ) -> OpenAIFileObject: ... + + if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from prisma.models import ( @@ -146,6 +159,19 @@ def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> return None +def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes: int) -> OpenAIFileObject: + filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id + return OpenAIFileObject( + id=unified_file_id, + object="file", + purpose="batch_output", + filename=filename, + created_at=int(time.time()), + bytes=size_bytes, + status="processed", + ) + + class _ManagedFileRow(Protocol): unified_file_id: str file_object: OpenAIFileObject @@ -595,14 +621,74 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for file_id in provider_file_ids: model_name = decode_model_from_file_id(file_id) raw_file_id = get_original_file_id(file_id) - await self.store_unified_file_id( - file_id=file_id, - file_object=None, + await self.store_batch_output_file( + unified_file_id=file_id, + provider_file_id=raw_file_id, + model_id=model_name or None, litellm_parent_otel_span=litellm_parent_otel_span, - model_mappings={model_name: raw_file_id} if model_name else {}, - user_api_key_dict=owner_identity, + owner=owner_identity, ) + async def _retrieve_batch_output_file_object( + self, + *, + unified_file_id: str, + provider_file_id: str, + model_id: str | None, + size_bytes: int | None, + ) -> OpenAIFileObject: + if model_id 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}) + 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) + + async def store_batch_output_file( + self, + *, + unified_file_id: str, + provider_file_id: str, + model_id: str | None, + owner: UserAPIKeyAuth, + litellm_parent_otel_span: Span | None, + size_bytes: int | None = None, + ) -> None: + stored_file: Final = await self.get_unified_file_id(unified_file_id, litellm_parent_otel_span) + if stored_file is not None and stored_file.file_object is not None: + return + + file_object: Final = await self._retrieve_batch_output_file_object( + unified_file_id=unified_file_id, + provider_file_id=provider_file_id, + model_id=model_id, + size_bytes=size_bytes, + ) + + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=file_object, + litellm_parent_otel_span=litellm_parent_otel_span, + model_mappings={model_id: provider_file_id} if model_id else {}, + user_api_key_dict=owner, + ) + async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, @@ -1474,42 +1560,12 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) setattr(response, file_attr, unified_file_id) - # Use llm_router credentials when available. Without credentials, - # Azure and other auth-required providers return 500/401. - file_object = None - try: - # Import module and use getattr for better testability with mocks - import litellm.proxy.proxy_server as proxy_server_module - - _llm_router = getattr(proxy_server_module, "llm_router", None) - if _llm_router is not None and model_id: - _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} - file_object = await litellm.afile_retrieve( - file_id=provider_file_id, - **cast(FileRetrieveCallOptions, _creds), - ) - else: - file_object = await litellm.afile_retrieve( - custom_llm_provider=cast( - FileRetrieveProvider, - model_name.split("/")[0] if model_name and "/" in model_name else "openai", - ), - file_id=provider_file_id, - ) - verbose_logger.debug( - f"Successfully retrieved file object for {file_attr}={provider_file_id}" - ) - except Exception as e: - verbose_logger.warning( - f"Failed to retrieve file object for {file_attr}={provider_file_id}: {str(e)}. Storing with None and will fetch on-demand." - ) - - await self.store_unified_file_id( - file_id=unified_file_id, - file_object=file_object, + await self.store_batch_output_file( + unified_file_id=unified_file_id, + provider_file_id=provider_file_id, + model_id=model_id, + owner=user_api_key_dict, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_mappings={model_id: provider_file_id}, - user_api_key_dict=user_api_key_dict, ) request_metadata: Final = data.get("litellm_metadata") await self.store_unified_object_id( @@ -1639,6 +1695,7 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id) or {} response = await litellm.afile_retrieve( file_id=model_file_id, + _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), **cast(FileRetrieveCallOptions, credentials), ) response.id = file_id # Replace with unified ID diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index e7ee2f89bbf..40e543c6903 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -8,9 +8,10 @@ from collections.abc import Iterable, Mapping, MutableMapping, Sequence from contextlib import suppress from dataclasses import dataclass from datetime import datetime +from email.utils import parsedate_to_datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final, Literal, TypeAlias, TypedDict +from typing import Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import quote, unquote, urlencode import httpx @@ -70,12 +71,52 @@ from ..common_utils import ( ) S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" +S3_RETRIEVE_FILE_ID_PARAM: Final = "_s3_retrieve_file_id" +S3_RETRIEVE_FILE_KEY_PARAM: Final = "_s3_retrieve_file_key" LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose" 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", "")) + range_parts: Final = content_range.removeprefix("bytes 0-0/") + if content_range.startswith("bytes 0-0/") and range_parts.isdigit(): + return int(range_parts) + raise BedrockError( + status_code=status_code, + message=f"Invalid S3 Content-Range header: {content_range}", + headers=raw_response.headers, + response=raw_response, + ) + if status_code == 200: + content_length: Final = cast(str, raw_response.headers.get("Content-Length", "")) + if content_length.isdigit(): + return int(content_length) + raise BedrockError( + status_code=status_code, + message=f"Invalid S3 Content-Length header: {content_length}", + headers=raw_response.headers, + response=raw_response, + ) + if status_code >= 400: + raise BedrockError( + status_code=status_code, + message=raw_response.text, + headers=raw_response.headers, + response=raw_response, + ) + raise BedrockError( + status_code=status_code, + message=f"S3 file retrieval returned HTTP {status_code}", + headers=raw_response.headers, + response=raw_response, + ) + + class _S3DeleteContext(LiteLLMBaseModel): file_id: str = Field(min_length=1) @@ -1276,18 +1317,52 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_retrieve_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file retrieval") + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + _, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params) + url, params = self._transform_s3_file_request( + file_id=file_id, + method="GET", + optional_params=optional_params, + litellm_params=litellm_params, + ) + signed_headers_object: Final = litellm_params.get(S3_SIGNED_REQUEST_HEADERS_PARAM) + if not isinstance(signed_headers_object, Mapping): + raise ValueError("S3 request signing did not produce request headers") + signed_headers: Final = cast(Mapping[str, str], signed_headers_object) + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = MappingProxyType({**signed_headers, "Range": "bytes=0-0"}) + litellm_params[S3_RETRIEVE_FILE_ID_PARAM] = file_id + litellm_params[S3_RETRIEVE_FILE_KEY_PARAM] = object_key + return url, params def transform_retrieve_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: - raise NotImplementedError("BedrockFilesConfig does not support file retrieval") + file_id: Final = litellm_params.get(S3_RETRIEVE_FILE_ID_PARAM) + object_key: Final = litellm_params.get(S3_RETRIEVE_FILE_KEY_PARAM) + if not isinstance(file_id, str) or not isinstance(object_key, str): + raise ValueError("S3 retrieve response is missing request context") + + file_size: Final = _retrieved_s3_file_size(raw_response) + + 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", "")) + created_at: Final = int(parsedate_to_datetime(last_modified).timestamp()) if last_modified else 0 + return OpenAIFileObject( + id=file_id, + bytes=file_size, + created_at=created_at, + filename=posixpath.basename(object_key), + object="file", + purpose="batch_output" if relative_key.startswith(BEDROCK_MANAGED_S3_OUTPUT_PREFIX) else "batch", + status="processed", + ) def transform_delete_file_request( self, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 7e7ab1dc1f0..42ce8d4cfe7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,4 +1,5 @@ import base64 +import logging import mimetypes import re from collections.abc import Mapping, Sequence @@ -27,6 +28,7 @@ from litellm.types.utils import SpecialEnums if TYPE_CHECKING: from fastapi import Request + from opentelemetry.trace import Span from prisma.models import LiteLLM_ManagedObjectTable from litellm.proxy._types import UserAPIKeyAuth @@ -104,6 +106,27 @@ class ManagedFileIdResolver(Protocol): ) -> Mapping[str, str]: ... +@runtime_checkable +class ManagedBatchOutputFileWriter(Protocol): + def get_unified_output_file_id( + self, + output_file_id: str, + model_id: str, + model_name: str | None = None, + ) -> str: ... + + async def store_batch_output_file( + self, + *, + unified_file_id: str, + provider_file_id: str, + model_id: str | None, + owner: "UserAPIKeyAuth", + litellm_parent_otel_span: "Span | None", + size_bytes: int | None = None, + ) -> None: ... + + def _is_base64_encoded_unified_file_id(b64_uid: object) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -1279,11 +1302,11 @@ def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mappi async def ensure_batch_response_managed_file_ids( - response, - managed_files_obj, - prisma_client, - verbose_proxy_logger, - user_api_key_dict=None, + response: "LiteLLMBatch", + managed_files_obj: object | None, + prisma_client: "PrismaClient | None", + verbose_proxy_logger: logging.Logger, + user_api_key_dict: "UserAPIKeyAuth | None" = None, db_batch_object: object | None = None, unified_batch_id: str | Literal[False] | None = None, ) -> None: @@ -1291,7 +1314,7 @@ async def ensure_batch_response_managed_file_ids( await resolve_input_file_id_to_unified(response, prisma_client) await resolve_output_file_ids_to_unified(response, prisma_client) - if managed_files_obj is None: + if not isinstance(managed_files_obj, ManagedBatchOutputFileWriter): return model_id: Final = _model_id_for_batch_response(response, unified_batch_id) @@ -1305,8 +1328,10 @@ async def ensure_batch_response_managed_file_ids( if effective_auth is None: return - for file_attr in ("output_file_id", "error_file_id"): - raw_file_id = getattr(response, file_attr, None) + for file_attr, raw_file_id in ( + ("output_file_id", response.output_file_id), + ("error_file_id", response.error_file_id), + ): if not raw_file_id or _is_base64_encoded_unified_file_id(raw_file_id): continue try: @@ -1315,12 +1340,13 @@ async def ensure_batch_response_managed_file_ids( model_id=model_id, model_name=model_name, ) - await managed_files_obj.store_unified_file_id( - file_id=new_unified_file_id, - file_object=None, - litellm_parent_otel_span=getattr(effective_auth, "parent_otel_span", None), - model_mappings={model_id: raw_file_id}, - user_api_key_dict=effective_auth, + await managed_files_obj.store_batch_output_file( + unified_file_id=new_unified_file_id, + provider_file_id=raw_file_id, + model_id=model_id, + owner=effective_auth, + litellm_parent_otel_span=effective_auth.parent_otel_span, + size_bytes=None, ) 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 754173a8c37..f1b31ba1c63 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -525,6 +525,8 @@ async def test_output_file_id_for_batch_retrieve(): proxy_managed_files = PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=AsyncMock() ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) + proxy_managed_files.store_unified_file_id = AsyncMock() response = await proxy_managed_files.async_post_call_success_hook( data={}, @@ -585,6 +587,8 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi proxy_managed_files = PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=AsyncMock() ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) + proxy_managed_files.store_unified_file_id = AsyncMock() provider_output_file = OpenAIFileObject( id="file-provider-output-id", @@ -661,6 +665,8 @@ async def test_error_file_id_for_failed_batch(): proxy_managed_files = PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=AsyncMock() ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) + proxy_managed_files.store_unified_file_id = AsyncMock() # Create a proper OpenAIFileObject for the error file error_file_object = OpenAIFileObject( @@ -1514,6 +1520,145 @@ async def test_store_unified_file_id_updates_file_metadata_on_existing_row(): assert second_update["storage_url"] == "s3://bucket/output.jsonl" +@pytest.mark.asyncio +async def test_store_batch_output_file_skips_existing_file_object(): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import OpenAIFileObject + + proxy_managed_files = PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=AsyncMock(), + ) + proxy_managed_files.get_unified_file_id = AsyncMock( + return_value=MagicMock( + file_object=OpenAIFileObject( + id="existing", + object="file", + bytes=1, + created_at=1, + filename="output.jsonl", + purpose="batch_output", + ) + ) + ) + proxy_managed_files.store_unified_file_id = AsyncMock() + + with 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() + proxy_managed_files.store_unified_file_id.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_store_batch_output_file_stores_provider_object_with_unified_id(): + import litellm.proxy.proxy_server as proxy_server_module + 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", + ) + router = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + proxy_managed_files = PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=AsyncMock(), + ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None)) + proxy_managed_files.store_unified_file_id = AsyncMock() + + with ( + patch.object(proxy_server_module, "llm_router", router), + patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=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="model-123", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + retrieve.assert_awaited_once() + stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"] + assert stored_object.id == "unified-output" + assert stored_object.bytes == 123 + + +@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 + from litellm.proxy._types import UserAPIKeyAuth + + router = MagicMock() + router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"} + proxy_managed_files = PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=AsyncMock(), + ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None)) + proxy_managed_files.store_unified_file_id = AsyncMock() + + with ( + patch.object(proxy_server_module, "llm_router", router), + patch("litellm.afile_retrieve", new_callable=AsyncMock, side_effect=RuntimeError("provider unavailable")), + ): + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-output", + provider_file_id="s3://bucket/output.jsonl", + model_id="model-123", + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + size_bytes=321, + ) + + stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"] + assert stored_object.id == "unified-output" + assert stored_object.filename == "output.jsonl" + assert stored_object.bytes == 321 + assert stored_object.purpose == "batch_output" + + +@pytest.mark.asyncio +async def test_store_batch_output_file_falls_back_without_model_id(): + from litellm.proxy._types import UserAPIKeyAuth + + proxy_managed_files = PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=AsyncMock(), + ) + proxy_managed_files.get_unified_file_id = AsyncMock(return_value=MagicMock(file_object=None)) + proxy_managed_files.store_unified_file_id = AsyncMock() + + with patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve: + await proxy_managed_files.store_batch_output_file( + unified_file_id="unified-error", + provider_file_id="error.jsonl", + model_id=None, + owner=UserAPIKeyAuth(user_id="user-123"), + litellm_parent_otel_span=None, + ) + + retrieve.assert_not_called() + stored_object = proxy_managed_files.store_unified_file_id.await_args.kwargs["file_object"] + assert stored_object.id == "unified-error" + assert stored_object.filename == "error.jsonl" + assert stored_object.bytes == 0 + assert stored_object.purpose == "batch_output" + + @pytest.mark.asyncio async def test_afile_delete_returns_provider_response_when_stored_file_object_none(): """ @@ -1956,7 +2101,7 @@ async def test_afile_list_returns_persisted_batch_output_file(): prisma_client = MagicMock() prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row]) - proxy_managed_files = _PROXY_LiteLLMManagedFiles( + proxy_managed_files = PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) @@ -3491,8 +3636,10 @@ async def test_post_call_batch_sync_does_not_claim_ownership(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.update_many.return_value = 0 proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()), + prisma_client=prisma_client, ) + prisma_client.db.litellm_managedfiletable.find_first.return_value = None await proxy_managed_files.async_post_call_success_hook( data={"batch_id": MODEL_ENCODED_BATCH_ID}, @@ -3515,22 +3662,16 @@ async def test_post_call_batch_sync_updates_existing_row(): prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( _owned_record(created_by="user_a", team_id="team_a") ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) await proxy_managed_files.async_post_call_success_hook( data={"batch_id": MODEL_ENCODED_BATCH_ID}, - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), response=_batch_response(MODEL_ENCODED_BATCH_ID), ) update_call = prisma_client.db.litellm_managedobjecttable.update_many.await_args - assert update_call.kwargs["where"] == { - "unified_object_id": MODEL_ENCODED_BATCH_ID - } + assert update_call.kwargs["where"] == {"unified_object_id": MODEL_ENCODED_BATCH_ID} assert update_call.kwargs["data"]["status"] == "completed" prisma_client.db.litellm_managedobjecttable.upsert.assert_not_awaited() @@ -3550,8 +3691,10 @@ async def test_post_call_batch_sync_stores_output_file_ownership_from_batch_row( _owned_record(created_by="user_a", team_id="team_a") ) proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()), + prisma_client=prisma_client, ) + prisma_client.db.litellm_managedfiletable.find_first.return_value = None await proxy_managed_files.async_post_call_success_hook( data={"batch_id": MODEL_ENCODED_BATCH_ID}, diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 4bfac61545d..a647b94b6e8 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2216,6 +2216,139 @@ class TestBedrockFileContentTransformation: assert "x-amz-content-sha256" in authorization assert "X-Amz-Date" in signed_headers + def test_transform_retrieve_file_request_adds_unsigned_range(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + url, params = BedrockFilesConfig().transform_retrieve_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.EXPECTED_URL + assert params == {} + assert litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Range"] == "bytes=0-0" + + @pytest.mark.parametrize( + ("file_id", "purpose"), + [ + ("s3://my-bucket/litellm-batch-outputs/job-123/output.jsonl", "batch_output"), + ("s3://my-bucket/litellm-bedrock-files-job-123/input.jsonl", "batch"), + ], + ) + def test_transform_retrieve_file_response_parses_metadata( + self, file_id: str, purpose: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + BedrockFilesConfig().transform_retrieve_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) + response = BedrockFilesConfig().transform_retrieve_file_response( + raw_response=httpx.Response( + 206, + headers={ + "Content-Range": "bytes 0-0/4321", + "Last-Modified": "Wed, 21 Oct 2015 07:28:00 GMT", + }, + request=httpx.Request("GET", file_id), + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + + assert response.id == file_id + assert response.bytes == 4321 + assert response.created_at == 1445412480 + assert response.filename == file_id.rsplit("/", 1)[-1] + assert response.purpose == purpose + assert response.status == "processed" + assert response.object == "file" + + def test_transform_retrieve_file_response_raises_for_empty_object(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + BedrockFilesConfig().transform_retrieve_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=litellm_params, + ) + with pytest.raises(BedrockError): + BedrockFilesConfig().transform_retrieve_file_response( + raw_response=httpx.Response( + 416, + headers={"Content-Range": "bytes */0"}, + request=httpx.Request("GET", self.S3_URI), + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + + def test_transform_retrieve_file_response_uses_content_length_when_range_is_ignored(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + BedrockFilesConfig().transform_retrieve_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=litellm_params, + ) + response = BedrockFilesConfig().transform_retrieve_file_response( + raw_response=httpx.Response( + 200, + headers={"Content-Length": "4321"}, + request=httpx.Request("GET", self.S3_URI), + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + + assert response.bytes == 4321 + + def test_transform_retrieve_file_response_raises_on_s3_error(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + BedrockFilesConfig().transform_retrieve_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=litellm_params, + ) + + with pytest.raises(BedrockError, match="AccessDenied"): + BedrockFilesConfig().transform_retrieve_file_response( + raw_response=httpx.Response( + 403, + content=b"AccessDenied", + request=httpx.Request("GET", self.S3_URI), + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + def test_transform_file_content_request_decodes_unified_file_id(self, monkeypatch): """Base64 unified ids carrying llm_output_file_id must resolve to their S3 object.""" import base64 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 bf113784720..466147f2f94 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -24,7 +24,7 @@ _CLAIM_OUTPUT_FILE_ID = "file-output-123" def test_batch_output_file_object_derives_metadata(): - from litellm_enterprise.proxy.common_utils.check_batch_cost import _batch_output_file_object + from litellm_enterprise.proxy.hooks.managed_files import _batch_output_file_object output_file = _batch_output_file_object( unified_file_id="unified-output", @@ -1039,9 +1039,11 @@ class TestCheckBatchCost: ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=response) - mock_hook = MagicMock() + from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter + + mock_hook = MagicMock(spec=ManagedBatchOutputFileWriter) mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] - mock_hook.store_unified_file_id = AsyncMock() + mock_hook.store_batch_output_file = AsyncMock() check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook await check_batch_cost_instance.check_batch_cost() @@ -1051,14 +1053,17 @@ class TestCheckBatchCost: model_id="model-123", model_name="gpt-5-batch", ) - stored = { - next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"] - for c in mock_hook.store_unified_file_id.call_args_list - } - assert stored == {raw_error_file_id: unified_error_file_id} - for store_call in mock_hook.store_unified_file_id.call_args_list: - assert store_call.kwargs["user_api_key_dict"].user_id == "user-1" - assert store_call.kwargs["user_api_key_dict"].team_id == "team-1" + mock_hook.store_batch_output_file.assert_awaited_once_with( + unified_file_id=unified_error_file_id, + provider_file_id=raw_error_file_id, + model_id="model-123", + owner=mock_hook.store_batch_output_file.await_args.kwargs["owner"], + litellm_parent_otel_span=None, + size_bytes=None, + ) + owner = mock_hook.store_batch_output_file.await_args.kwargs["owner"] + assert owner.user_id == "user-1" + assert owner.team_id == "team-1" assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args @@ -1529,6 +1534,8 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ + from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1563,12 +1570,12 @@ class TestCheckBatchCost: mock_deployment.model_info.model_dump.return_value = {} mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) - mock_hook = MagicMock() + mock_hook = MagicMock(spec=ManagedBatchOutputFileWriter) mock_hook.get_unified_output_file_id.side_effect = [ fake_managed_output_id, fake_managed_error_id, ] - mock_hook.store_unified_file_id = AsyncMock() + mock_hook.store_batch_output_file = AsyncMock() check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() @@ -1631,28 +1638,21 @@ class TestCheckBatchCost: model_id="model-123", model_name="gpt-5-batch", ) - assert mock_hook.store_unified_file_id.await_count == 2 - # {raw_file_id: managed_file_id} for each store call + assert mock_hook.store_batch_output_file.await_count == 2 stored = { - next(iter(c[1]["model_mappings"].values())): c[1]["file_id"] - for c in mock_hook.store_unified_file_id.call_args_list + c.kwargs["provider_file_id"]: c.kwargs["unified_file_id"] + for c in mock_hook.store_batch_output_file.call_args_list } assert stored == { raw_output_file_id: fake_managed_output_id, raw_error_file_id: fake_managed_error_id, } - for store_call in mock_hook.store_unified_file_id.call_args_list: - file_object = store_call.kwargs["file_object"] - assert file_object is not None - assert file_object.object == "file" - assert file_object.purpose == "batch_output" - assert file_object.id == store_call.kwargs["file_id"] - stored_file_objects = { - next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_object"] - for c in mock_hook.store_unified_file_id.call_args_list + stored_sizes = { + c.kwargs["provider_file_id"]: c.kwargs["size_bytes"] + for c in mock_hook.store_batch_output_file.call_args_list } - assert stored_file_objects[raw_output_file_id].bytes == len(mock_file_content.content) - assert stored_file_objects[raw_error_file_id].bytes == 0 + assert stored_sizes[raw_output_file_id] == len(mock_file_content.content) + assert stored_sizes[raw_error_file_id] is None assert mock_response.output_file_id == fake_managed_output_id assert mock_response.error_file_id == fake_managed_error_id @@ -2227,10 +2227,10 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) + + from enterprise.litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles + from litellm.proxy.openai_files_endpoints.common_utils import ManagedBatchOutputFileWriter from litellm.types.utils import LiteLLMBatch - from enterprise.litellm_enterprise.proxy.hooks.managed_files import ( - PROXY_LiteLLMManagedFiles, - ) router = MagicMock() router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) @@ -2241,13 +2241,13 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: deployment.model_info.model_dump.return_value = {} router.get_deployment = MagicMock(return_value=deployment) - hook = MagicMock() + hook = MagicMock(spec=ManagedBatchOutputFileWriter) hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) - hook.store_unified_file_id = AsyncMock() + hook.store_batch_output_file = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.get_proxy_hook.return_value = hook