scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

fix(proxy): stop dropping bedrock batch credentials and model in batch endpoints

Browse source 
CredentialLiteLLMParams now whitelists s3_region_name, s3_encryption_key_id,
and aws_batch_role_arn so they survive the credential round-trip, and the
proxy batch endpoints opt in to receiving the deployment's underlying model
so litellm.acreate_batch can resolve the bedrock provider config. Router-bound
callers (vector store files) keep the router alias untouched.

Fixes #25104
This commit is contained in:
mateo-berri 2026-08-05 21:31:01 -07:00
parent d26ef670e2
commit d210cd272c
9 changed files with 124 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -189,6 +189,7 @@ async def create_batch(
llm_router=llm_router,
model_id=model_from_file_id,
operation_context="batch creation (file created with model)",
include_model=True,
)
original_file_id: Final = get_original_file_id(input_file_id)
@ -280,6 +281,7 @@ async def create_batch(
llm_router=llm_router,
model_id=model_param,
operation_context="batch creation",
include_model=True,
)
prepare_data_with_credentials(
@ -489,6 +491,7 @@ async def retrieve_batch(
llm_router=llm_router,
model_id=model_from_id,
operation_context="batch retrieval (batch created with model)",
include_model=True,
)
original_batch_id: Final = get_original_file_id(batch_id)
@ -699,6 +702,7 @@ async def list_batches(
llm_router=llm_router,
model_id=model_param,
operation_context="batch listing",
include_model=True,
)
data.update(credentials)
@ -880,6 +884,7 @@ async def cancel_batch(
llm_router=llm_router,
model_id=model_from_id,
operation_context="batch cancellation (batch created with model)",
include_model=True,
)
original_batch_id: Final = get_original_file_id(batch_id)

8
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -272,6 +272,7 @@ def get_credentials_for_model(
llm_router, # Router instance
model_id: str,
operation_context: str = "file operation",
include_model: bool = False,
):
"""
Retrieve API credentials for a model from the LLM Router.
@ -280,6 +281,9 @@ def get_credentials_for_model(
llm_router: LiteLLM Router instance
model_id: Model name or deployment ID
operation_context: Description for error messages (e.g., "file upload", "batch creation")
include_model: Include the deployment's underlying model in the returned credentials.
Only safe for call sites that dispatch directly to litellm.* functions; router-bound
paths must keep the router alias in "model" for deployment lookup to succeed.
Returns:
Dictionary with credentials (api_key, api_base, custom_llm_provider, etc.)
@ -295,7 +299,9 @@ def get_credentials_for_model(
detail={"error": "Router not initialized. Cannot use model-based routing."},
)
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
credentials: Final = llm_router.get_deployment_credentials_with_provider(
model_id=model_id, include_model=include_model
)
if credentials is None:
raise HTTPException(

12
litellm/router.py
View file

@ -8712,7 +8712,7 @@ class Router:
)
def get_deployment_credentials_with_provider(
self, model_id: str, team_id: str | None = None
self, model_id: str, team_id: str | None = None, include_model: bool = False
) -> dict[str, Any] | None:
"""
Get API credentials and provider info from a model name in model_list.
@ -8729,6 +8729,13 @@ class Router:
wildcard lookups never resolve a deployment owned by a
different team, so shared model names can't leak another
team's credentials.
include_model: When True, the deployment's `litellm_params.model`
is returned under the "model" key. Batch and file operations
need it so `litellm.acreate_batch` / `aretrieve_batch` can
resolve provider configs (e.g. Bedrock) and set the provider
model id. Off by default because callers that merge these
credentials into router-bound request data (e.g. vector store
files) must keep the router alias in "model".
Returns:
Dictionary containing api_key, api_base, custom_llm_provider, etc.
@ -8797,6 +8804,9 @@ class Router:
# Remove the credential name since we've resolved it
credentials.pop("litellm_credential_name", None)
if include_model:
credentials["model"] = deployment.litellm_params.model
# Add custom_llm_provider
if deployment.litellm_params.custom_llm_provider:
credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider

8
litellm/types/router.py
View file

@ -200,6 +200,9 @@ class CredentialLiteLLMParams(BaseModel):
aws_bedrock_runtime_endpoint: str | None = None
aws_bedrock_project_id: str | None = None
s3_bucket_name: str | None = None
s3_region_name: str | None = None
s3_encryption_key_id: str | None = None
aws_batch_role_arn: str | None = None
## IBM WATSONX ##
watsonx_region_name: str | None = None
@ -272,11 +275,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
quality_router_config: dict | None = None
quality_router_default_model: str | None = None
# Batch/File API Params
s3_bucket_name: str | None = None
s3_encryption_key_id: str | None = None
gcs_bucket_name: str | None = None
# Vector Store Params
vector_store_id: str | None = None
milvus_text_field: str | None = None

33
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -146,9 +146,12 @@ class Harness:
return dict(self.router_acreate.call_args.kwargs)
def _creds_lookup(*, model_id: str) -> Dict[str, str]:
def _creds_lookup(*, model_id: str, include_model: bool = False) -> Dict[str, str]:
# KeyError on an unknown/hardcoded model_id - the bug cannot hide.
return dict(CREDS[model_id])
creds = dict(CREDS[model_id])
if not include_model:
creds.pop("model")
return creds
@pytest.fixture
@ -259,7 +262,7 @@ async def test_create__model_encoded_file_id(harness):
harness.router_acreate.assert_not_called()
# 2. CREDENTIALS - resolved for the model decoded FROM the file id.
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this.
assert harness.acreate_kwargs() == {
@ -315,7 +318,7 @@ async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harnes
await call_create(harness)
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# =========================================================================== #
@ -339,7 +342,7 @@ async def test_create__model_from_body(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", include_model=True)
payload = harness.acreate_kwargs()
assert payload["custom_llm_provider"] == "vertex_ai"
assert payload["input_file_id"] == "file-plain"
@ -359,7 +362,7 @@ async def test_create__model_from_header(harness):
await call_create(harness, headers={"x-litellm-model": "vertex-model"})
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", include_model=True)
harness.router_acreate.assert_not_called()
@ -376,7 +379,7 @@ async def test_create__model_from_query(harness):
await call_create(harness, query={"model": "vertex-model"})
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", include_model=True)
harness.router_acreate.assert_not_called()
@ -399,7 +402,7 @@ async def test_create__body_model_beats_header_and_query(harness):
query={"model": "vertex-model"},
)
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# =========================================================================== #
@ -707,7 +710,7 @@ async def test_create__model_encoded_beats_unified(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# =========================================================================== #
@ -753,7 +756,7 @@ async def test_create__model_encoded_beats_loadbalancing(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# =========================================================================== #
@ -1060,7 +1063,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
retrieve_harness.router_aretrieve.assert_not_called()
# 2. CREDENTIALS - resolved for the model decoded FROM the batch id.
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call.
# Note `model` is the DECODED model, not the deployment from creds: the
@ -1118,7 +1121,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
assert retrieve_harness.litellm_aretrieve.call_count == 1
retrieve_harness.router_aretrieve.assert_not_called()
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# --------------------------------------------------------------------------- #
@ -1536,7 +1539,7 @@ async def test_list__model_from_body_routes_and_encodes(list_harness):
assert list_harness.litellm_alist.call_count == 1
list_harness.router_alist.assert_not_called()
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch")
assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch")
@ -1826,7 +1829,7 @@ async def test_cancel__model_encoded_id(cancel_harness):
cancel_harness.router_acancel.assert_not_called()
# CREDENTIALS - resolved for the model decoded from the batch id.
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the
# DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike
@ -1864,7 +1867,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness):
assert cancel_harness.litellm_acancel.call_count == 1
cancel_harness.router_acancel.assert_not_called()
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", include_model=True)
# --------------------------------------------------------------------------- #

5
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -158,7 +158,6 @@ async def test_vector_store_file_list_resolves_credentials_from_model_query_para
"api_key": "sk-team-openai",
"api_base": "https://api.openai.com/v1",
"custom_llm_provider": "openai",
"model": "openai/gpt-4o-mini",
}
data = {
@ -174,10 +173,10 @@ async def test_vector_store_file_list_resolves_credentials_from_model_query_para
assert result["api_key"] == "sk-team-openai"
assert result["api_base"] == "https://api.openai.com/v1"
assert result["model"] == "openai/gpt-4o-mini"
assert "model" not in result
assert "custom_llm_provider" not in result
llm_router.get_deployment_credentials_with_provider.assert_called_once_with(
model_id="team-openai"
model_id="team-openai", include_model=False
)

8
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
View file

@ -138,12 +138,12 @@ async def test_vector_store_file_list_resolves_managed_vector_store_before_team_
llm_router = MagicMock()
def get_credentials(model_id):
def get_credentials(model_id, include_model=False):
assert include_model is False
return {
"api_key": f"sk-{model_id}",
"api_base": "https://api.openai.com/v1",
"custom_llm_provider": "openai",
"model": f"openai/{model_id}",
}
llm_router.get_deployment_credentials_with_provider.side_effect = get_credentials
@ -169,9 +169,9 @@ async def test_vector_store_file_list_resolves_managed_vector_store_before_team_
assert response == {"ok": True}
assert captured_data["vector_store_id"] == "vs_provider_native"
assert captured_data["api_key"] == "sk-managed-deployment"
assert captured_data["model"] == "openai/managed-deployment"
assert captured_data["model"] == "managed-deployment"
llm_router.get_deployment_credentials_with_provider.assert_called_once_with(
model_id="managed-deployment"
model_id="managed-deployment", include_model=False
)

66
tests/test_litellm/test_router.py
View file

@ -3976,6 +3976,72 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name():
assert credentials["custom_llm_provider"] == "vertex_ai"
def test_get_deployment_credentials_with_provider_includes_bedrock_batch_fields():
"""
Regression (#25104): s3_region_name, s3_encryption_key_id, and
aws_batch_role_arn must survive the CredentialLiteLLMParams filter.
Previously they were dropped, so Bedrock batch file uploads failed with
"S3 bucket_name is required" style errors even when configured in
litellm_params.
"""
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-batch-model",
"litellm_params": {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"aws_region_name": "us-east-1",
"s3_bucket_name": "my-batch-bucket",
"s3_region_name": "us-west-2",
"s3_encryption_key_id": "my-kms-key",
"aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch",
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(
model_id="bedrock-batch-model"
)
assert credentials is not None
assert credentials["s3_bucket_name"] == "my-batch-bucket"
assert credentials["s3_region_name"] == "us-west-2"
assert credentials["s3_encryption_key_id"] == "my-kms-key"
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123456789012:role/bedrock-batch"
assert credentials["custom_llm_provider"] == "bedrock"
assert "model" not in credentials
def test_get_deployment_credentials_with_provider_include_model():
"""
Regression (#25104): with include_model=True the deployment's
litellm_params.model is returned under "model" so batch/file operations can
resolve provider configs (litellm.acreate_batch needs a model kwarg to load
BedrockBatchesConfig and set the Bedrock modelId). Off by default so
router-bound callers keep the router alias in "model".
"""
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-batch-model",
"litellm_params": {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"aws_region_name": "us-east-1",
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(
model_id="bedrock-batch-model", include_model=True
)
assert credentials is not None
assert credentials["model"] == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
assert credentials["custom_llm_provider"] == "bedrock"
def test_get_deployment_credentials_with_provider_resolves_credential_name():
"""
Test that get_deployment_credentials_with_provider correctly resolves

8
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -26716,6 +26716,8 @@ export interface components {
auto_router_max_input_chars?: number | null;
/** Aws Access Key Id */
aws_access_key_id?: string | null;
/** Aws Batch Role Arn */
aws_batch_role_arn?: string | null;
/** Aws Bedrock Project Id */
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
@ -26945,6 +26947,8 @@ export interface components {
s3_bucket_name?: string | null;
/** S3 Encryption Key Id */
s3_encryption_key_id?: string | null;
/** S3 Region Name */
s3_region_name?: string | null;
/** Search Context Cost Per Query */
search_context_cost_per_query?: {
[key: string]: unknown;
@ -35295,6 +35299,8 @@ export interface components {
auto_router_max_input_chars?: number | null;
/** Aws Access Key Id */
aws_access_key_id?: string | null;
/** Aws Batch Role Arn */
aws_batch_role_arn?: string | null;
/** Aws Bedrock Project Id */
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
@ -35524,6 +35530,8 @@ export interface components {
s3_bucket_name?: string | null;
/** S3 Encryption Key Id */
s3_encryption_key_id?: string | null;
/** S3 Region Name */
s3_region_name?: string | null;
/** Search Context Cost Per Query */
search_context_cost_per_query?: {
[key: string]: unknown;