phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy_server.py): return all teams, user is a member of in /user/info

Browse source 
https://github.com/BerriAI/litellm/issues/2014
This commit is contained in:
Krrish Dholakia 2024-02-17 19:18:41 -08:00
parent 045d84e167
commit d0813fd27f
2 changed files with 15 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/proxy_server.py
View file

@ -3981,6 +3981,10 @@ async def user_info(
user_info = await prisma_client.get_data(user_id=user_id) user_info = await prisma_client.get_data(user_id=user_id)
else: else:
user_info = None 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 ## ## GET ALL KEYS ##
keys = await prisma_client.get_data( keys = await prisma_client.get_data(
user_id=user_id, user_id=user_id,
@ -4004,7 +4008,12 @@ async def user_info(
# if using pydantic v1 # if using pydantic v1
key = key.dict() key = key.dict()
key.pop("token", None) 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: except Exception as e:
if isinstance(e, HTTPException): if isinstance(e, HTTPException):
raise ProxyException( raise ProxyException(

17
litellm/proxy/utils.py
View file

@ -602,7 +602,7 @@ class PrismaClient:
status_code=status.HTTP_401_UNAUTHORIZED, status_code=status.HTTP_401_UNAUTHORIZED,
detail="Authentication Error: invalid user key - token does not exist", 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" table_name is not None and table_name == "user"
): ):
if query_type == "find_unique": if query_type == "find_unique":
@ -672,17 +672,10 @@ class PrismaClient:
response = await self.db.litellm_teamtable.find_unique( response = await self.db.litellm_teamtable.find_unique(
where={"team_id": team_id} # type: ignore where={"team_id": team_id} # type: ignore
) )
if query_type == "find_all" and team_id is not None: elif query_type == "find_all" and user_id is not None:
user_id_values = str(tuple(team_id)) response = await self.db.litellm_teamtable.find_many(
sql_query = f""" where={"members": {"has": user_id}}
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)
return response return response
except Exception as e: except Exception as e:
print_verbose(f"LiteLLM Prisma Client Exception: {e}") print_verbose(f"LiteLLM Prisma Client Exception: {e}")