mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
commit
0acca3e86a
5 changed files with 51 additions and 7 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue