From 2873365f020e7f54f16b548ca3f3f282b1943f9b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:30:26 -0800 Subject: [PATCH 1/5] (fix) proxy - always use hashed_token as /key cache key --- litellm/proxy/proxy_server.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4a84847e0d9..bb4f2c7e300 100644 --- a/litellm/proxy/proxy_server.py +++ b/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) From 0ee8e8f081f9dbcbbce882136909a85873bd6903 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:31:04 -0800 Subject: [PATCH 2/5] (fix) test - prisma key budget tracking --- litellm/tests/test_key_generate_prisma.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 98a056730d6..ab490063f5f 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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, } }, From 2c79585e184618a339c77a2f1c4957bfbc5b3c56 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:32:00 -0800 Subject: [PATCH 3/5] (fix) dynamo - use hashed_api keys in table --- litellm/proxy/db/dynamo_db.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 534adbddc94..28bfa442c7e 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/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))) From 8b5e397abe643ebdfc0a3149d323e17ffa985192 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:32:34 -0800 Subject: [PATCH 4/5] (fix) dynamo - use hashed tokens budget tracking --- litellm/tests/test_key_generate_dynamodb.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index be55595fadf..5729c0e87bd 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/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, } }, From 2a60af7164b5770a2db1f249db44cf11ec7cf0e5 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:43:16 -0800 Subject: [PATCH 5/5] (test) key crossing budget --- tests/test_keys.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/test_keys.py b/tests/test_keys.py index a296ef13eb3..283776b96d5 100644 --- a/tests/test_keys.py +++ b/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)