From cff9f7fee694d848cec4f18c8d6b0e418fd4eab7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 17 Jan 2024 17:28:23 -0800 Subject: [PATCH] fix(proxy_server.py): handle empty insert_data response --- litellm/proxy/db/dynamo_db.py | 2 +- litellm/proxy/proxy_server.py | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 1ebcf97232..eb1c085286 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -171,7 +171,7 @@ class DynamoDBWrapper(CustomDB): if isinstance(v, datetime): value[k] = v.isoformat() - return await table.put_item(item=value) + return await table.put_item(item=value, return_values=ReturnValues.all_old) async def get_data(self, key: str, table_name: Literal["user", "key", "config"]): from aiodynamo.client import Client diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0647ef0f75..6d74607b9b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1110,10 +1110,15 @@ async def generate_key_helper_fn( user_row = await custom_db_client.insert_data( value=user_data, table_name="user" ) + if user_row is None: + # GET USER ROW + user_row = await custom_db_client.get_data( + key=user_id, table_name="user" + ) ## use default user model list if no key-specific model list provided - if len(user_row["models"]) > 0 and len(key_data["models"]) == 0: # type: ignore - key_data["models"] = user_row["models"] + if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore + key_data["models"] = user_row.models ## CREATE KEY verbose_proxy_logger.debug(f"CustomDBClient: Creating Key={key_data}") await custom_db_client.insert_data(value=key_data, table_name="key")