From 9c13c65ce24b27aa8a6c24ed0355d5c5a094af17 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 2 Feb 2024 10:57:06 -0800 Subject: [PATCH] (dynamo) use assume_role_with_web_identity --- litellm/proxy/_types.py | 5 +++++ litellm/proxy/db/dynamo_db.py | 7 ++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 88ee88bcd0f..10290fc3c11 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -210,6 +210,11 @@ class DynamoDBArgs(LiteLLMBase): spend_table_name: str = "LiteLLM_SpendLogs" aws_role_name: Optional[str] = None aws_session_name: Optional[str] = None + aws_web_identity_token: Optional[str] = None + aws_provider_id: Optional[str] = None + aws_policy_arns: Optional[List[str]] = None + aws_policy: Optional[str] = None + aws_duration_seconds: Optional[int] = None class ConfigGeneralSettings(LiteLLMBase): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 50330ccf3bf..3ea2b285ae3 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -61,9 +61,14 @@ class DynamoDBWrapper(CustomDB): import boto3, os sts_client = boto3.client("sts") - assumed_role = sts_client.assume_role( + assumed_role = sts_client.assume_role_with_web_identity( RoleArn=self.database_arguments.aws_role_name, RoleSessionName=self.database_arguments.aws_session_name, + WebIdentityToken=self.database_arguments.aws_web_identity_token, + ProviderId=self.database_arguments.aws_provider_id, + PolicyArns=self.database_arguments.aws_policy_arns, + Policy=self.database_arguments.aws_policy, + DurationSeconds=self.database_arguments.aws_duration_seconds, ) aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"]