[Feat] UI - Allow editing team member rpm/tpm limits (#13669)

* show team member tpm/rpm limits

* ui - allow setting team settings

* fix better debugging

* fix types: TeamMemberUpdateRequest

* add _upsert_budget_and_membership

* allow updating team member RPM/TPM in teamMemberUpdateCall

* editing team member rpm/tpm

* UI - fixes for team member component

* fix info

* test_upsert_rpm_only_creates_new_budget

Author: ishaan-jaff

litellm/proxy/_types.py

@@ -2762,11 +2762,22 @@ class TeamMemberDeleteRequest(MemberDeleteRequest):
 class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
     max_budget_in_team: Optional[float] = None
     role: Optional[Literal["admin", "user"]] = None
+    tpm_limit: Optional[int] = Field(
+        default=None,
+        description="Tokens per minute limit for this team member"
+    )
+    rpm_limit: Optional[int] = Field(
+        default=None,
+        description="Requests per minute limit for this team member"
+    )
+    
 
 
 class TeamMemberUpdateResponse(MemberUpdateResponse):
     team_id: str
     max_budget_in_team: Optional[float] = None
+    tpm_limit: Optional[int] = None
+    rpm_limit: Optional[int] = None
 
 
 class TeamModelAddRequest(BaseModel):

litellm/proxy/management_endpoints/common_utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 from litellm.proxy._types import (
     KeyRequestBase,
@@ -56,6 +56,8 @@ async def _upsert_budget_and_membership(
     max_budget: Optional[float],
     existing_budget_id: Optional[str],
     user_api_key_dict: UserAPIKeyAuth,
+    tpm_limit: Optional[int] = None,
+    rpm_limit: Optional[int] = None,
 ):
     """
     Helper function to Create/Update or Delete the budget within the team membership
@@ -66,33 +68,34 @@ async def _upsert_budget_and_membership(
         max_budget: The maximum budget for the team
         existing_budget_id: The ID of the existing budget, if any
         user_api_key_dict: User API Key dictionary containing user information
+        tpm_limit: Tokens per minute limit for the team member
+        rpm_limit: Requests per minute limit for the team member
 
-    If max_budget is None, the user's budget is removed from the team membership.
-    If max_budget exists, a budget is updated or created and linked to the team membership.
+    If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership.
+    If any of these values exist, a budget is updated or created and linked to the team membership.
     """
-    if max_budget is None:
-        # disconnect the budget since max_budget is None
+    if max_budget is None and tpm_limit is None and rpm_limit is None:
+        # disconnect the budget since all limits are None
         await tx.litellm_teammembership.update(
             where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
             data={"litellm_budget_table": {"disconnect": True}},
         )
         return
 
-    if existing_budget_id:
-        # update the existing budget
-        await tx.litellm_budgettable.update(
-            where={"budget_id": existing_budget_id},
-            data={"max_budget": max_budget},
-        )
-        return
-
     # create a new budget
+    create_data: Dict[str, Any] = {
+        "created_by": user_api_key_dict.user_id or "",
+        "updated_by": user_api_key_dict.user_id or "",
+    }
+    if max_budget is not None:
+        create_data["max_budget"] = max_budget
+    if tpm_limit is not None:
+        create_data["tpm_limit"] = tpm_limit
+    if rpm_limit is not None:
+        create_data["rpm_limit"] = rpm_limit
+    
     new_budget = await tx.litellm_budgettable.create(
-        data={
-            "max_budget": max_budget,
-            "created_by": user_api_key_dict.user_id or "",
-            "updated_by": user_api_key_dict.user_id or "",
-        },
+        data=create_data,
         include={"team_membership": True},
     )
     # upsert the team membership with the new/updated budget

litellm/proxy/management_endpoints/team_endpoints.py

@@ -1617,6 +1617,8 @@ async def team_member_update(
             max_budget=data.max_budget_in_team,
             existing_budget_id=identified_budget_id,
             user_api_key_dict=user_api_key_dict,
+            tpm_limit=data.tpm_limit,
+            rpm_limit=data.rpm_limit,
         )
 
     ### update team member role
@@ -1647,6 +1649,8 @@ async def team_member_update(
         user_id=received_user_id,
         user_email=data.user_email,
         max_budget_in_team=data.max_budget_in_team,
+        tpm_limit=data.tpm_limit,
+        rpm_limit=data.rpm_limit,
     )
 
 

tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py

@@ -1,13 +1,13 @@
 # tests/litellm/proxy/common_utils/test_upsert_budget_membership.py
 import types
-import pytest
 from unittest.mock import AsyncMock, MagicMock
 
+import pytest
+
 from litellm.proxy.management_endpoints.common_utils import (
     _upsert_budget_and_membership,
 )
 
-
 # ---------------------------------------------------------------------------
 # Fixtures: a fake Prisma transaction and a fake UserAPIKeyAuth object
 # ---------------------------------------------------------------------------
@@ -63,25 +63,51 @@ async def test_upsert_disconnect(mock_tx, fake_user):
     mock_tx.litellm_teammembership.upsert.assert_not_called()
 
 
-# TEST: existing budget id, update only
+# TEST: existing budget id, creates new budget (current behavior)
 @pytest.mark.asyncio
-async def test_upsert_update_existing(mock_tx, fake_user):
+async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user):
+    """
+    Test that even when existing_budget_id is provided, the function creates a new budget.
+    This reflects the current implementation behavior.
+    """
     await _upsert_budget_and_membership(
         mock_tx,
         team_id="team-2",
         user_id="user-2",
         max_budget=42.0,
-        existing_budget_id="bud-999",
+        existing_budget_id="bud-999",  # This parameter is currently unused
         user_api_key_dict=fake_user,
     )
 
-    mock_tx.litellm_budgettable.update.assert_awaited_once_with(
-        where={"budget_id": "bud-999"},
-        data={"max_budget": 42.0},
+    # Should create a new budget, not update existing
+    mock_tx.litellm_budgettable.create.assert_awaited_once_with(
+        data={
+            "max_budget": 42.0,
+            "created_by": fake_user.user_id,
+            "updated_by": fake_user.user_id,
+        },
+        include={"team_membership": True},
+    )
+
+    # Should upsert team membership with the new budget ID
+    new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
+    mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
+        where={"user_id_team_id": {"user_id": "user-2", "team_id": "team-2"}},
+        data={
+            "create": {
+                "user_id": "user-2",
+                "team_id": "team-2",
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+            "update": {
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+        },
     )
+
+    # Should NOT update existing budget
+    mock_tx.litellm_budgettable.update.assert_not_called()
     mock_tx.litellm_teammembership.update.assert_not_called()
-    mock_tx.litellm_budgettable.create.assert_not_called()
-    mock_tx.litellm_teammembership.upsert.assert_not_called()
 
 
 # TEST: create new budget and link membership
@@ -126,9 +152,13 @@ async def test_upsert_create_and_link(mock_tx, fake_user):
     mock_tx.litellm_budgettable.update.assert_not_called()
 
 
-# TEST: create new budget and link membership, then update
+# TEST: create new budget and link membership, then create another new budget
 @pytest.mark.asyncio
-async def test_upsert_create_then_update(mock_tx, fake_user):
+async def test_upsert_create_then_create_another(mock_tx, fake_user):
+    """
+    Test that multiple calls to _upsert_budget_and_membership create separate budgets,
+    reflecting the current implementation behavior.
+    """
     # FIRST CALL – create new budget and link membership
     await _upsert_budget_and_membership(
         mock_tx,
@@ -146,25 +176,143 @@ async def test_upsert_create_then_update(mock_tx, fake_user):
     mock_tx.litellm_budgettable.create.assert_awaited_once()
     mock_tx.litellm_teammembership.upsert.assert_awaited_once()
 
-    # SECOND CALL – pretend the same membership already exists, and
-    # reset call history so the next assertions are clear
+    # SECOND CALL – reset call history and create another budget
     mock_tx.litellm_budgettable.create.reset_mock()
     mock_tx.litellm_teammembership.upsert.reset_mock()
     mock_tx.litellm_budgettable.update.reset_mock()
 
+    # Set up a new budget ID for the second create call
+    mock_tx.litellm_budgettable.create.return_value = types.SimpleNamespace(budget_id="new-budget-456")
+
     await _upsert_budget_and_membership(
         mock_tx,
         team_id="team-42",
         user_id="user-42",
         max_budget=25.0,                # new limit
-        existing_budget_id=created_bid, # now we say it exists
+        existing_budget_id=created_bid, # this is ignored in current implementation
         user_api_key_dict=fake_user,
     )
 
-    # Now we expect ONLY an update to fire
-    mock_tx.litellm_budgettable.update.assert_awaited_once_with(
-        where={"budget_id": created_bid},
-        data={"max_budget": 25.0},
+    # Should create another new budget (not update existing)
+    mock_tx.litellm_budgettable.create.assert_awaited_once_with(
+        data={
+            "max_budget": 25.0,
+            "created_by": fake_user.user_id,
+            "updated_by": fake_user.user_id,
+        },
+        include={"team_membership": True},
+    )
+
+    # Should upsert team membership with the new budget ID
+    new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
+    mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
+        where={"user_id_team_id": {"user_id": "user-42", "team_id": "team-42"}},
+        data={
+            "create": {
+                "user_id": "user-42",
+                "team_id": "team-42",
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+            "update": {
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+        },
+    )
+
+    # Should NOT call update
+    mock_tx.litellm_budgettable.update.assert_not_called()
+
+
+# TEST: update rpm_limit for member with existing budget_id
+@pytest.mark.asyncio
+async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
+    """
+    Test that updating rpm_limit for a member with an existing budget_id
+    creates a new budget with the new rpm/tpm limits and assigns it to the user.
+    """
+    existing_budget_id = "existing-budget-456"
+    
+    await _upsert_budget_and_membership(
+        mock_tx,
+        team_id="team-rpm-test",
+        user_id="user-rpm-test",
+        max_budget=50.0,
+        existing_budget_id=existing_budget_id,
+        user_api_key_dict=fake_user,
+        tpm_limit=1000,
+        rpm_limit=100,  # updating rpm_limit
+    )
+
+    # Should create a new budget with all the specified limits
+    mock_tx.litellm_budgettable.create.assert_awaited_once_with(
+        data={
+            "max_budget": 50.0,
+            "tpm_limit": 1000,
+            "rpm_limit": 100,
+            "created_by": fake_user.user_id,
+            "updated_by": fake_user.user_id,
+        },
+        include={"team_membership": True},
+    )
+
+    # Should NOT update the existing budget
+    mock_tx.litellm_budgettable.update.assert_not_called()
+
+    # Should upsert team membership with the new budget ID
+    new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
+    mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
+        where={"user_id_team_id": {"user_id": "user-rpm-test", "team_id": "team-rpm-test"}},
+        data={
+            "create": {
+                "user_id": "user-rpm-test",
+                "team_id": "team-rpm-test",
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+            "update": {
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+        },
+    )
+
+
+# TEST: create new budget with only rpm_limit (no max_budget)
+@pytest.mark.asyncio
+async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
+    """
+    Test that setting only rpm_limit creates a new budget with just the rpm_limit.
+    """
+    await _upsert_budget_and_membership(
+        mock_tx,
+        team_id="team-rpm-only",
+        user_id="user-rpm-only", 
+        max_budget=None,
+        existing_budget_id=None,
+        user_api_key_dict=fake_user,
+        rpm_limit=50,
+    )
+
+    # Should create a new budget with only rpm_limit
+    mock_tx.litellm_budgettable.create.assert_awaited_once_with(
+        data={
+            "rpm_limit": 50,
+            "created_by": fake_user.user_id,
+            "updated_by": fake_user.user_id,
+        },
+        include={"team_membership": True},
+    )
+
+    # Should upsert team membership with the new budget ID
+    new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
+    mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
+        where={"user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"}},
+        data={
+            "create": {
+                "user_id": "user-rpm-only",
+                "team_id": "team-rpm-only",
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+            "update": {
+                "litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
+            },
+        },
     )
-    mock_tx.litellm_budgettable.create.assert_not_called()
-    mock_tx.litellm_teammembership.upsert.assert_not_called()