From 64130e368c8403c865b90683d4eb68748ec3ce01 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 21 May 2024 18:04:17 -0700
Subject: [PATCH] feat - create keys with permissions

---
 litellm/proxy/_types.py       |  1 +
 litellm/proxy/proxy_server.py | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 63d82e7095e..a6640e4d3db 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -157,6 +157,7 @@ class LiteLLMRoutes(enum.Enum):
         "/global/spend/end_users",
         "/global/spend/models",
         "/global/predict/spend/logs",
+        "/global/spend/report",
     ]
 
     public_routes: List = [
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 08c2fee0747..cd9251942b8 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1211,6 +1211,13 @@ async def user_api_key_auth(
                     _has_user_setup_sso()
                     and route in LiteLLMRoutes.sso_only_routes.value
                 ):
+                    pass
+                elif (
+                    route in LiteLLMRoutes.global_spend_tracking_routes.value
+                    and getattr(valid_token, "permissions", None) is not None
+                    and "get_spend_routes" in getattr(valid_token, "permissions", None)
+                ):
+
                     pass
                 else:
                     user_role = "unknown"
@@ -2967,7 +2974,7 @@ async def generate_key_helper_fn(
     organization_id: Optional[str] = None,
     table_name: Optional[Literal["key", "user"]] = None,
 ):
-    global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name
+    global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name, premium_user
 
     if prisma_client is None and custom_db_client is None:
         raise Exception(
@@ -3062,6 +3069,14 @@ async def generate_key_helper_fn(
         if isinstance(saved_token["metadata"], str):
             saved_token["metadata"] = json.loads(saved_token["metadata"])
         if isinstance(saved_token["permissions"], str):
+            if (
+                "get_spend_routes" in saved_token["permissions"]
+                and premium_user != True
+            ):
+                raise Exception(
+                    "get_spend_routes permission is only available for LiteLLM Enterprise users"
+                )
+
             saved_token["permissions"] = json.loads(saved_token["permissions"])
         if isinstance(saved_token["model_max_budget"], str):
             saved_token["model_max_budget"] = json.loads(