phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

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)
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(

17
litellm/proxy/utils.py
View file

@ -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}")