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 6fe37f0aacb..cfe60a79eed 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple from litellm._logging import verbose_proxy_logger @@ -537,6 +538,7 @@ class CheckBatchCost: credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} _file_content = await afile_content( file_id=raw_output_file_id, + _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), **credentials, ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2432ea3729..1cbed2a68f9 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,6 +7,7 @@ import asyncio import traceback +from types import MappingProxyType from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -706,11 +707,18 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: + deployment_credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model) + trusted_model_credentials: Final = ( + {"_litellm_internal_model_credentials": MappingProxyType(dict(deployment_credentials))} + if deployment_credentials is not None + else {} + ) response = await llm_router.afile_content( **{ "model": model, "file_id": file_id, **data, + **trusted_model_credentials, } ) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fa274324fd6..20390159665 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -312,6 +312,102 @@ class TestCheckBatchCost: ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" + @pytest.mark.asyncio + async def test_output_fetch_passes_deployment_credentials_as_trusted_snapshot( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Bedrock resolves the output bucket ONLY from the immutable snapshot kwarg. + + Spreading the credentials as plain kwargs is not enough: get_litellm_params drops + s3_bucket_name, so without _litellm_internal_model_credentials the cost poller + cannot read the output file and every completed Bedrock batch stays unbilled. + """ + from types import MappingProxyType + from unittest.mock import patch + + 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) + + mock_job = MagicMock() + mock_job.id = "job-bedrock-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "custom_llm_provider": "bedrock", + "s3_bucket_name": "configured-batch-bucket", + "aws_region_name": "us-east-1", + } + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "bedrock" + mock_deployment.litellm_params.model = "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"recordId":"req-1"}' + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"recordId": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("anthropic.claude-haiku-4-5-20251001-v1:0", "bedrock", None, None), + ), + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + mock_afile_content.assert_awaited() + passed_kwargs = mock_afile_content.await_args[1] + snapshot = passed_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "cost poller must pass the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index f27c8dfd2f4..0c26e5f7695 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3149,6 +3149,96 @@ def test_require_managed_files_rejects_raw_provider_file_id( mock_call.assert_not_called() +def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkeypatch): + """A managed batch output id routes by model, and that branch must build the snapshot. + + The managed-files pre-call hook sets data["model"] for any id carrying + llm_output_file_id, so batch output retrieval always takes the model-routed branch + and never reaches managed_files_obj.afile_content. Bedrock resolves its output + bucket only from _litellm_internal_model_credentials, so without the snapshot every + Bedrock batch output retrieval fails with "S3 bucket_name is required". + """ + import base64 + from types import MappingProxyType + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.types.utils import SpecialEnums + + router = Router( + model_list=[ + { + "model_name": "anthropic.batch.claude-4.5-haiku", + "litellm_params": { + "model": "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-east-1", + "s3_bucket_name": "configured-batch-bucket", + }, + "model_info": {"id": "bedrock-batch-deployment-id"}, + } + ] + ) + + from unittest.mock import MagicMock + + managed_file_row = MagicMock() + managed_file_row.created_by = "test-user" + managed_file_row.team_id = None + managed_file_row.storage_backend = None + managed_file_row.storage_url = None + prisma_stub = MagicMock() + prisma_stub.db.litellm_managedfiletable.find_first = AsyncMock(return_value=managed_file_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_stub) + setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + + captured_kwargs: dict = {} + + async def _mock_router_afile_content(**kwargs): + captured_kwargs.update(kwargs) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=b'{"recordId":"req-1"}', + headers={"content-type": "application/octet-stream"}, + ) + ) + + monkeypatch.setattr(router, "afile_content", _mock_router_afile_content) + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/jsonl", + "unified-output-id", + "anthropic.batch.claude-4.5-haiku", + "llm_output_file_id,s3://configured-batch-bucket/out/batch.jsonl", + "bedrock-batch-deployment-id", + ) + encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + try: + response = client.get( + f"/v1/files/{encoded_id}/content", + headers={"Authorization": "Bearer test-key", "custom-llm-provider": "bedrock"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + snapshot = captured_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "model-routed branch must attach the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + + def _unified_managed_file_id() -> str: import base64