diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d62bfbb7d50..91a953c217e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1672,6 +1672,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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d4488b8a78..3643373be65 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -724,6 +724,7 @@ async def new_team( # noqa: PLR0915 - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member. + - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members. - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members. - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" @@ -934,6 +935,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, @@ -942,6 +944,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 diff --git a/poetry.lock b/poetry.lock index dc258644429..d43cdea726c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3473,15 +3473,15 @@ files = [ [[package]] name = "mcp" -version = "1.25.0" +version = "1.26.0" description = "Model Context Protocol SDK" optional = true python-versions = ">=3.10" groups = ["main"] markers = "python_version >= \"3.10\" and extra == \"proxy\"" files = [ - {file = "mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a"}, - {file = "mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802"}, + {file = "mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca"}, + {file = "mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66"}, ] [package.dependencies] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index c325b0b6fce..366f659bdab 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6441,3 +6441,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"