phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) dyanmo fixes

Browse source
This commit is contained in:
ishaan-jaff 2024-02-12 19:54:39 -08:00
parent 551a4c60a7
commit c4038e38e5
2 changed files with 11 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -215,6 +215,8 @@ class DynamoDBArgs(LiteLLMBase):
aws_policy_arns: Optional[List[str]] = None
aws_policy: Optional[str] = None
aws_duration_seconds: Optional[int] = None
assume_role_aws_role_name: Optional[str] = None
assume_role_aws_session_name: Optional[str] = None
class ConfigGeneralSettings(LiteLLMBase):

10
litellm/proxy/db/dynamo_db.py
View file

@ -61,12 +61,20 @@ class DynamoDBWrapper(CustomDB):
import boto3, os
sts_client = boto3.client("sts")
assumed_role = sts_client.assume_role_with_web_identity(
# call 1
non_used_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,
)
# call 2
assumed_role = sts_client.assume_role(
RoleArn=self.database_arguments.assume_role_aws_role_name,
RoleSessionName=self.database_arguments.assume_role_aws_session_name,
)
aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"]
aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"]
aws_session_token = assumed_role["Credentials"]["SessionToken"]