phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy_server.py): handle empty insert_data response

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-17 17:28:23 -08:00
parent 98b83fa780
commit cff9f7fee6
2 changed files with 8 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -171,7 +171,7 @@ class DynamoDBWrapper(CustomDB):
if isinstance(v, datetime): if isinstance(v, datetime):
value[k] = v.isoformat() 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"]): async def get_data(self, key: str, table_name: Literal["user", "key", "config"]):
from aiodynamo.client import Client from aiodynamo.client import Client

9
litellm/proxy/proxy_server.py
View file

@ -1110,10 +1110,15 @@ async def generate_key_helper_fn(
user_row = await custom_db_client.insert_data( user_row = await custom_db_client.insert_data(
value=user_data, table_name="user" 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 ## 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 if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore
key_data["models"] = user_row["models"] key_data["models"] = user_row.models
## CREATE KEY ## CREATE KEY
verbose_proxy_logger.debug(f"CustomDBClient: Creating Key={key_data}") verbose_proxy_logger.debug(f"CustomDBClient: Creating Key={key_data}")
await custom_db_client.insert_data(value=key_data, table_name="key") await custom_db_client.insert_data(value=key_data, table_name="key")