diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 46305ac70d6..3e155135d79 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index fcab831e3e3..5bdddfc4c61 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c2849f86584..7b0df5a5622 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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, ) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index f00b2b1c636..a36dc7ff2e3 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/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() diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 96d879295a2..b164dc685f2 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3951,6 +3951,9 @@ export interface Member { role: string; user_id: string | null; user_email?: string | null; + max_budget_in_team?: number | null; + tpm_limit?: number | null; + rpm_limit?: number | null; } export const teamMemberAddCall = async ( @@ -4075,21 +4078,43 @@ export const teamMemberUpdateCall = async ( ) => { try { console.log("Form Values in teamMemberUpdateCall:", formValues); + console.log("Budget value:", formValues.max_budget_in_team); + console.log("TPM limit:", formValues.tpm_limit); + console.log("RPM limit:", formValues.rpm_limit); const url = proxyBaseUrl ? `${proxyBaseUrl}/team/member_update` : `/team/member_update`; + + const requestBody: any = { + team_id: teamId, + role: formValues.role, + user_id: formValues.user_id, + }; + + // Add optional budget and rate limit fields + if (formValues.user_email !== undefined) { + requestBody.user_email = formValues.user_email; + } + if (formValues.max_budget_in_team !== undefined && formValues.max_budget_in_team !== null) { + requestBody.max_budget_in_team = formValues.max_budget_in_team; + } + if (formValues.tpm_limit !== undefined && formValues.tpm_limit !== null) { + requestBody.tpm_limit = formValues.tpm_limit; + } + if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { + requestBody.rpm_limit = formValues.rpm_limit; + } + + console.log("Final request body:", requestBody); + const response = await fetch(url, { method: "POST", headers: { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, - body: JSON.stringify({ - team_id: teamId, - role: formValues.role, - user_id: formValues.user_id, - }), + body: JSON.stringify(requestBody), }); if (!response.ok) { diff --git a/ui/litellm-dashboard/src/components/team/edit_membership.tsx b/ui/litellm-dashboard/src/components/team/edit_membership.tsx index 4a608772470..990e1d1b8c1 100644 --- a/ui/litellm-dashboard/src/components/team/edit_membership.tsx +++ b/ui/litellm-dashboard/src/components/team/edit_membership.tsx @@ -1,8 +1,9 @@ import React, { useState, useEffect } from 'react'; -import { Modal, Form, Input, Select as AntSelect, Button as AntButton, message } from 'antd'; -import { Select, SelectItem } from "@tremor/react"; +import { Modal, Form, Button as AntButton, message } from 'antd'; +import { Select, SelectItem, TextInput } from "@tremor/react"; import { Card, Text } from "@tremor/react"; import NotificationManager from "../molecules/notifications_manager"; +import NumericalInput from "../shared/numerical_input"; interface BaseMember { user_email?: string; @@ -21,10 +22,13 @@ interface ModalConfig { showUserId?: boolean; additionalFields?: Array<{ name: string; - label: string; - type: 'input' | 'select'; + label: string | React.ReactNode; + type: 'input' | 'select' | 'numerical'; options?: Array<{ label: string; value: string }>; rules?: any[]; + step?: number; + min?: number; + placeholder?: string; }>; } @@ -54,11 +58,17 @@ const MemberModal = ({ if (visible) { if (mode === 'edit' && initialData) { // For edit mode, use the initialData values - form.setFieldsValue({ + const formValues = { ...initialData, // Ensure role is set correctly for editing - role: initialData.role || config.defaultRole - }); + role: initialData.role || config.defaultRole, + // Keep numeric values as numbers for NumericalInput components + max_budget_in_team: (initialData as any).max_budget_in_team || null, + tpm_limit: (initialData as any).tpm_limit || null, + rpm_limit: (initialData as any).rpm_limit || null, + }; + console.log("Setting form values:", formValues); + form.setFieldsValue(formValues); } else { // For add mode, reset to defaults form.resetFields(); @@ -71,12 +81,21 @@ const MemberModal = ({ const handleSubmit = async (values: any) => { try { - // Trim string values - const formData = Object.entries(values).reduce((acc, [key, value]) => ({ - ...acc, - [key]: typeof value === 'string' ? value.trim() : value - }), {}) as T; + // Trim string values and clean up form data + const formData = Object.entries(values).reduce((acc, [key, value]) => { + if (typeof value === 'string') { + const trimmedValue = value.trim(); + // For empty strings on optional numeric fields, set to null + if (trimmedValue === '' && (key === 'max_budget_in_team' || key === 'tpm_limit' || key === 'rpm_limit')) { + return { ...acc, [key]: null }; + } + return { ...acc, [key]: trimmedValue }; + } + // For numeric values from NumericalInput, use as-is (already numbers) + return { ...acc, [key]: value }; + }, {}) as T; + console.log("Submitting form data:", formData); onSubmit(formData); form.resetFields(); // message.success(`Successfully ${mode === 'add' ? 'added' : 'updated'} member`); @@ -93,30 +112,39 @@ const MemberModal = ({ const renderField = (field: { name: string; - label: string; - type: 'input' | 'select'; + label: string | React.ReactNode; + type: 'input' | 'select' | 'numerical'; options?: Array<{ label: string; value: string }>; rules?: any[]; + step?: number; + min?: number; + placeholder?: string; }) => { switch (field.type) { case 'input': return ( - { - e.target.value = e.target.value.trim(); - }} + + ); + case 'numerical': + return ( + ); case 'select': return ( - + ); default: return null; @@ -127,7 +155,7 @@ const MemberModal = ({ @@ -147,12 +175,8 @@ const MemberModal = ({ { type: 'email', message: 'Please enter a valid email!' } ]} > - { - e.target.value = e.target.value.trim(); - }} /> )} @@ -169,12 +193,8 @@ const MemberModal = ({ name="user_id" className="mb-4" > - { - e.target.value = e.target.value.trim(); - }} /> )} @@ -196,7 +216,7 @@ const MemberModal = ({ { required: true, message: 'Please select a role!' } ]} > - + {config.additionalFields?.map(field => ( diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx index cce59e0ec20..676f813b801 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -254,7 +254,11 @@ const TeamInfoView: React.FC = ({ user_email: values.user_email, user_id: values.user_id, role: values.role, + max_budget_in_team: values.max_budget_in_team, + tpm_limit: values.tpm_limit, + rpm_limit: values.rpm_limit, }; + console.log("Updating member with values:", member); message.destroy(); // Remove all existing toasts await teamMemberUpdateCall(accessToken, teamId, member); @@ -898,6 +902,53 @@ const TeamInfoView: React.FC = ({ { label: "Admin", value: "admin" }, { label: "User", value: "user" }, ], + additionalFields: [ + { + name: "max_budget_in_team", + label: ( + + Team Member Budget (USD){' '} + + + + + ), + type: "numerical" as const, + step: 0.01, + min: 0, + placeholder: "Budget limit for this member within this team" + }, + { + name: "tpm_limit", + label: ( + + Team Member TPM Limit{' '} + + + + + ), + type: "numerical" as const, + step: 1, + min: 0, + placeholder: "Tokens per minute limit for this member in this team" + }, + { + name: "rpm_limit", + label: ( + + Team Member RPM Limit{' '} + + + + + ), + type: "numerical" as const, + step: 1, + min: 0, + placeholder: "Requests per minute limit for this member in this team" + } + ] }} /> diff --git a/ui/litellm-dashboard/src/components/team/team_member_view.tsx b/ui/litellm-dashboard/src/components/team/team_member_view.tsx index 84a38faa75d..ba39b2112ab 100644 --- a/ui/litellm-dashboard/src/components/team/team_member_view.tsx +++ b/ui/litellm-dashboard/src/components/team/team_member_view.tsx @@ -144,7 +144,15 @@ const TeamMembersComponent: React.FC = ({ icon={PencilAltIcon} size="sm" onClick={() => { - setSelectedEditMember(member); + // Get budget and rate limit data from team membership + const membership = teamData.team_memberships.find(tm => tm.user_id === member.user_id); + const enhancedMember = { + ...member, + max_budget_in_team: membership?.litellm_budget_table?.max_budget || null, + tpm_limit: membership?.litellm_budget_table?.tpm_limit || null, + rpm_limit: membership?.litellm_budget_table?.rpm_limit || null, + }; + setSelectedEditMember(enhancedMember); setIsEditMemberModalVisible(true); }} className="cursor-pointer hover:text-blue-600"