scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-01 02:02:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 92

(dynamo) use assume_role_with_web_identity

Browse source
This commit is contained in:
ishaan-jaff 2024-02-02 10:57:06 -08:00
parent d2c1e1554f
commit 9c13c65ce2
2 changed files with 11 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/_types.py
View file

@ -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):

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

@ -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"]