scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

v0 set arn number on dynamoDB /key/gen

Browse source
This commit is contained in:
ishaan-jaff 2024-01-24 10:29:28 -08:00
parent d9fd8f28d6
commit 9b39893980
4 changed files with 46 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
docs/my-website/docs/proxy/virtual_keys.md
View file

@ -554,7 +554,9 @@ general_settings:
"region_name": "us-west-2"
"user_table_name": "your-user-table",
"key_table_name": "your-token-table",
"config_table_name": "your-config-table"
"config_table_name": "your-config-table",
"aws_role_name": "your-aws_role_name",
"aws_session_name": "your-aws_session_name",
}
```

2
litellm/proxy/_types.py
View file

@ -208,6 +208,8 @@ class DynamoDBArgs(LiteLLMBase):
key_table_name: str = "LiteLLM_VerificationToken"
config_table_name: str = "LiteLLM_Config"
spend_table_name: str = "LiteLLM_SpendLogs"
aws_role_name: Optional[str] = None
aws_session_name: Optional[str] = None
class ConfigGeneralSettings(LiteLLMBase):

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

@ -52,6 +52,32 @@ class DynamoDBWrapper(CustomDB):
self.database_arguments = database_arguments
self.region_name = database_arguments.region_name
def set_env_vars_based_on_arn(self):
if self.database_arguments.aws_role_name is None:
return
verbose_proxy_logger.debug(
f"DynamoDB: setting env vars based on arn={self.database_arguments.aws_role_name}"
)
import boto3, os
sts_client = boto3.client("sts")
assumed_role = sts_client.assume_role(
RoleArn=self.database_arguments.aws_role_name,
RoleSessionName=self.database_arguments.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"]
verbose_proxy_logger.debug(
f"Got STS assumed Role, aws_access_key_id={aws_access_key_id}"
)
# set these in the env so aiodynamo can use them
os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
os.environ["AWS_SESSION_TOKEN"] = aws_session_token
async def connect(self):
"""
Connect to DB, and creating / updating any tables
@ -74,6 +100,7 @@ class DynamoDBWrapper(CustomDB):
import aiohttp
verbose_proxy_logger.debug("DynamoDB Wrapper - Attempting to connect")
self.set_env_vars_based_on_arn()
# before making ClientSession check if ssl_verify=False
if self.database_arguments.ssl_verify == False:
client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
@ -170,6 +197,8 @@ class DynamoDBWrapper(CustomDB):
from aiohttp import ClientSession
import aiohttp
self.set_env_vars_based_on_arn()
if self.database_arguments.ssl_verify == False:
client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
else:
@ -210,6 +239,8 @@ class DynamoDBWrapper(CustomDB):
from aiohttp import ClientSession
import aiohttp
self.set_env_vars_based_on_arn()
if self.database_arguments.ssl_verify == False:
client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
else:
@ -253,6 +284,7 @@ class DynamoDBWrapper(CustomDB):
async def update_data(
self, key: str, value: dict, table_name: Literal["user", "key", "config"]
):
self.set_env_vars_based_on_arn()
from aiodynamo.client import Client
from aiodynamo.credentials import Credentials, StaticCredentials
from aiodynamo.http.httpx import HTTPX
@ -323,4 +355,5 @@ class DynamoDBWrapper(CustomDB):
"""
Not Implemented yet.
"""
self.set_env_vars_based_on_arn()
return super().delete_data(keys, table_name)

14
litellm/proxy/proxy_config.yaml
View file

@ -67,12 +67,14 @@ litellm_settings:
general_settings:
master_key: sk-1234
# database_type: "dynamo_db"
# database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190
# "billing_mode": "PAY_PER_REQUEST",
# "region_name": "us-west-2",
# "ssl_verify": False
# }
database_type: "dynamo_db"
database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190
"billing_mode": "PAY_PER_REQUEST",
"region_name": "us-west-2",
"ssl_verify": False,
"aws_role_name": "<your-aws_role_name>",
"aws_session_name": "<your-aws_session_name>",
}