scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #37708 from BerriAI/litellm_team_member_budget_no_reset

Browse source 
fix(team): allow no-reset default budgets for team members
This commit is contained in:
ryan-crabbe-berri 2026-08-26 15:14:30 -07:00 • committed by GitHub
commit 52b7bea6f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 223 additions and 132 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.

134
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,138 @@ 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(mock_db_client):
"""
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_db_client.db.litellm_budgettable.update = AsyncMock(
side_effect=lambda where, data: SimpleNamespace(**data)
)
result = 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",
},
)
written = mock_db_client.db.litellm_budgettable.update.call_args.kwargs["data"]
assert written["max_budget"] == 100.0
assert written["budget_duration"] is None
assert written["budget_reset_at"] is None
assert "rpm_limit" not in written
assert "tpm_limit" not in written
assert result["metadata"]["team_member_budget_id"] == "existing_budget_123"
assert "team_member_budget" not in result
assert "team_member_budget_duration" not in result
@pytest.mark.asyncio
async def test_create_team_member_budget_table_explicit_null_duration_does_not_inherit_team_duration(
mock_db_client,
):
"""
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_db_client.db.litellm_budgettable.create = AsyncMock(
side_effect=lambda data: SimpleNamespace(**data)
)
result = 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",
},
)
written = mock_db_client.db.litellm_budgettable.create.call_args.kwargs["data"]
assert written["max_budget"] == 100.0
assert "budget_duration" not in written
assert "budget_reset_at" not in written
assert result["metadata"]["team_member_budget_id"] == written["budget_id"]
assert "team_member_budget" not in result
@pytest.mark.asyncio
async def test_create_team_member_budget_table_inherits_team_duration_when_duration_omitted(
mock_db_client,
):
"""
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_db_client.db.litellm_budgettable.create = AsyncMock(
side_effect=lambda data: SimpleNamespace(**data)
)
result = 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"},
)
written = mock_db_client.db.litellm_budgettable.create.call_args.kwargs["data"]
assert written["budget_duration"] == "30d"
assert written["budget_reset_at"] is not None
assert result["metadata"]["team_member_budget_id"] == written["budget_id"]
@pytest.mark.asyncio
async def test_update_team_with_team_member_budget_duration(
disable_audit_logging_for_mocked_team,
@ -2799,6 +2932,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>
);
}

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

@ -4,7 +4,7 @@ import { fireEvent, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils";
import TeamInfoView from "./TeamInfo";
import TeamInfoView, { type TeamData } from "./TeamInfo";
const authState = vi.hoisted(() => ({ userRole: "Admin" }));
@ -1613,10 +1613,18 @@ describe("TeamInfoView - which team member fields reach the update payload depen
vi.clearAllMocks();
});
const openEditor = async (user: ReturnType<typeof userEvent.setup>) => {
const openEditor = async (
user: ReturnType<typeof userEvent.setup>,
teamMemberBudgetTable: TeamData["team_info"]["team_member_budget_table"] = {
max_budget: 42,
budget_duration: "30d",
tpm_limit: 11,
rpm_limit: 22,
},
) => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
team_member_budget_table: { max_budget: 42, budget_duration: "30d", tpm_limit: 11, rpm_limit: 22 },
team_member_budget_table: teamMemberBudgetTable,
default_team_member_models: ["gpt-4"],
}),
);
@ -1667,6 +1675,46 @@ 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("shows Never resets for a stored member budget whose duration is null", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user, { max_budget: 42, budget_duration: null, tpm_limit: null, rpm_limit: null });
await user.click(screen.getByText("Team Member Settings"));
expect(await screen.findByLabelText("Default Budget Duration")).toHaveTextContent("Never resets");
});
it("omits team_member_budget_duration when the dropdown is left untouched on a team with no member budget", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user, null);
await user.click(screen.getByText("Team Member Settings"));
const durationSelect = await screen.findByLabelText("Default Budget Duration");
expect(durationSelect).toHaveTextContent("Inherit team reset period");
expect(durationSelect).not.toHaveTextContent("Never resets");
await user.type(screen.getByLabelText("Default Budget (USD)"), "100");
const payload = await save(user);
expect(payload.team_member_budget).toBe(100);
expect(JSON.parse(JSON.stringify(payload))).not.toHaveProperty("team_member_budget_duration");
});
it("omits object_permission.search_tools while Search Tool Settings is closed", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user);

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

@ -48,7 +48,7 @@ import { z } from "zod/v4";
import GuardrailsSelect from "./GuardrailsSelect";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import BudgetDurationDropdown from "../common_components/budget_duration_dropdown";
import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown";
import {
computeTeamModelBadges,
normalizeTeamModelSelection,
@ -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,15 @@ 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}
showNeverResets
placeholder="Inherit team reset period"
value={value === null ? NEVER_RESETS_BUDGET_DURATION : value}
onChange={(next) => onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next)}
/>
)}
</FormField>
<FormField
control={form.control}