From d0813fd27f2d2a16ba8be0a3066a66a8e03c2099 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 17 Feb 2024 19:18:41 -0800
Subject: [PATCH] feat(proxy_server.py): return all teams, user is a member of
 in /user/info

https://github.com/BerriAI/litellm/issues/2014
---
 litellm/proxy/proxy_server.py | 11 ++++++++++-
 litellm/proxy/utils.py        | 17 +++++------------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index ebbd91f8d1..ac44195531 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -3981,6 +3981,10 @@ async def user_info(
             user_info = await prisma_client.get_data(user_id=user_id)
         else:
             user_info = None
+        ## GET ALL TEAMS ##
+        teams = await prisma_client.get_data(
+            user_id=user_id, table_name="team", query_type="find_all"
+        )
         ## GET ALL KEYS ##
         keys = await prisma_client.get_data(
             user_id=user_id,
@@ -4004,7 +4008,12 @@ async def user_info(
                 # if using pydantic v1
                 key = key.dict()
             key.pop("token", None)
-        return {"user_id": user_id, "user_info": user_info, "keys": keys}
+        return {
+            "user_id": user_id,
+            "user_info": user_info,
+            "keys": keys,
+            "teams": teams,
+        }
     except Exception as e:
         if isinstance(e, HTTPException):
             raise ProxyException(
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 51e7c3c15c..fb7ff98f94 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -602,7 +602,7 @@ class PrismaClient:
                         status_code=status.HTTP_401_UNAUTHORIZED,
                         detail="Authentication Error: invalid user key - token does not exist",
                     )
-            elif user_id is not None or (
+            elif (user_id is not None and table_name is None) or (
                 table_name is not None and table_name == "user"
             ):
                 if query_type == "find_unique":
@@ -672,17 +672,10 @@ class PrismaClient:
                     response = await self.db.litellm_teamtable.find_unique(
                         where={"team_id": team_id}  # type: ignore
                     )
-                if query_type == "find_all" and team_id is not None:
-                    user_id_values = str(tuple(team_id))
-                    sql_query = f"""
-                    SELECT *
-                    FROM "LiteLLM_TeamTable"
-                    WHERE "team_id" IN {team_id}
-                    """
-
-                    # Execute the raw query
-                    # The asterisk before `team_id` unpacks the list into separate arguments
-                    response = await self.db.query_raw(sql_query)
+                elif query_type == "find_all" and user_id is not None:
+                    response = await self.db.litellm_teamtable.find_many(
+                        where={"members": {"has": user_id}}
+                    )
                 return response
         except Exception as e:
             print_verbose(f"LiteLLM Prisma Client Exception: {e}")