mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #37047 from BerriAI/litellm_bedrock_batch_output_bucket
fix(bedrock): resolve the managed-batch output bucket on every path that reads it
This commit is contained in:
Mateo Wang
GitHub
commit
5e9adb9ed7
14 changed files with 429 additions and 15 deletions
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue