From fd282ea9325da9a1a4ef9e4b932b075a07d48e7c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 20 Apr 2024 11:48:41 -0700
Subject: [PATCH] fix testing fixes

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

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index b697b6e97..ca9926cef 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -87,6 +87,14 @@ class LiteLLMRoutes(enum.Enum):
         "/v2/key/info",
     ]
 
+    sso_only_routes: List = [
+        "/key/generate",
+        "/key/update",
+        "/key/delete",
+        "/global/spend/logs",
+        "/global/predict/spend/logs",
+    ]
+
     management_routes: List = [  # key
         "/key/generate",
         "/key/update",
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 2aab7e453..ebeea120d 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1053,6 +1053,11 @@ async def user_api_key_auth(
                                 status_code=status.HTTP_403_FORBIDDEN,
                                 detail="key not allowed to access this team's info",
                             )
+                elif (
+                    _has_user_setup_sso()
+                    and route in LiteLLMRoutes.sso_only_routes.value
+                ):
+                    pass
                 else:
                     raise Exception(
                         f"Only master key can be used to generate, delete, update info for new keys/users/teams. Route={route}"
@@ -1102,6 +1107,13 @@ async def user_api_key_auth(
                     return UserAPIKeyAuth(
                         api_key=api_key, user_role="proxy_admin", **valid_token_dict
                     )
+                elif (
+                    _has_user_setup_sso()
+                    and route in LiteLLMRoutes.sso_only_routes.value
+                ):
+                    return UserAPIKeyAuth(
+                        api_key=api_key, user_role="app_owner", **valid_token_dict
+                    )
                 else:
                     raise Exception(
                         f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
@@ -6545,7 +6557,7 @@ async def team_member_add(
     existing_team_row.members_with_roles.append(new_member)
 
     complete_team_data = LiteLLM_TeamTable(
-        **existing_team_row.model_dump(),
+        **_get_pydantic_json_dict(existing_team_row),
     )
 
     team_row = await prisma_client.update_data(