mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-23 18:25:22 +00:00
fix(dynamo_db.py): allow user to specify table names
This commit is contained in:
Krrish Dholakia
+17
-9
@@ -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]
|
||||
|
||||
@@ -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)}")
|
||||
|
||||
Reference in New Issue
Block a user