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")