diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d5443bfe4b..77f9aabced 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4,6 +4,7 @@ from typing import Optional, List, Union, Dict, Literal from datetime import datetime import uuid, json + class LiteLLMBase(BaseModel): """ Implements default functions, all pydantic objects should have. @@ -15,13 +16,13 @@ class LiteLLMBase(BaseModel): except: # if using pydantic v1 return self.dict() - + def fields_set(self): try: - return self.model_fields_set # noqa + return self.model_fields_set # noqa except: # if using pydantic v1 - return self.__fields_set__ + return self.__fields_set__ ######### Request Class Definition ###### @@ -187,11 +188,15 @@ class KeyManagementSystem(enum.Enum): AZURE_KEY_VAULT = "azure_key_vault" LOCAL = "local" + class DynamoDBArgs(LiteLLMBase): billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"] read_capacity_units: Optional[int] = None write_capacity_units: Optional[int] = None region_name: str + user_table_name: str = "LiteLLM_UserTable" + key_table_name: str = "LiteLLM_VerificationToken" + config_table_name: str = "LiteLLM_Config" class ConfigGeneralSettings(LiteLLMBase): @@ -218,8 +223,13 @@ class ConfigGeneralSettings(LiteLLMBase): None, description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", ) - database_type: Optional[Literal["dynamo_db"]] = Field(None, description="to use dynamodb instead of postgres db") - database_args: Optional[DynamoDBArgs] = Field(None, description="custom args for instantiating dynamodb client - e.g. billing provision") + database_type: Optional[Literal["dynamo_db"]] = Field( + None, description="to use dynamodb instead of postgres db" + ) + database_args: Optional[DynamoDBArgs] = Field( + None, + description="custom args for instantiating dynamodb client - e.g. billing provision", + ) otel: Optional[bool] = Field( None, description="[BETA] OpenTelemetry support - this might change, use with caution.", @@ -273,10 +283,6 @@ class ConfigYAML(LiteLLMBase): class Config: protected_namespaces = () -class DBTableNames(enum.Enum): - user = "LiteLLM_UserTable" - key = "LiteLLM_VerificationToken" - config = "LiteLLM_Config" class LiteLLM_VerificationToken(LiteLLMBase): token: str @@ -289,10 +295,12 @@ class LiteLLM_VerificationToken(LiteLLMBase): max_parallel_requests: Union[int, None] metadata: Dict[str, str] = {} + class LiteLLM_Config(LiteLLMBase): param_name: str param_value: Dict + class LiteLLM_UserTable(LiteLLMBase): user_id: str max_budget: Optional[float] diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 149e730666..0196b16cf9 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -7,7 +7,6 @@ from yarl import URL from litellm.proxy.db.base_client import CustomDB from litellm.proxy._types import ( DynamoDBArgs, - DBTableNames, LiteLLM_VerificationToken, LiteLLM_Config, LiteLLM_UserTable, @@ -40,6 +39,7 @@ class DynamoDBWrapper(CustomDB): raise Exception( f"Invalid args passed in. Need to set both read_capacity_units and write_capacity_units. Args passed in - {database_arguments}" ) + self.database_arguments = database_arguments self.region_name = database_arguments.region_name async def connect(self): @@ -49,42 +49,49 @@ class DynamoDBWrapper(CustomDB): async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) ## User - table = client.table(DBTableNames.user.value) - if not await table.exists(): - sample_code_snippet = f""" - table = client.table({DBTableNames.user.value}) - await table.create( - self.throughput_type, - KeySchema(hash_key=KeySpec("user_id", KeyType.string)), - ) - """ + try: + error_occurred = False + table = client.table(self.database_arguments.user_table_name) + if not await table.exists(): + await table.create( + self.throughput_type, + KeySchema(hash_key=KeySpec("user_id", KeyType.string)), + ) + except Exception as e: + error_occurred = True + if error_occurred == True: raise Exception( - f"Failed to create table - {DBTableNames.user.value}.\nPlease create a new table called {DBTableNames.user.value}\nAND set `hash_key` as 'user_id'\n\nEg.: {sample_code_snippet}" + f"Failed to create table - {self.database_arguments.user_table_name}.\nPlease create a new table called {self.database_arguments.user_table_name}\nAND set `hash_key` as 'user_id'" ) ## Token - if not await table.exists(): - sample_code_snippet = f""" - table = client.table({DBTableNames.key.value}) - await table.create( - self.throughput_type, - KeySchema(hash_key=KeySpec("token", KeyType.string)), - ) - """ + try: + error_occurred = False + table = client.table(self.database_arguments.key_table_name) + if not await table.exists(): + await table.create( + self.throughput_type, + KeySchema(hash_key=KeySpec("token", KeyType.string)), + ) + except Exception as e: + error_occurred = True + if error_occurred == True: raise Exception( - f"Failed to create table - {DBTableNames.key.value}.\nPlease create a new table called {DBTableNames.key.value}\nAND set `hash_key` as 'token'\n\nE.g.: {sample_code_snippet}" + f"Failed to create table - {self.database_arguments.key_table_name}.\nPlease create a new table called {self.database_arguments.key_table_name}\nAND set `hash_key` as 'token'" ) ## Config - table = client.table(DBTableNames.config.value) - if not await table.exists(): - sample_code_snippet = f""" - table = client.table({DBTableNames.config.value}) - await table.create( - self.throughput_type, - KeySchema(hash_key=KeySpec("param_name", KeyType.string)), - ) - """ + try: + error_occurred = False + table = client.table(self.database_arguments.config_table_name) + if not await table.exists(): + await table.create( + self.throughput_type, + KeySchema(hash_key=KeySpec("param_name", KeyType.string)), + ) + except Exception as e: + error_occurred = True + if error_occurred == True: raise Exception( - f"Failed to create table - {DBTableNames.config.value}.\nPlease create a new table called {DBTableNames.config.value}\nAND set `hash_key` as 'param_name'\n\nE.g.: {sample_code_snippet}" + f"Failed to create table - {self.database_arguments.config_table_name}.\nPlease create a new table called {self.database_arguments.config_table_name}\nAND set `hash_key` as 'param_name'" ) async def insert_data( @@ -93,12 +100,12 @@ class DynamoDBWrapper(CustomDB): async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) table = None - if table_name == DBTableNames.user.name: - table = client.table(DBTableNames.user.value) - elif table_name == DBTableNames.key.name: - table = client.table(DBTableNames.key.value) - elif table_name == DBTableNames.config.name: - table = client.table(DBTableNames.config.value) + if table_name == "user": + table = client.table(self.database_arguments.user_table_name) + elif table_name == "key": + table = client.table(self.database_arguments.key_table_name) + elif table_name == "config": + table = client.table(self.database_arguments.config_table_name) for k, v in value.items(): if isinstance(v, datetime): @@ -111,22 +118,22 @@ class DynamoDBWrapper(CustomDB): client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) table = None key_name = None - if table_name == DBTableNames.user.name: - table = client.table(DBTableNames.user.value) + if table_name == "user": + table = client.table(self.database_arguments.user_table_name) key_name = "user_id" - elif table_name == DBTableNames.key.name: - table = client.table(DBTableNames.key.value) + elif table_name == "key": + table = client.table(self.database_arguments.key_table_name) key_name = "token" - elif table_name == DBTableNames.config.name: - table = client.table(DBTableNames.config.value) + elif table_name == "config": + table = client.table(self.database_arguments.config_table_name) key_name = "param_name" response = await table.get_item({key_name: key}) new_response: Any = None - if table_name == DBTableNames.user.name: + if table_name == "user": new_response = LiteLLM_UserTable(**response) - elif table_name == DBTableNames.key.name: + elif table_name == "key": new_response = {} for k, v in response.items(): # handle json string if ( @@ -138,7 +145,7 @@ class DynamoDBWrapper(CustomDB): else: new_response[k] = v new_response = LiteLLM_VerificationToken(**new_response) - elif table_name == DBTableNames.config.name: + elif table_name == "config": new_response = LiteLLM_Config(**response) return new_response @@ -150,20 +157,20 @@ class DynamoDBWrapper(CustomDB): table = None key_name = None try: - if table_name == DBTableNames.user.name: - table = client.table(DBTableNames.user.value) + if table_name == "user": + table = client.table(self.database_arguments.user_table_name) key_name = "user_id" - elif table_name == DBTableNames.key.name: - table = client.table(DBTableNames.key.value) + elif table_name == "key": + table = client.table(self.database_arguments.key_table_name) key_name = "token" - elif table_name == DBTableNames.config.name: - table = client.table(DBTableNames.config.value) + elif table_name == "config": + table = client.table(self.database_arguments.config_table_name) key_name = "param_name" else: raise Exception( - f"Invalid table name. Needs to be one of - {DBTableNames.user.name}, {DBTableNames.key.name}, {DBTableNames.config.name}" + f"Invalid table name. Needs to be one of - {self.database_arguments.user_table_name}, {self.database_arguments.key_table_name}, {self.database_arguments.config_table_name}" ) except Exception as e: raise Exception(f"Error connecting to table - {str(e)}")