Merge pull request #1676 from BerriAI/litellm_patch_dynamo_key

[Fix]Litellm patch dynamoDB

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

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!")

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)}")