From 0a4320bbfa406c24ad32a420f82152da7bdd7323 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 20 Jun 2025 18:10:06 -0700
Subject: [PATCH] feat(team_endpoints.py): return team member budget on team
 info

allows ui to display this to admin / team member
---
 litellm/proxy/_types.py                       |  6 ++++-
 .../management_endpoints/team_endpoints.py    | 24 ++++++++++++++++---
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 269c77627c8..6ae200b9e2a 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2645,9 +2645,13 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse):
 ##########################################
 
 
+class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
+    team_member_budget_table: Optional[LiteLLM_BudgetTable] = None
+
+
 class TeamInfoResponseObject(TypedDict):
     team_id: str
-    team_info: LiteLLM_TeamTable
+    team_info: TeamInfoResponseObjectTeamTable
     keys: List
     team_memberships: List[LiteLLM_TeamMembership]
 
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index c6637d2aee5..ef7d6f40b1a 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -1606,6 +1606,7 @@ async def team_info(
     --header 'Authorization: Bearer your_api_key_here'
     ```
     """
+    from litellm.proxy._types import TeamInfoResponseObjectTeamTable
     from litellm.proxy.proxy_server import prisma_client
 
     try:
@@ -1674,11 +1675,28 @@ async def team_info(
         )
 
         if isinstance(team_info, dict):
-            _team_info = LiteLLM_TeamTable(**team_info)
+            _team_info = TeamInfoResponseObjectTeamTable(**team_info)
         elif isinstance(team_info, BaseModel):
-            _team_info = LiteLLM_TeamTable(**team_info.model_dump())
+            _team_info = TeamInfoResponseObjectTeamTable(**team_info.model_dump())
         else:
-            _team_info = LiteLLM_TeamTable()
+            _team_info = TeamInfoResponseObjectTeamTable()
+
+        ## GET TEAM BUDGET (if exists) ##
+        team_member_budget_id = (
+            _team_info.metadata.get("team_member_budget_id")
+            if _team_info.metadata is not None
+            else None
+        )
+        if team_member_budget_id is not None:
+            try:
+                team_budget = await prisma_client.db.litellm_budgettable.find_unique(
+                    where={"budget_id": team_member_budget_id}
+                )
+                _team_info.team_member_budget_table = team_budget
+            except Exception:
+                verbose_proxy_logger.info(
+                    f"Team member budget table not found, passed team_member_budget_id={team_member_budget_id}"
+                )
 
         # ## UNFURL 'all-proxy-models' into the team_info.models list ##
         # if llm_router is not None: