mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
fix(router): return model and Bedrock batch fields in deployment credentials
get_deployment_credentials_with_provider dropped s3_region_name, s3_encryption_key_id, and aws_batch_role_arn because CredentialLiteLLMParams never declared them, and it never returned the deployment's model, so proxy batch creation against Bedrock failed with "LiteLLM doesn't support custom_llm_provider=bedrock for 'create_batch'" or "AWS IAM role ARN is required" (#25104) Provider-only file and batch calls keep their no-model contract: get_team_provider_credentials strips the model key so a provider-scoped request is not pinned to an arbitrary matching deployment
This commit is contained in:
parent
ba91768146
commit
3d275d97fe
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:
|
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)
|
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:
|
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
|
return None
|
||||||
|
|
||||||
# 1. Prefer the team's own BYOK deployment, matched by model_info.team_id.
|
# 1. Prefer the team's own BYOK deployment, matched by model_info.team_id.
|
||||||
|
|
|
||||||
|
|
@ -8738,7 +8738,7 @@ class Router:
|
||||||
|
|
||||||
Example:
|
Example:
|
||||||
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
|
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
|
# Try to get deployment by model_id first
|
||||||
deployment = self.get_deployment(model_id=model_id)
|
deployment = self.get_deployment(model_id=model_id)
|
||||||
|
|
@ -8797,6 +8797,8 @@ class Router:
|
||||||
# Remove the credential name since we've resolved it
|
# Remove the credential name since we've resolved it
|
||||||
credentials.pop("litellm_credential_name", None)
|
credentials.pop("litellm_credential_name", None)
|
||||||
|
|
||||||
|
credentials["model"] = deployment.litellm_params.model
|
||||||
|
|
||||||
# Add custom_llm_provider
|
# Add custom_llm_provider
|
||||||
if deployment.litellm_params.custom_llm_provider:
|
if deployment.litellm_params.custom_llm_provider:
|
||||||
credentials["custom_llm_provider"] = 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_runtime_endpoint: str | None = None
|
||||||
aws_bedrock_project_id: str | None = None
|
aws_bedrock_project_id: str | None = None
|
||||||
s3_bucket_name: 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 ##
|
## IBM WATSONX ##
|
||||||
watsonx_region_name: str | None = None
|
watsonx_region_name: str | None = None
|
||||||
|
|
||||||
|
|
@ -272,11 +275,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
||||||
quality_router_config: dict | None = None
|
quality_router_config: dict | None = None
|
||||||
quality_router_default_model: str | 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 Params
|
||||||
vector_store_id: str | None = None
|
vector_store_id: str | None = None
|
||||||
milvus_text_field: 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 = []
|
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:
|
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
|
||||||
return {
|
return {
|
||||||
"model_name": f"model_name_team-1_{model_id}",
|
"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;
|
auto_router_max_input_chars?: number | null;
|
||||||
/** Aws Access Key Id */
|
/** Aws Access Key Id */
|
||||||
aws_access_key_id?: string | null;
|
aws_access_key_id?: string | null;
|
||||||
|
/** Aws Batch Role Arn */
|
||||||
|
aws_batch_role_arn?: string | null;
|
||||||
/** Aws Bedrock Project Id */
|
/** Aws Bedrock Project Id */
|
||||||
aws_bedrock_project_id?: string | null;
|
aws_bedrock_project_id?: string | null;
|
||||||
/** Aws Bedrock Runtime Endpoint */
|
/** Aws Bedrock Runtime Endpoint */
|
||||||
|
|
@ -26949,6 +26951,8 @@ export interface components {
|
||||||
s3_bucket_name?: string | null;
|
s3_bucket_name?: string | null;
|
||||||
/** S3 Encryption Key Id */
|
/** S3 Encryption Key Id */
|
||||||
s3_encryption_key_id?: string | null;
|
s3_encryption_key_id?: string | null;
|
||||||
|
/** S3 Region Name */
|
||||||
|
s3_region_name?: string | null;
|
||||||
/** Search Context Cost Per Query */
|
/** Search Context Cost Per Query */
|
||||||
search_context_cost_per_query?: {
|
search_context_cost_per_query?: {
|
||||||
[key: string]: unknown;
|
[key: string]: unknown;
|
||||||
|
|
@ -35316,6 +35320,8 @@ export interface components {
|
||||||
auto_router_max_input_chars?: number | null;
|
auto_router_max_input_chars?: number | null;
|
||||||
/** Aws Access Key Id */
|
/** Aws Access Key Id */
|
||||||
aws_access_key_id?: string | null;
|
aws_access_key_id?: string | null;
|
||||||
|
/** Aws Batch Role Arn */
|
||||||
|
aws_batch_role_arn?: string | null;
|
||||||
/** Aws Bedrock Project Id */
|
/** Aws Bedrock Project Id */
|
||||||
aws_bedrock_project_id?: string | null;
|
aws_bedrock_project_id?: string | null;
|
||||||
/** Aws Bedrock Runtime Endpoint */
|
/** Aws Bedrock Runtime Endpoint */
|
||||||
|
|
@ -35549,6 +35555,8 @@ export interface components {
|
||||||
s3_bucket_name?: string | null;
|
s3_bucket_name?: string | null;
|
||||||
/** S3 Encryption Key Id */
|
/** S3 Encryption Key Id */
|
||||||
s3_encryption_key_id?: string | null;
|
s3_encryption_key_id?: string | null;
|
||||||
|
/** S3 Region Name */
|
||||||
|
s3_region_name?: string | null;
|
||||||
/** Search Context Cost Per Query */
|
/** Search Context Cost Per Query */
|
||||||
search_context_cost_per_query?: {
|
search_context_cost_per_query?: {
|
||||||
[key: string]: unknown;
|
[key: string]: unknown;
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue