From 96c9fd3f02d1015af004f34b60d29c860ea6a7df Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 29 Jan 2024 16:52:31 -0800 Subject: [PATCH 1/2] (fix) patch -For DynamoDB Backwards Compatibility --- litellm/proxy/db/dynamo_db.py | 4 ---- litellm/proxy/proxy_server.py | 14 ++++++++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index b570aa0ca88..b8d2b09ca4e 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -233,10 +233,6 @@ 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dd9b509d984..1c2f7509b58 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -292,6 +292,7 @@ async def user_api_key_auth( raise Exception("No connected db.") ## check for cache hit (In-Memory Cache) + original_api_key = api_key # (Patch: For DynamoDB Backwards Compatibility) if api_key.startswith("sk-"): api_key = hash_token(token=api_key) valid_token = user_api_key_cache.get_cache(key=api_key) @@ -304,10 +305,15 @@ async def user_api_key_auth( ) elif custom_db_client is not None: - valid_token = await custom_db_client.get_data( - key=api_key, table_name="key" - ) - + try: + valid_token = await custom_db_client.get_data( + key=api_key, table_name="key" + ) + except: + # (Patch: For DynamoDB Backwards Compatibility) + valid_token = await custom_db_client.get_data( + key=original_api_key, table_name="key" + ) verbose_proxy_logger.debug(f"Token from db: {valid_token}") elif valid_token is not None: verbose_proxy_logger.debug(f"API Key Cache Hit!") From e2a79101134b622d37f37e403fde93937d8d389b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 29 Jan 2024 16:53:27 -0800 Subject: [PATCH 2/2] (test) dynamoDB backwards compat --- litellm/tests/test_key_generate_dynamodb.py | 28 +++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index f482f35437f..4a715b34232 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -472,3 +472,31 @@ def test_call_with_key_over_budget_stream(custom_db_client): error_detail = e.message assert "Authentication Error, ExceededTokenBudget:" in error_detail print(vars(e)) + + +def test_dynamo_db_migration(custom_db_client): + # Tests the temporary patch we have in place + setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + bearer_token = ( + "Bearer " + "sk-elJDL2pOEjcAuC7zD4psAg" + ) # this works with ishaan's db, it's a never expiring key + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return b'{"model": "azure-models"}' + + request.body = return_body + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print("result from user auth with new key", result) + + asyncio.run(test()) + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}")