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

feat(proxy_server.py): enable /team/info endpoint

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-14 18:12:18 -08:00
parent c094db7160
commit 8717ee6d9a
3 changed files with 108 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/proxy/utils.py
View file

@ -509,8 +509,9 @@ class PrismaClient:
token: Optional[Union[str, list]] = None,
user_id: Optional[str] = None,
user_id_list: Optional[list] = None,
team_id: Optional[str] = None,
key_val: Optional[dict] = None,
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
query_type: Literal["find_unique", "find_all"] = "find_unique",
expires: Optional[datetime] = None,
reset_at: Optional[datetime] = None,
@ -545,6 +546,14 @@ class PrismaClient:
for r in response:
if isinstance(r.expires, datetime):
r.expires = r.expires.isoformat()
elif query_type == "find_all" and team_id is not None:
response = await self.db.litellm_verificationtoken.find_many(
where={"team_id": team_id}
)
if response is not None and len(response) > 0:
for r in response:
if isinstance(r.expires, datetime):
r.expires = r.expires.isoformat()
elif (
query_type == "find_all"
and expires is not None
@ -657,7 +666,23 @@ class PrismaClient:
order={"startTime": "desc"},
)
return response
elif table_name == "team":
if query_type == "find_unique":
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)
return response
except Exception as e:
print_verbose(f"LiteLLM Prisma Client Exception: {e}")
import traceback