mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(proxy): add team_member_budget_duration to NewTeamRequest
NewTeamRequest was missing the team_member_budget_duration field, causing Pydantic to silently drop the value when creating a team via POST /team/new. The template budget row was created without budget_duration or budget_reset_at, so the ResetBudgetJob never found it and team member spend was never reset. Add the field to NewTeamRequest and pass it through to should_create_budget and create_team_member_budget_table in the new_team handler (matching the existing update_team path which already works correctly). Fixes #16057
This commit is contained in:
parent
89d8401d72
commit
78927138d2
3 changed files with 53 additions and 0 deletions
|
|
@ -1669,6 +1669,7 @@ class NewTeamRequest(TeamBase):
|
|||
int
|
||||
] = None # allow user to set TPM limit for all team members
|
||||
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
|
||||
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
|
||||
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
|
||||
enforced_batch_output_expires_after: Optional[dict] = None
|
||||
enforced_file_expires_after: Optional[dict] = None
|
||||
|
|
|
|||
|
|
@ -898,6 +898,7 @@ async def new_team( # noqa: PLR0915
|
|||
team_member_budget=data.team_member_budget,
|
||||
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,
|
||||
):
|
||||
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=data,
|
||||
|
|
@ -906,6 +907,7 @@ async def new_team( # noqa: PLR0915
|
|||
team_member_budget=data.team_member_budget,
|
||||
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,
|
||||
)
|
||||
|
||||
## ADD TO TEAM TABLE
|
||||
|
|
|
|||
|
|
@ -6157,3 +6157,53 @@ async def test_list_team_v1_batches_key_queries():
|
|||
assert result[0].keys == [key1, key2]
|
||||
assert result[1].team_id == "team-2"
|
||||
assert result[1].keys == [key3]
|
||||
|
||||
|
||||
def test_new_team_request_accepts_team_member_budget_duration():
|
||||
"""Test that NewTeamRequest does not silently drop team_member_budget_duration."""
|
||||
from litellm.proxy._types import NewTeamRequest
|
||||
|
||||
request = NewTeamRequest(
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
assert request.team_member_budget == 20.0
|
||||
assert request.team_member_budget_duration == "30d"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_team_member_budget_table_with_duration():
|
||||
"""Verify that create_team_member_budget_table passes budget_duration
|
||||
through to the new_budget call when team_member_budget_duration is provided."""
|
||||
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
mock_budget_response = MagicMock(budget_id="budget-abc")
|
||||
mock_admin = UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
data = NewTeamRequest(
|
||||
team_alias="test-team",
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_budget_response,
|
||||
) as mock_new_budget:
|
||||
result = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=data,
|
||||
new_team_data_json={"metadata": None},
|
||||
user_api_key_dict=mock_admin,
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
|
||||
mock_new_budget.assert_awaited_once()
|
||||
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
|
||||
assert budget_request.budget_duration == "30d"
|
||||
assert budget_request.max_budget == 20.0
|
||||
assert result["metadata"]["team_member_budget_id"] == "budget-abc"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue