phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(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_arns: Optional[List[str]] = None
aws_policy: Optional[str] = None aws_policy: Optional[str] = None
aws_duration_seconds: Optional[int] = 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): class ConfigGeneralSettings(LiteLLMBase):

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

@ -61,12 +61,20 @@ class DynamoDBWrapper(CustomDB):
import boto3, os import boto3, os
sts_client = boto3.client("sts") 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, RoleArn=self.database_arguments.aws_role_name,
RoleSessionName=self.database_arguments.aws_session_name, RoleSessionName=self.database_arguments.aws_session_name,
WebIdentityToken=self.database_arguments.aws_web_identity_token, 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_access_key_id = assumed_role["Credentials"]["AccessKeyId"]
aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"] aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"]
aws_session_token = assumed_role["Credentials"]["SessionToken"] aws_session_token = assumed_role["Credentials"]["SessionToken"]