phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(dynamo_db.py): don't auto-create tables, allow database_type == 'dynamodb'

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-12 11:33:40 +05:30
parent aeff32ac62
commit 352f943dcf
2 changed files with 33 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -51,39 +51,41 @@ class DynamoDBWrapper(CustomDB):
## User
table = client.table(DBTableNames.user.value)
if not await table.exists():
try:
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
)
except:
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'"
)
sample_code_snippet = f"""
table = client.table({DBTableNames.user.value})
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
)
"""
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}"
)
## Token
table = client.table(DBTableNames.key.value)
if not await table.exists():
try:
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("token", KeyType.string)),
)
except:
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'"
)
sample_code_snippet = f"""
table = client.table({DBTableNames.key.value})
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("token", KeyType.string)),
)
"""
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}"
)
## Config
table = client.table(DBTableNames.config.value)
if not await table.exists():
try:
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("param_name", KeyType.string)),
)
except:
raise Exception(
f"Failed to create table - {DBTableNames.config.value}.\nPlease create a new table called {DBTableNames.config.value}\nAND set `hash_key` as 'token'"
)
sample_code_snippet = f"""
table = client.table({DBTableNames.config.value})
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("param_name", KeyType.string)),
)
"""
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}"
)
async def insert_data(
self, value: Any, table_name: Literal["user", "key", "config"]

4
litellm/proxy/proxy_server.py
View file

@ -853,7 +853,9 @@ class ProxyConfig:
)
## dynamodb
database_type = general_settings.get("database_type", None)
if database_type is not None and database_type == "dynamo_db":
if database_type is not None and (
database_type == "dynamo_db" or database_type == "dynamodb"
):
database_args = general_settings.get("database_args", None)
custom_db_client = DBClient(
custom_db_args=database_args, custom_db_type=database_type