From cccd577e756445e0f2c5c328ec081e13b79c20ae Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 15 Feb 2024 20:03:32 -0800
Subject: [PATCH] feat(proxy_server.py): expose new permissions field for keys

---
 .pre-commit-config.yaml       |  4 ++--
 litellm/proxy/_types.py       |  4 ++--
 litellm/proxy/proxy_server.py | 14 ++++++--------
 litellm/proxy/schema.prisma   |  1 +
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8978e0d1a..2a84048e0 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,10 +1,10 @@
 repos:
 -   repo: https://github.com/psf/black
-    rev: stable
+    rev: 24.2.0
     hooks:
     - id: black
 -   repo: https://github.com/pycqa/flake8
-    rev: 3.8.4  # The version of flake8 to use
+    rev: 7.0.0  # The version of flake8 to use
     hooks:
     -  id: flake8
        exclude: ^litellm/tests/|^litellm/proxy/proxy_cli.py|^litellm/integrations/|^litellm/proxy/tests/
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 411b61a40..e2f3e696f 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -154,7 +154,7 @@ class GenerateKeyRequest(GenerateRequestBase):
     duration: Optional[str] = None
     aliases: Optional[dict] = {}
     config: Optional[dict] = {}
-    permissions: Optional[dict] = None
+    permissions: Optional[dict] = {}
 
 
 class GenerateKeyResponse(GenerateKeyRequest):
@@ -167,7 +167,7 @@ class GenerateKeyResponse(GenerateKeyRequest):
     def set_model_info(cls, values):
         if values.get("token") is not None:
             values.update({"key": values.get("token")})
-        dict_fields = ["metadata", "aliases", "config"]
+        dict_fields = ["metadata", "aliases", "config", "permissions"]
         for field in dict_fields:
             value = values.get(field)
             if value is not None and isinstance(value, str):
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index e3efb5e15..48271be0b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1567,7 +1567,7 @@ async def generate_key_helper_fn(
     update_key_values: Optional[dict] = None,
     key_alias: Optional[str] = None,
     allowed_cache_controls: Optional[list] = [],
-    permissions: Optional[dict] = None,
+    permissions: Optional[dict] = {},
 ):
     global prisma_client, custom_db_client, user_api_key_cache
 
@@ -1597,11 +1597,9 @@ async def generate_key_helper_fn(
         duration_s = _duration_in_seconds(duration=budget_duration)
         reset_at = datetime.utcnow() + timedelta(seconds=duration_s)
 
-    if permissions is not None and isinstance(permissions, dict):
-        permissions = json.dumps(permissions)  # type: ignore
-
     aliases_json = json.dumps(aliases)
     config_json = json.dumps(config)
+    permissions_json = json.dumps(permissions)
     metadata_json = json.dumps(metadata)
     user_id = user_id or str(uuid.uuid4())
     user_role = user_role or "app_user"
@@ -1644,7 +1642,7 @@ async def generate_key_helper_fn(
             "budget_duration": key_budget_duration,
             "budget_reset_at": key_reset_at,
             "allowed_cache_controls": allowed_cache_controls,
-            "permissions": permissions,
+            "permissions": permissions_json,
         }
         if (
             general_settings.get("allow_user_auth", False) == True
@@ -1824,9 +1822,9 @@ async def initialize(
         user_api_base = api_base
         dynamic_config[user_model]["api_base"] = api_base
     if api_version:
-        os.environ[
-            "AZURE_API_VERSION"
-        ] = api_version  # set this for azure - litellm can read this from the env
+        os.environ["AZURE_API_VERSION"] = (
+            api_version  # set this for azure - litellm can read this from the env
+        )
     if max_tokens:  # model-specific param
         user_max_tokens = max_tokens
         dynamic_config[user_model]["max_tokens"] = max_tokens
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index e04ed580a..a047951dc 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -55,6 +55,7 @@ model LiteLLM_VerificationToken {
     config     Json  @default("{}")
     user_id    String?
     team_id    String?
+    permissions Json @default("{}")
     max_parallel_requests Int?
     metadata   Json  @default("{}")
     tpm_limit     BigInt?