From 8ed0ffea541374886f5b85a1619d4572ffff7b49 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 30 Aug 2024 16:22:28 -0700
Subject: [PATCH] fix use existing custom_auth.py

---
 .../proxy/example_config_yaml/custom_auth.py  | 50 ++++++++++++++++---
 1 file changed, 44 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/example_config_yaml/custom_auth.py b/litellm/proxy/example_config_yaml/custom_auth.py
index 2726b1c3d..66c213dcf 100644
--- a/litellm/proxy/example_config_yaml/custom_auth.py
+++ b/litellm/proxy/example_config_yaml/custom_auth.py
@@ -1,14 +1,52 @@
+import os
+
 from fastapi import Request
 
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import GenerateKeyRequest, UserAPIKeyAuth
 
 
 async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
     try:
-        return UserAPIKeyAuth(
-            api_key="best-api-key-ever",
-            user_id="best-user-id-ever",
-            team_id="best-team-id-ever",
-        )
+        modified_master_key = f"{os.getenv('PROXY_MASTER_KEY')}-1234"
+        if api_key == modified_master_key:
+            return UserAPIKeyAuth(api_key=api_key)
+        raise Exception
     except:
         raise Exception
+
+
+async def generate_key_fn(data: GenerateKeyRequest):
+    """
+    Asynchronously decides if a key should be generated or not based on the provided data.
+
+    Args:
+        data (GenerateKeyRequest): The data to be used for decision making.
+
+    Returns:
+        bool: True if a key should be generated, False otherwise.
+    """
+    # decide if a key should be generated or not
+    data_json = data.json()  # type: ignore
+
+    # Unpacking variables
+    team_id = data_json.get("team_id")
+    duration = data_json.get("duration")
+    models = data_json.get("models")
+    aliases = data_json.get("aliases")
+    config = data_json.get("config")
+    spend = data_json.get("spend")
+    user_id = data_json.get("user_id")
+    max_parallel_requests = data_json.get("max_parallel_requests")
+    metadata = data_json.get("metadata")
+    tpm_limit = data_json.get("tpm_limit")
+    rpm_limit = data_json.get("rpm_limit")
+
+    if team_id is not None and len(team_id) > 0:
+        return {
+            "decision": True,
+        }
+    else:
+        return {
+            "decision": True,
+            "message": "This violates LiteLLM Proxy Rules. No team id provided.",
+        }