From 96c9fd3f02d1015af004f34b60d29c860ea6a7df Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 29 Jan 2024 16:52:31 -0800 Subject: [PATCH] (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!")