scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

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

Browse source 
* 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
This commit is contained in:
Ishaan Jaff 2025-08-15 17:29:44 -07:00 • committed by GitHub
parent 9f17bedb48
commit 40395598da
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 355 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/_types.py
View file

@ -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):

39
litellm/proxy/management_endpoints/common_utils.py
View file

@ -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

4
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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,
)

190
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View file

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

35
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -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) {

100
ui/litellm-dashboard/src/components/team/edit_membership.tsx
View file

@ -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 = <T extends BaseMember>({
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 = <T extends BaseMember>({
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 = <T extends BaseMember>({
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 (
<Input
className="px-3 py-2 border rounded-md w-full"
onChange={(e) => {
e.target.value = e.target.value.trim();
}}
<TextInput
placeholder={field.placeholder}
/>
);
case 'numerical':
return (
<NumericalInput
step={field.step || 1}
min={field.min || 0}
style={{ width: "100%" }}
placeholder={field.placeholder || "Enter a numerical value"}
/>
);
case 'select':
return (
<AntSelect>
<Select>
{field.options?.map(option => (
<AntSelect.Option key={option.value} value={option.value}>
<SelectItem key={option.value} value={option.value}>
{option.label}
</AntSelect.Option>
</SelectItem>
))}
</AntSelect>
</Select>
);
default:
return null;
@ -127,7 +155,7 @@ const MemberModal = <T extends BaseMember>({
<Modal
title={config.title || (mode === 'add' ? "Add Member" : "Edit Member")}
open={visible}
width={800}
width={1000}
footer={null}
onCancel={onCancel}
>
@ -147,12 +175,8 @@ const MemberModal = <T extends BaseMember>({
{ type: 'email', message: 'Please enter a valid email!' }
]}
>
<Input
className="px-3 py-2 border rounded-md w-full"
<TextInput
placeholder="user@example.com"
onChange={(e) => {
e.target.value = e.target.value.trim();
}}
/>
</Form.Item>
)}
@ -169,12 +193,8 @@ const MemberModal = <T extends BaseMember>({
name="user_id"
className="mb-4"
>
<Input
className="px-3 py-2 border rounded-md w-full"
<TextInput
placeholder="user_123"
onChange={(e) => {
e.target.value = e.target.value.trim();
}}
/>
</Form.Item>
)}
@ -196,7 +216,7 @@ const MemberModal = <T extends BaseMember>({
{ required: true, message: 'Please select a role!' }
]}
>
<AntSelect>
<Select>
{mode === 'edit' && initialData
? [
// Current role first
@ -204,16 +224,16 @@ const MemberModal = <T extends BaseMember>({
// Then all other roles
...config.roleOptions.filter(option => option.value !== initialData.role)
].map(option => (
<AntSelect.Option key={option.value} value={option.value}>
<SelectItem key={option.value} value={option.value}>
{option.label}
</AntSelect.Option>
</SelectItem>
))
: config.roleOptions.map(option => (
<AntSelect.Option key={option.value} value={option.value}>
<SelectItem key={option.value} value={option.value}>
{option.label}
</AntSelect.Option>
</SelectItem>
))}
</AntSelect>
</Select>
</Form.Item>
{config.additionalFields?.map(field => (

51
ui/litellm-dashboard/src/components/team/team_info.tsx
View file

@ -254,7 +254,11 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
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<TeamInfoProps> = ({
{ label: "Admin", value: "admin" },
{ label: "User", value: "user" },
],
additionalFields: [
{
name: "max_budget_in_team",
label: (
<span>
Team Member Budget (USD){' '}
<Tooltip title="Maximum amount in USD this member can spend within this team. This is separate from any global user budget limits">
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</Tooltip>
</span>
),
type: "numerical" as const,
step: 0.01,
min: 0,
placeholder: "Budget limit for this member within this team"
},
{
name: "tpm_limit",
label: (
<span>
Team Member TPM Limit{' '}
<Tooltip title="Maximum tokens per minute this member can use within this team. This is separate from any global user TPM limit">
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</Tooltip>
</span>
),
type: "numerical" as const,
step: 1,
min: 0,
placeholder: "Tokens per minute limit for this member in this team"
},
{
name: "rpm_limit",
label: (
<span>
Team Member RPM Limit{' '}
<Tooltip title="Maximum requests per minute this member can make within this team. This is separate from any global user RPM limit">
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</Tooltip>
</span>
),
type: "numerical" as const,
step: 1,
min: 0,
placeholder: "Requests per minute limit for this member in this team"
}
]
}}
/>

10
ui/litellm-dashboard/src/components/team/team_member_view.tsx
View file

@ -144,7 +144,15 @@ const TeamMembersComponent: React.FC<TeamMembersComponentProps> = ({
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"