mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(router): keep bedrock batch litellm_params through credential resolution
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
8177230a29
commit
8c5f96e940
2 changed files with 44 additions and 0 deletions
|
|
@ -211,6 +211,11 @@ class CredentialLiteLLMParams(BaseModel):
|
|||
aws_bedrock_runtime_endpoint: Optional[str] = None
|
||||
aws_bedrock_project_id: Optional[str] = None
|
||||
s3_bucket_name: Optional[str] = None
|
||||
s3_output_bucket_name: Optional[str] = None
|
||||
s3_region_name: Optional[str] = None
|
||||
s3_encryption_key_id: Optional[str] = None
|
||||
aws_batch_role_arn: Optional[str] = None
|
||||
bedrock_tags: Optional[List[Dict[str, str]]] = None
|
||||
## IBM WATSONX ##
|
||||
watsonx_region_name: Optional[str] = None
|
||||
|
||||
|
|
|
|||
|
|
@ -3666,6 +3666,45 @@ 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_params():
|
||||
"""
|
||||
Regression: the Bedrock batch/files params must survive the CredentialLiteLLMParams
|
||||
filter. When they were dropped, model-routed POST /v1/batches failed with
|
||||
"Set 'aws_batch_role_arn' in litellm_params or AWS_BATCH_ROLE_ARN env var" even
|
||||
though the deployment declared it.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-batch",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_region_name": "us-east-1",
|
||||
"s3_bucket_name": "my-input-bucket",
|
||||
"s3_output_bucket_name": "my-output-bucket",
|
||||
"s3_region_name": "us-west-2",
|
||||
"s3_encryption_key_id": "arn:aws:kms:us-east-1:123:key/abc",
|
||||
"aws_batch_role_arn": "arn:aws:iam::123:role/BedrockBatchRole",
|
||||
"bedrock_tags": [{"key": "team", "value": "platform"}],
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="bedrock-batch"
|
||||
)
|
||||
|
||||
assert credentials is not None
|
||||
assert credentials["s3_bucket_name"] == "my-input-bucket"
|
||||
assert credentials["s3_output_bucket_name"] == "my-output-bucket"
|
||||
assert credentials["s3_region_name"] == "us-west-2"
|
||||
assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-east-1:123:key/abc"
|
||||
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/BedrockBatchRole"
|
||||
assert credentials["bedrock_tags"] == [{"key": "team", "value": "platform"}]
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue