fix(proxy/utils.py): fix resetting budget logic

litellm/proxy/proxy_server.py

@@ -3464,7 +3464,7 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest):
         response = await prisma_client.update_data(
             token=key, data={**non_default_values, "token": key}
         )
-        return {"key": key, **non_default_values}
+        return {"key": key, **response["data"]}
         # update based on remaining passed in values
     except Exception as e:
         if isinstance(e, HTTPException):

litellm/proxy/utils.py

@@ -1140,7 +1140,13 @@ class PrismaClient:
                     + f"DB Token Table update succeeded {response}"
                     + "\033[0m"
                 )
-                return {"token": token, "data": db_data}
+                _data: dict = {}
+                if response is not None:
+                    try:
+                        _data = response.model_dump()
+                    except:
+                        _data = response.dict()
+                return {"token": token, "data": _data}
             elif (
                 user_id is not None
                 or (table_name is not None and table_name == "user")
@@ -1228,9 +1234,11 @@ class PrismaClient:
                     if t.token.startswith("sk-"):  # type: ignore
                         t.token = self.hash_token(token=t.token)  # type: ignore
                     try:
-                        data_json = self.jsonify_object(data=t.model_dump())
+                        data_json = self.jsonify_object(
+                            data=t.model_dump(exclude_none=True)
+                        )
                     except:
-                        data_json = self.jsonify_object(data=t.dict())
+                        data_json = self.jsonify_object(data=t.dict(exclude_none=True))
                     batcher.litellm_verificationtoken.update(
                         where={"token": t.token},  # type: ignore
                         data={**data_json},  # type: ignore

proxy_server_config.yaml

@@ -40,8 +40,8 @@ litellm_settings:
   budget_duration: 30d
 general_settings: 
   master_key: sk-1234 # [OPTIONAL] Only use this if you to require all calls to contain this key (Authorization: Bearer sk-1234)
-  proxy_budget_rescheduler_min_time: 30
-  proxy_budget_rescheduler_max_time: 60
+  proxy_budget_rescheduler_min_time: 3
+  proxy_budget_rescheduler_max_time: 6
   # database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy
 
 environment_variables:

tests/test_keys.py

@@ -438,6 +438,15 @@ async def test_key_with_budgets():
     """
     from litellm.proxy.utils import hash_token
 
+    async def retry_request(func, *args, _max_attempts=5, **kwargs):
+        for attempt in range(_max_attempts):
+            try:
+                return await func(*args, **kwargs)
+            except aiohttp.client_exceptions.ClientOSError as e:
+                if attempt + 1 == _max_attempts:
+                    raise  # re-raise the last ClientOSError if all attempts failed
+                print(f"Attempt {attempt+1} failed, retrying...")
+
     async with aiohttp.ClientSession() as session:
         key_gen = await generate_key(
             session=session, i=0, budget=10, budget_duration="5s"
@@ -449,9 +458,11 @@ async def test_key_with_budgets():
         reset_at_init_value = key_info["info"]["budget_reset_at"]
         reset_at_new_value = None
         i = 0
-        await asyncio.sleep(120)
-        while i < 3:
-            key_info = await get_key_info(session=session, get_key=key, call_key=key)
+        await asyncio.sleep(10)
+        for i in range(3):
+            key_info = await retry_request(
+                get_key_info, session=session, get_key=key, call_key=key
+            )
             reset_at_new_value = key_info["info"]["budget_reset_at"]
             try:
                 assert reset_at_init_value != reset_at_new_value