scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(team): allow no-reset default budgets for team members

Browse source 
The Default Budget Duration field in Team Member Settings only offered daily, weekly and monthly, so a team member budget could never be set to never reset. It now uses the shared BudgetDurationDropdown, and /team/update writes an explicitly null duration through to the member budget row along with its reset time.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-08-20 20:36:07 +00:00 • committed by ryan-crabbe-berri
parent e4037978f1
commit 99d4741586
6 changed files with 189 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -14,6 +14,7 @@ import json
import math
import traceback
from collections.abc import Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
@ -494,8 +495,13 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit: int | None = None,
team_member_tpm_limit: int | None = None,
team_member_budget_duration: str | None = None,
explicitly_set_fields: AbstractSet[str] = frozenset(),
) -> dict:
"""Create team member budget table with provided limits"""
"""Create team member budget table with provided limits.
The team's own reset period is only inherited when the caller left the
member duration out, so an explicit null means "never resets".
"""
from litellm.proxy._types import BudgetNewRequest
from litellm.proxy.management_endpoints.budget_management_endpoints import (
new_budget,
@ -509,7 +515,11 @@ class TeamMemberBudgetHandler:
# Create budget request with all provided limits
budget_request: Final = BudgetNewRequest(
budget_id=budget_id,
budget_duration=data.budget_duration or team_member_budget_duration,
budget_duration=(
team_member_budget_duration
if "team_member_budget_duration" in explicitly_set_fields
else data.budget_duration or team_member_budget_duration
),
)
if team_member_budget is not None:
@ -545,8 +555,13 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit: int | None = None,
team_member_tpm_limit: int | None = None,
team_member_budget_duration: str | None = None,
explicitly_set_fields: AbstractSet[str] = frozenset(),
) -> dict:
"""Upsert team member budget table with provided limits"""
"""Upsert team member budget table with provided limits.
A field the caller explicitly sent as null is written as null, so a
team can keep a member budget while dropping its reset period.
"""
from litellm.proxy._types import BudgetNewRequest
from litellm.proxy.management_endpoints.budget_management_endpoints import (
update_budget,
@ -560,14 +575,16 @@ class TeamMemberBudgetHandler:
# Budget exists - create update request with only provided values
budget_request: Final = BudgetNewRequest(budget_id=team_member_budget_id)
if team_member_budget is not None:
if team_member_budget is not None or "team_member_budget" in explicitly_set_fields:
budget_request.max_budget = team_member_budget
if team_member_rpm_limit is not None:
if team_member_rpm_limit is not None or "team_member_rpm_limit" in explicitly_set_fields:
budget_request.rpm_limit = team_member_rpm_limit
if team_member_tpm_limit is not None:
if team_member_tpm_limit is not None or "team_member_tpm_limit" in explicitly_set_fields:
budget_request.tpm_limit = team_member_tpm_limit
if team_member_budget_duration is not None:
if team_member_budget_duration is not None or "team_member_budget_duration" in explicitly_set_fields:
budget_request.budget_duration = team_member_budget_duration
if team_member_budget_duration is None:
budget_request.budget_reset_at = None
budget_row: Final = await _as_budget_write(update_budget)(
budget_obj=budget_request,
@ -593,6 +610,7 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit=team_member_rpm_limit,
team_member_tpm_limit=team_member_tpm_limit,
team_member_budget_duration=team_member_budget_duration,
explicitly_set_fields=explicitly_set_fields,
)
# Remove team member fields from updated_kv
@ -1479,6 +1497,7 @@ async def new_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
explicitly_set_fields=data.model_fields_set,
)
## ADD TO TEAM TABLE
@ -2184,6 +2203,7 @@ async def update_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
explicitly_set_fields=_team_member_fields_in_request,
)
# Backfill team_memberships for members who joined before the
# budget was configured — they won't have a membership row yet.

138
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -2366,6 +2366,7 @@ async def test_update_team_team_member_budget_not_passed_to_db(
team_member_rpm_limit=None,
team_member_tpm_limit=None,
team_member_budget_duration=None,
explicitly_set_fields=frozenset(),
):
# Remove team_member_budget from updated_kv as the real function does
result_kv = updated_kv.copy()
@ -2738,6 +2739,142 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
assert "team_member_budget_duration" not in result
@pytest.mark.asyncio
async def test_upsert_team_member_budget_table_clears_duration_kept_budget():
"""
A request that keeps team_member_budget but explicitly nulls
team_member_budget_duration must clear the reset period and its reset time.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
mock_user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
)
team_table = MagicMock(spec=LiteLLM_TeamTable)
team_table.metadata = {"team_member_budget_id": "existing_budget_123"}
mock_budget_response = MagicMock()
mock_budget_response.budget_id = "existing_budget_123"
with patch(
"litellm.proxy.management_endpoints.budget_management_endpoints.update_budget",
new_callable=AsyncMock,
) as mock_update_budget:
mock_update_budget.return_value = mock_budget_response
await TeamMemberBudgetHandler.upsert_team_member_budget_table(
team_table=team_table,
user_api_key_dict=mock_user_api_key_dict,
updated_kv={
"team_id": "test_team_id",
"team_member_budget": 100.0,
"team_member_budget_duration": None,
},
team_member_budget=100.0,
team_member_budget_duration=None,
explicitly_set_fields={
"team_member_budget",
"team_member_budget_duration",
},
)
budget_request = mock_update_budget.call_args.kwargs["budget_obj"]
assert budget_request.max_budget == 100.0
assert "budget_duration" in budget_request.model_fields_set
assert budget_request.budget_duration is None
assert "budget_reset_at" in budget_request.model_fields_set
assert budget_request.budget_reset_at is None
assert "rpm_limit" not in budget_request.model_fields_set
assert "tpm_limit" not in budget_request.model_fields_set
@pytest.mark.asyncio
async def test_create_team_member_budget_table_explicit_null_duration_does_not_inherit_team_duration():
"""
A first-time member budget with an explicitly null duration must never
reset, even when the team itself has a reset period.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
mock_user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
)
team_table = MagicMock(spec=LiteLLM_TeamTable)
team_table.metadata = {}
team_table.team_alias = "Test Team"
team_table.budget_duration = "30d"
mock_budget_response = MagicMock()
mock_budget_response.budget_id = "new_budget_456"
with patch(
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
new_callable=AsyncMock,
) as mock_new_budget:
mock_new_budget.return_value = mock_budget_response
await TeamMemberBudgetHandler.create_team_member_budget_table(
data=team_table,
new_team_data_json={"team_id": "test_team_id"},
user_api_key_dict=mock_user_api_key_dict,
team_member_budget=100.0,
team_member_budget_duration=None,
explicitly_set_fields={
"team_member_budget",
"team_member_budget_duration",
},
)
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
assert budget_request.max_budget == 100.0
assert budget_request.budget_duration is None
@pytest.mark.asyncio
async def test_create_team_member_budget_table_inherits_team_duration_when_duration_omitted():
"""
Omitting team_member_budget_duration keeps the existing inheritance of the
team's own reset period.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
mock_user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
)
team_table = MagicMock(spec=LiteLLM_TeamTable)
team_table.metadata = {}
team_table.team_alias = "Test Team"
team_table.budget_duration = "30d"
mock_budget_response = MagicMock()
mock_budget_response.budget_id = "new_budget_789"
with patch(
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
new_callable=AsyncMock,
) as mock_new_budget:
mock_new_budget.return_value = mock_budget_response
await TeamMemberBudgetHandler.create_team_member_budget_table(
data=team_table,
new_team_data_json={"team_id": "test_team_id"},
user_api_key_dict=mock_user_api_key_dict,
team_member_budget=100.0,
explicitly_set_fields={"team_member_budget"},
)
assert mock_new_budget.call_args.kwargs["budget_obj"].budget_duration == "30d"
@pytest.mark.asyncio
async def test_update_team_with_team_member_budget_duration(
disable_audit_logging_for_mocked_team,
@ -2799,6 +2936,7 @@ async def test_update_team_with_team_member_budget_duration(
team_member_rpm_limit=None,
team_member_tpm_limit=None,
team_member_budget_duration=None,
explicitly_set_fields=frozenset(),
):
result_kv = updated_kv.copy()
result_kv.pop("team_member_budget", None)

79
ui/litellm-dashboard/src/components/common_components/DurationSelect.test.tsx
View file

@ -1,79 +0,0 @@
import { render, screen } from "@testing-library/react";
import { useState } from "react";
import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
import { describe, it, expect, vi } from "vitest";
import DurationSelect from "./DurationSelect";
describe("DurationSelect", () => {
it("should render", () => {
render(<DurationSelect />);
expect(screen.getByRole("combobox")).toBeInTheDocument();
});
it("should render all three duration options", async () => {
const user = userEvent.setup();
render(<DurationSelect />);
const select = screen.getByRole("combobox");
await user.click(select);
expect(screen.getByText("Daily")).toBeInTheDocument();
expect(screen.getByText("Weekly")).toBeInTheDocument();
expect(screen.getByText("Monthly")).toBeInTheDocument();
const dailyLabel = screen.getByText("Daily");
const dailyOption = dailyLabel.closest('[role="option"]') ?? dailyLabel;
await user.click(dailyOption);
});
it("should apply className prop", () => {
render(<DurationSelect className="test-class" />);
const select = screen.getByRole("combobox");
expect(select.closest(".test-class")).toBeInTheDocument();
});
it("should call onChange when an option is selected", async () => {
const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
const onChange = vi.fn();
render(<DurationSelect onChange={onChange} />);
const select = screen.getByRole("combobox");
await user.click(select);
const dailyLabel = screen.getByText("Daily");
const dailyOption = dailyLabel.closest('[role="option"]') ?? dailyLabel;
await user.click(dailyOption);
expect(onChange).toHaveBeenCalledWith("24h", expect.any(Object));
});
it("should accept and pass value prop to Select", () => {
render(<DurationSelect value="7d" />);
const select = screen.getByRole("combobox");
expect(select).toBeInTheDocument();
});
it.each([
["24h", "Daily"],
["7d", "Weekly"],
["30d", "Monthly"],
])("shows the human label on the trigger for %s", (value, label) => {
render(<DurationSelect value={value} />);
expect(screen.getByRole("combobox")).toHaveTextContent(label);
});
it("shows the human label on the trigger after the user picks an option", async () => {
const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
const Harness = () => {
const [value, setValue] = useState("24h");
return <DurationSelect value={value} onChange={setValue} />;
};
render(<Harness />);
await user.click(screen.getByRole("combobox"));
const monthly = screen.getByText("Monthly");
await user.click(monthly.closest('[role="option"]') ?? monthly);
expect(screen.getByRole("combobox")).toHaveTextContent("Monthly");
});
});

39
ui/litellm-dashboard/src/components/common_components/DurationSelect.tsx
View file

@ -1,39 +0,0 @@
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
interface DurationSelectProps {
className?: string;
value?: string;
onChange?: (value: string, option: { value: string; label: string }) => void;
}
const DURATION_OPTIONS = [
{ value: "24h", label: "Daily" },
{ value: "7d", label: "Weekly" },
{ value: "30d", label: "Monthly" },
];
export default function DurationSelect({ className, value, onChange }: DurationSelectProps) {
return (
<Select
items={DURATION_OPTIONS}
value={value}
onValueChange={(nextValue) => {
const selectedOption = DURATION_OPTIONS.find((option) => option.value === nextValue);
if (selectedOption) {
onChange?.(selectedOption.value, selectedOption);
}
}}
>
<SelectTrigger className={className}>
<SelectValue placeholder="Select duration" />
</SelectTrigger>
<SelectContent>
{DURATION_OPTIONS.map((option) => (
<SelectItem key={option.value} value={option.value}>
{option.label}
</SelectItem>
))}
</SelectContent>
</Select>
);
}

15
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
View file

@ -1667,6 +1667,21 @@ describe("TeamInfoView - which team member fields reach the update payload depen
expect(payload.default_team_member_models).toEqual(["gpt-4"]);
});
it("sends a null team_member_budget_duration when Default Budget Duration is set to never reset", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user);
await user.click(screen.getByText("Team Member Settings"));
await screen.findByLabelText("Default Budget (USD)");
await chooseSelectOption(user, screen.getByLabelText("Default Budget Duration"), "Never resets");
const payload = await save(user);
expect(payload.team_member_budget_duration).toBeNull();
expect(payload.team_member_budget).toBe(42);
expect(JSON.stringify(payload)).toContain('"team_member_budget_duration":null');
});
it("omits object_permission.search_tools while Search Tool Settings is closed", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user);

12
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -64,7 +64,6 @@ import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMeta
import ModelAliasManager from "../common_components/ModelAliasManager";
import AgentSelector from "../agent_management/AgentSelector";
import DeleteResourceModal from "../common_components/DeleteResourceModal";
import DurationSelect from "../common_components/DurationSelect";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
import { unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key";
import GuardrailSettingsView from "../GuardrailSettingsView";
@ -169,7 +168,7 @@ export interface TeamData {
object_permission?: ObjectPermission | null;
team_member_budget_table: {
max_budget: number;
budget_duration: string;
budget_duration: string | null;
tpm_limit: number | null;
rpm_limit: number | null;
} | null;
@ -1256,7 +1255,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
name="team_member_budget_duration"
label="Default Budget Duration"
>
{({ value, onChange }) => <DurationSelect value={value ?? undefined} onChange={onChange} />}
{({ id, value, onChange }) => (
<BudgetDurationDropdown
id={id}
placeholder="Never resets"
value={value}
onChange={(next) => onChange(next ?? null)}
/>
)}
</FormField>
<FormField
control={form.control}