From 5b54bcc712c2d92f9709dc6961d8029ceedd3498 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 12:39:11 -0800 Subject: [PATCH] (feat) spendLogs table DynamoDB --- litellm/proxy/db/dynamo_db.py | 21 ++++++++++++++++++++- litellm/proxy/proxy_server.py | 3 +++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index eb1c08528..83cf6b157 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -131,10 +131,27 @@ class DynamoDBWrapper(CustomDB): raise Exception( 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'" ) + + ## Spend + try: + verbose_proxy_logger.debug("DynamoDB Wrapper - Creating Spend Table") + error_occurred = False + table = client.table(self.database_arguments.spend_table_name) + if not await table.exists(): + await table.create( + self.throughput_type, + KeySchema(hash_key=KeySpec("request_id", KeyType.string)), + ) + except Exception as e: + error_occurred = True + if error_occurred == True: + raise Exception( + 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'" + ) verbose_proxy_logger.debug("DynamoDB Wrapper - Done connecting()") async def insert_data( - self, value: Any, table_name: Literal["user", "key", "config"] + self, value: Any, table_name: Literal["user", "key", "config", "spend"] ): from aiodynamo.client import Client from aiodynamo.credentials import Credentials, StaticCredentials @@ -166,6 +183,8 @@ class DynamoDBWrapper(CustomDB): table = client.table(self.database_arguments.key_table_name) elif table_name == "config": table = client.table(self.database_arguments.config_table_name) + elif table_name == "spend": + table = client.table(self.database_arguments.spend_table_name) for k, v in value.items(): if isinstance(v, datetime): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 10efd495b..32e985113 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -658,6 +658,9 @@ async def update_database( if prisma_client is not None: await prisma_client.insert_data(data=payload, table_name="spend") + elif custom_db_client is not None: + await custom_db_client.insert_data(payload, table_name="spend") + tasks = [] tasks.append(_update_user_db()) tasks.append(_update_key_db())