Merge pull request #24548 from mpcusack-altos/fix/bedrock-batch-credential-fields

fix(router): include Bedrock batch/S3 fields and model in deployment credentials
This commit is contained in:
Mateo Wang 2026-08-05 22:39:39 -07:00 committed by GitHub
commit 0acca3e86a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 51 additions and 7 deletions

View file

@ -373,7 +373,7 @@ def get_team_provider_credentials(
def _provider_credentials(model_id: str) -> dict | None:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider:
return credentials
return {key: value for key, value in credentials.items() if key != "model"}
return None
# 1. Prefer the team's own BYOK deployment, matched by model_info.team_id.

View file

@ -8738,7 +8738,7 @@ class Router:
Example:
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...}
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", "model": "gpt-4o", ...}
"""
# Try to get deployment by model_id first
deployment = self.get_deployment(model_id=model_id)
@ -8797,6 +8797,8 @@ class Router:
# Remove the credential name since we've resolved it
credentials.pop("litellm_credential_name", None)
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

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

View file

@ -4024,6 +4024,42 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
litellm.credential_list = []
def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
"""
Test that get_deployment_credentials_with_provider returns the deployment's
model and the Bedrock batch/S3 fields (s3_region_name, s3_encryption_key_id,
aws_batch_role_arn) instead of silently dropping them (#25104).
"""
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-west-2",
"s3_bucket_name": "my-batch-bucket",
"s3_region_name": "us-east-1",
"s3_encryption_key_id": "arn:aws:kms:us-west-2:123:key/abc",
"aws_batch_role_arn": "arn:aws:iam::123:role/batch-role",
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(
model_id="bedrock-batch-model"
)
assert credentials is not None
assert credentials["custom_llm_provider"] == "bedrock"
assert credentials["model"] == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
assert credentials["aws_region_name"] == "us-west-2"
assert credentials["s3_bucket_name"] == "my-batch-bucket"
assert credentials["s3_region_name"] == "us-east-1"
assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-west-2:123:key/abc"
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role"
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
return {
"model_name": f"model_name_team-1_{model_id}",

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 */
@ -26949,6 +26951,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;
@ -35316,6 +35320,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 */
@ -35549,6 +35555,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;