ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

(fix) proxy + dynamo - add some debugging

Browse Source
This commit is contained in:
ishaan-jaff 2024-01-15 12:03:10 -08:00
parent 61a1211bba
commit 1f62637197
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
litellm/proxy/db/dynamo_db.py
View File

@ -9,6 +9,7 @@ from litellm import get_secret
from typing import Any, List, Literal, Optional, Union
import json
from datetime import datetime
from litellm._logging import verbose_proxy_logger
class DynamoDBWrapper(CustomDB):
@ -71,13 +72,19 @@ class DynamoDBWrapper(CustomDB):
from aiodynamo.http.aiohttp import AIOHTTP
from aiohttp import ClientSession
verbose_proxy_logger.debug("DynamoDB Wrapper - Attempting to connect")
async with ClientSession() as session:
client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
## User
try:
error_occurred = False
verbose_proxy_logger.debug("DynamoDB Wrapper - Creating User Table")
table = client.table(self.database_arguments.user_table_name)
verbose_proxy_logger.debug(f"DynamoDB Wrapper - Created Table, {table}")
if not await table.exists():
verbose_proxy_logger.debug(
f"DynamoDB Wrapper - {table} does not exist"
)
await table.create(
self.throughput_type,
KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
@ -118,6 +125,7 @@ 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'"
)
verbose_proxy_logger.debug("DynamoDB Wrapper - Done connecting()")
async def insert_data(
self, value: Any, table_name: Literal["user", "key", "config"]

6
litellm/proxy/proxy_server.py
View File

@ -1281,7 +1281,9 @@ async def startup_event():
if prisma_client is not None:
await prisma_client.connect()
verbose_proxy_logger.debug(f"custom_db_client client - {custom_db_client}")
if custom_db_client is not None:
verbose_proxy_logger.debug(f"custom_db_client connecting - {custom_db_client}")
await custom_db_client.connect()
if prisma_client is not None and master_key is not None:
@ -1289,7 +1291,9 @@ async def startup_event():
await generate_key_helper_fn(
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
)
verbose_proxy_logger.debug(
f"custom_db_client client - Inserting master key {custom_db_client}. Master_key: {master_key}"
)
if custom_db_client is not None and master_key is not None:
# add master key to db
await generate_key_helper_fn(