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 7745938b9e7..1e58de0146c 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/batches/batch_utils.py b/litellm/batches/batch_utils.py index baf522c0bb1..9681d64f656 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -5,6 +5,7 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch from litellm.types.utils import CallTypes, ModelInfo, Usage @@ -295,7 +296,7 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: if litellm_params: # List of credential keys that should be passed to file operations - credential_keys: Final = [ + credential_keys: Final = ( "api_key", "api_base", "api_version", @@ -309,7 +310,9 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "bucket_name", "timeout", "max_retries", - ] + "_litellm_internal_model_credentials", + *AWS_CREDENTIAL_KWARGS_KEYS, + ) for key in credential_keys: if key in litellm_params: credentials[key] = litellm_params[key] diff --git a/litellm/batches/main.py b/litellm/batches/main.py index ce52c12818e..20d38bbb77f 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -22,6 +22,7 @@ from openai.types.batch import BatchRequestCounts import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -527,6 +528,7 @@ def retrieve_batch( custom_llm_provider=custom_llm_provider, **kwargs, ) + add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) if litellm_logging_obj is not None: litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/files/main.py b/litellm/files/main.py index 34421d13761..9a64c78552b 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -11,7 +11,6 @@ import time import uuid as uuid_module from collections.abc import Coroutine from functools import partial -from types import MappingProxyType from typing import Any, Final, Literal, cast import httpx @@ -34,6 +33,7 @@ import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse from litellm.files.types import FileContentProvider, FileContentStreamingResult +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials @@ -85,14 +85,6 @@ bedrock_files_instance: Final = BedrockFilesHandler() ################################################# -def _add_trusted_model_credentials_to_litellm_params( - litellm_params_dict: dict[str, Any], kwargs: dict[str, Any] -) -> None: - trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") - if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials - - @client async def acreate_file( file: FileTypes, @@ -372,7 +364,7 @@ def file_retrieve( ) if provider_config is not None: litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -494,7 +486,7 @@ def file_delete( pass optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -834,7 +826,7 @@ def file_content( try: optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index f251ab4d74a..3eb8c163d5c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping, MutableMapping +from types import MappingProxyType from typing import Final from litellm.llms.openai.data_residency import infer_openai_data_residency @@ -184,3 +186,19 @@ def get_litellm_params( litellm_params[key] = kwargs[key] return litellm_params + + +def add_trusted_model_credentials_to_litellm_params( + litellm_params_dict: MutableMapping[str, object], kwargs: Mapping[str, object] +) -> None: + """ + Carry the immutable server-side credential snapshot into litellm_params. + + get_litellm_params has a fixed signature, so callers that need the snapshot to + survive into the logging object and the downstream file read have to re-add it. Only + a MappingProxyType is accepted, since providers resolve trusted configuration such + as a Bedrock file bucket from it and must not read a request-supplied mapping. + """ + trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") + if isinstance(trusted_model_credentials, MappingProxyType): + litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e442cefa360..87f9927b191 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + add_internal_model_credentials, apply_team_provider_credentials, decode_model_from_file_id, encode_batch_response_ids, @@ -537,6 +538,13 @@ async def retrieve_batch( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) + if unified_batch_id: + add_internal_model_credentials( + data=data, + llm_router=llm_router, + model_id=get_model_id_from_unified_batch_id(unified_batch_id), + ) + response = await llm_router.aretrieve_batch(**data) response._hidden_params["unified_batch_id"] = unified_batch_id if unified_batch_id: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 56e986c89cf..f2e6fb633e1 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -465,6 +465,31 @@ def apply_team_provider_credentials( prepare_data_with_credentials(data=data, credentials=credentials) +def add_internal_model_credentials( + data: dict, + llm_router: "Router", + model_id: str | None, +) -> None: + """ + Attach the deployment's immutable server-side credential snapshot to a router-routed + batch call (in-place). + + Cost accounting for a completed batch reads the batch's output file, and the Bedrock + file config resolves its bucket only from this snapshot, never from a request param, + because the bucket is what managed file ids are validated against. Without it that + read fails and the batch's cost is never recorded. + """ + if model_id is None: + return + try: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + except Exception: # noqa: BLE001 # the snapshot only enables cost accounting; a batch whose deployment no longer resolves must still be retrievable + return + if credentials is None: + return + data["_litellm_internal_model_credentials"] = MappingProxyType(dict(credentials)) + + def prepare_data_with_credentials( data: dict, credentials: dict, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2432ea3729..361b5b920e2 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + add_internal_model_credentials, apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, @@ -706,6 +707,7 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: + add_internal_model_credentials(data=data, llm_router=llm_router, model_id=model) response = await llm_router.afile_content( **{ "model": model, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index b4c8aad81ed..6bb19a07d6c 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/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 523b512e4cf..d2074853f2b 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. import json import os import sys +from types import MappingProxyType import httpx import pytest @@ -1229,3 +1230,72 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) assert models == ["claude-sonnet-4-5"] + + +def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): + """Bedrock resolves a batch's output bucket only from the immutable server-side + snapshot, never from a request param, so cost accounting on the retrieve path cannot + read the output file unless this key is forwarded. Without it the accounting raises + "S3 bucket_name is required" for a bucket the deployment has configured, and the + batch's cost is never recorded.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"}) + + credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot}) + + assert credentials["_litellm_internal_model_credentials"] is snapshot + + +def test_extract_credentials_forwards_the_deployment_aws_credentials(): + """The retrieve path's logging object carries the deployment's AWS keys in its + litellm_params, and the S3 read of the output file signs with whatever afile_content + receives. Dropping them here sent the read to the ambient credential chain, so a + deployment whose only AWS credentials live in its litellm_params never recorded + batch cost on retrieve even once the bucket resolved.""" + params = { + "aws_access_key_id": "AKIA-deployment", + "aws_secret_access_key": "secret-deployment", + "aws_session_token": "token-deployment", + "aws_region_name": "us-west-2", + "aws_role_name": "arn:aws:iam::123456789012:role/batch-reader", + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + } + + credentials = bu._extract_file_access_credentials(params) + + assert credentials == {key: value for key, value in params.items() if key != "model"} + + +@pytest.mark.asyncio +async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b""})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-west-2"}) + + await bu._fetch_batch_output_file_content( + _batch("s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"), + custom_llm_provider="bedrock", + litellm_params={ + "aws_access_key_id": "AKIA-deployment", + "aws_secret_access_key": "secret-deployment", + "aws_session_token": "token-deployment", + "aws_region_name": "us-west-2", + "_litellm_internal_model_credentials": snapshot, + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + }, + ) + + assert captured["file_id"] == "s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out" + assert captured["custom_llm_provider"] == "bedrock" + assert captured["aws_access_key_id"] == "AKIA-deployment" + assert captured["aws_secret_access_key"] == "secret-deployment" + assert captured["aws_session_token"] == "token-deployment" + assert captured["aws_region_name"] == "us-west-2" + assert captured["_litellm_internal_model_credentials"] is snapshot + assert "model" not in captured diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index 1f7a91a5511..17e9ee29d4d 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -28,6 +28,7 @@ import sys from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict +from types import MappingProxyType from unittest.mock import MagicMock, patch import pytest @@ -742,3 +743,33 @@ def test_resolve_timeout__httpx_timeout_returns_float_read(): resolved = bm._resolve_timeout(_params(timeout=t), {}, "openai") assert isinstance(resolved, float) assert resolved == 99.0 + + +def test_retrieve__forwards_trusted_model_credentials_into_litellm_params(seams): + """The batch's cost is computed by reading its output file after the retrieve, and + Bedrock resolves that bucket only from this immutable snapshot. get_litellm_params has + a fixed signature that drops it, so without re-adding it here the snapshot never + reaches the logging object and cost accounting fails on a bucket that is configured.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket"}) + logging_obj = MagicMock() + + bm.retrieve_batch( + batch_id="batch-1", + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + _litellm_internal_model_credentials=snapshot, + ) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is snapshot + + +def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): + """A retrieve with no snapshot must not invent an empty one, which would read as a + configured bucket of nothing.""" + logging_obj = MagicMock() + + bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai", litellm_logging_obj=logging_obj) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert "_litellm_internal_model_credentials" not in litellm_params diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a80c19f0708..aa5c63280b8 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1138,7 +1138,11 @@ async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness): # DISPATCH - router fired, direct litellm did not. assert retrieve_harness.router_aretrieve.call_count == 1 retrieve_harness.litellm_aretrieve.assert_not_called() - retrieve_harness.creds_resolver.assert_not_called() + + # Credentials are resolved for the deployment behind the unified id so the batch's + # output file can be read for cost accounting. This id resolves to nothing here, and + # the retrieve must still serve the batch rather than fail on the lookup. + retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini") # router receives the (still-encoded) batch id verbatim - this layer does # not decode it for the unified path. diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4a021627c3e..ad7f5e4725a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -95,3 +95,71 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): "unified-out", "unified-err", ) + + +# =========================================================================== # +# add_internal_model_credentials - the snapshot that lets a completed +# batch's output file be read, and therefore its cost be recorded +# =========================================================================== # + + +def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): + """Cost accounting for a completed batch reads its output file, and Bedrock resolves + that bucket only from this snapshot. It must be immutable so nothing downstream can + redirect the bucket that managed file ids are validated against.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"} + ) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-1") + + snapshot = data["_litellm_internal_model_credentials"] + assert snapshot["s3_bucket_name"] == "configured-bucket" + assert isinstance(snapshot, MappingProxyType) + with pytest.raises(TypeError): + snapshot["s3_bucket_name"] = "attacker-bucket" + router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="deployment-1") + + +@pytest.mark.parametrize( + "model_id, credentials", + [(None, {"s3_bucket_name": "b"}), ("deployment-1", None)], + ids=["no-model-id", "deployment-has-no-credentials"], +) +def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deployment(model_id, credentials): + """An unroutable batch must be left alone rather than given an empty snapshot, which + would look like a configured bucket of nothing.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials(data=data, llm_router=router, model_id=model_id) + + assert "_litellm_internal_model_credentials" not in data + + +def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): + """The snapshot only enables cost accounting, so a batch whose deployment no longer + resolves, which happens when a model group is removed while batches are in flight, + must still be retrievable rather than failing the request on the lookup.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone")) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") + + assert data == {"batch_id": "unified-batch-id"} 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..e363a266688 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,99 @@ 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) + + # One frozen snapshot per call rather than one dict merged across calls, so a second + # invocation is visible instead of silently overwriting the first. + calls: list[MappingProxyType] = [] + + async def _mock_router_afile_content(**kwargs): + calls.append(MappingProxyType(dict(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 + assert len(calls) == 1, f"expected exactly one routed retrieval, got {len(calls)}" + snapshot = calls[0].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