Merge pull request #1640 from BerriAI/litellm_fix_key_gen_budget_crossed

[FIX] Fixes Bug where Keys could cross their Key max_budget

litellm/proxy/db/dynamo_db.py

@@ -5,6 +5,7 @@ from litellm.proxy._types import (
     LiteLLM_Config,
     LiteLLM_UserTable,
 )
+from litellm.proxy.utils import hash_token
 from litellm import get_secret
 from typing import Any, List, Literal, Optional, Union
 import json
@@ -187,6 +188,8 @@ class DynamoDBWrapper(CustomDB):
                 table = client.table(self.database_arguments.spend_table_name)
 
             for k, v in value.items():
+                if k == "token" and value[k].startswith("sk-"):
+                    value[k] = hash_token(token=v)
                 if isinstance(v, datetime):
                     value[k] = v.isoformat()
 
@@ -229,6 +232,10 @@ class DynamoDBWrapper(CustomDB):
                 table = client.table(self.database_arguments.config_table_name)
                 key_name = "param_name"
 
+            if key_name == "token" and key.startswith("sk-"):
+                # ensure it's hashed
+                key = hash_token(token=key)
+
             response = await table.get_item({key_name: key})
 
             new_response: Any = None
@@ -308,6 +315,8 @@ class DynamoDBWrapper(CustomDB):
                 # Convert datetime object to ISO8601 string
                 if isinstance(v, datetime):
                     v = v.isoformat()
+                if k == "token" and value[k].startswith("sk-"):
+                    value[k] = hash_token(token=v)
 
                 # Accumulate updates
                 actions.append((F(k), Value(value=v)))

litellm/proxy/proxy_server.py

@@ -75,6 +75,7 @@ from litellm.proxy.utils import (
     send_email,
     get_logging_payload,
     reset_budget,
+    hash_token,
 )
 from litellm.proxy.secret_managers.google_kms import load_google_kms
 import pydantic
@@ -288,8 +289,9 @@ async def user_api_key_auth(
             raise Exception("No connected db.")
 
         ## check for cache hit (In-Memory Cache)
+        if api_key.startswith("sk-"):
+            api_key = hash_token(token=api_key)
         valid_token = user_api_key_cache.get_cache(key=api_key)
-        verbose_proxy_logger.debug(f"valid_token from cache: {valid_token}")
         if valid_token is None:
             ## check db
             verbose_proxy_logger.debug(f"api key: {api_key}")
@@ -482,10 +484,10 @@ async def user_api_key_auth(
                     )
 
             # Token passed all checks
-            # Add token to cache
-            user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
-
             api_key = valid_token.token
+
+            # Add hashed token to cache
+            user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
             valid_token_dict = _get_pydantic_json_dict(valid_token)
             valid_token_dict.pop("token", None)
             """
@@ -748,6 +750,9 @@ async def update_database(
 
         ### UPDATE KEY SPEND ###
         async def _update_key_db():
+            verbose_proxy_logger.debug(
+                f"adding spend to key db. Response cost: {response_cost}. Token: {token}."
+            )
             if prisma_client is not None:
                 # Fetch the existing cost for the given token
                 existing_spend_obj = await prisma_client.get_data(token=token)

litellm/tests/test_key_generate_dynamodb.py

@@ -33,7 +33,7 @@ from litellm.proxy.proxy_server import (
 )
 
 from litellm.proxy._types import NewUserRequest, DynamoDBArgs, GenerateKeyRequest
-from litellm.proxy.utils import DBClient
+from litellm.proxy.utils import DBClient, hash_token
 from starlette.datastructures import URL
 
 
@@ -232,7 +232,7 @@ def test_call_with_user_over_budget(custom_db_client):
                     "stream": False,
                     "litellm_params": {
                         "metadata": {
-                            "user_api_key": generated_key,
+                            "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
                         }
                     },
@@ -305,7 +305,7 @@ def test_call_with_user_over_budget_stream(custom_db_client):
                     "complete_streaming_response": resp,
                     "litellm_params": {
                         "metadata": {
-                            "user_api_key": generated_key,
+                            "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
                         }
                     },
@@ -376,7 +376,7 @@ def test_call_with_user_key_budget(custom_db_client):
                     "stream": False,
                     "litellm_params": {
                         "metadata": {
-                            "user_api_key": generated_key,
+                            "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
                         }
                     },
@@ -449,7 +449,7 @@ def test_call_with_key_over_budget_stream(custom_db_client):
                     "complete_streaming_response": resp,
                     "litellm_params": {
                         "metadata": {
-                            "user_api_key": generated_key,
+                            "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
                         }
                     },

litellm/tests/test_key_generate_prisma.py

@@ -46,7 +46,7 @@ from litellm.proxy.proxy_server import (
     spend_key_fn,
     view_spend_logs,
 )
-from litellm.proxy.utils import PrismaClient, ProxyLogging
+from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
 from litellm._logging import verbose_proxy_logger
 
 verbose_proxy_logger.setLevel(level=logging.DEBUG)
@@ -918,7 +918,7 @@ def test_call_with_key_over_budget(prisma_client):
                     "stream": False,
                     "litellm_params": {
                         "metadata": {
-                            "user_api_key": generated_key,
+                            "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
                         }
                     },
@@ -1009,7 +1009,7 @@ async def test_call_with_key_never_over_budget(prisma_client):
                 "stream": False,
                 "litellm_params": {
                     "metadata": {
-                        "user_api_key": generated_key,
+                        "user_api_key": hash_token(generated_key),
                         "user_api_key_user_id": user_id,
                     }
                 },
@@ -1083,7 +1083,7 @@ async def test_call_with_key_over_budget_stream(prisma_client):
                 "complete_streaming_response": resp,
                 "litellm_params": {
                     "metadata": {
-                        "user_api_key": generated_key,
+                        "user_api_key": hash_token(generated_key),
                         "user_api_key_user_id": user_id,
                     }
                 },

tests/test_keys.py

@@ -115,7 +115,9 @@ async def chat_completion(session, key, model="gpt-4"):
         print()
 
         if status != 200:
-            raise Exception(f"Request did not return a 200 status code: {status}")
+            raise Exception(
+                f"Request did not return a 200 status code: {status}. Response: {response_text}"
+            )
 
         return await response.json()
 
@@ -386,3 +388,31 @@ async def test_key_with_budgets():
         key_info = await get_key_info(session=session, get_key=key, call_key=key)
         reset_at_new_value = key_info["info"]["budget_reset_at"]
         assert reset_at_init_value != reset_at_new_value
+
+
+@pytest.mark.asyncio
+async def test_key_crossing_budget():
+    """
+    - Create key with budget with budget=0.00000001
+    - make a /chat/completions call
+    - wait 5s
+    - make a /chat/completions call - should fail with key crossed it's budget
+
+    - Check if value updated
+    """
+    from litellm.proxy.utils import hash_token
+
+    async with aiohttp.ClientSession() as session:
+        key_gen = await generate_key(session=session, i=0, budget=0.0000001)
+        key = key_gen["key"]
+        hashed_token = hash_token(token=key)
+        print(f"hashed_token: {hashed_token}")
+
+        response = await chat_completion(session=session, key=key)
+        print("response 1: ", response)
+        await asyncio.sleep(2)
+        try:
+            response = await chat_completion(session=session, key=key)
+            pytest.fail("Should have failed - Key crossed it's budget")
+        except Exception as e:
+            assert "ExceededTokenBudget: Current spend for token:" in str(e)