From 2a6f2a3fb8ba807feaec4f963a1a17299b4e93b2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 6 Jan 2026 17:43:17 -0800 Subject: [PATCH] add team member budget duration in team/update --- litellm/proxy/_types.py | 1 + .../management_endpoints/team_endpoints.py | 15 +- .../test_team_endpoints.py | 366 +++++++++++++++++- 3 files changed, 380 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 954c26e2cb..cff3bee1ca 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1532,6 +1532,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): guardrails: Optional[List[str]] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None team_member_budget: Optional[float] = None + team_member_budget_duration: Optional[str] = None team_member_rpm_limit: Optional[int] = None team_member_tpm_limit: Optional[int] = None team_member_key_duration: Optional[str] = None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 76c607f5c4..920105edc1 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -112,6 +112,7 @@ class TeamMemberBudgetHandler: team_member_budget: Optional[float] = None, team_member_rpm_limit: Optional[int] = None, team_member_tpm_limit: Optional[int] = None, + team_member_budget_duration: Optional[str] = None, ) -> bool: """Check if any team member limits are provided""" return any( @@ -119,6 +120,7 @@ class TeamMemberBudgetHandler: team_member_budget is not None, team_member_rpm_limit is not None, team_member_tpm_limit is not None, + team_member_budget_duration is not None, ] ) @@ -130,6 +132,7 @@ class TeamMemberBudgetHandler: team_member_budget: Optional[float] = None, team_member_rpm_limit: Optional[int] = None, team_member_tpm_limit: Optional[int] = None, + team_member_budget_duration: Optional[str] = None, ) -> dict: """Create team member budget table with provided limits""" from litellm.proxy._types import BudgetNewRequest @@ -147,7 +150,7 @@ class TeamMemberBudgetHandler: # Create budget request with all provided limits budget_request = BudgetNewRequest( budget_id=budget_id, - budget_duration=data.budget_duration, + budget_duration=data.budget_duration or team_member_budget_duration, ) if team_member_budget is not None: @@ -156,6 +159,8 @@ class TeamMemberBudgetHandler: budget_request.rpm_limit = team_member_rpm_limit if team_member_tpm_limit is not None: budget_request.tpm_limit = team_member_tpm_limit + if team_member_budget_duration is not None: + budget_request.budget_duration = team_member_budget_duration team_member_budget_table = await new_budget( budget_obj=budget_request, @@ -182,6 +187,7 @@ class TeamMemberBudgetHandler: team_member_budget: Optional[float] = None, team_member_rpm_limit: Optional[int] = None, team_member_tpm_limit: Optional[int] = None, + team_member_budget_duration: Optional[str] = None, ) -> dict: """Upsert team member budget table with provided limits""" from litellm.proxy._types import BudgetNewRequest @@ -203,6 +209,8 @@ class TeamMemberBudgetHandler: budget_request.rpm_limit = team_member_rpm_limit if team_member_tpm_limit is not None: budget_request.tpm_limit = team_member_tpm_limit + if team_member_budget_duration is not None: + budget_request.budget_duration = team_member_budget_duration budget_row = await update_budget( budget_obj=budget_request, @@ -223,6 +231,7 @@ class TeamMemberBudgetHandler: team_member_budget=team_member_budget, team_member_rpm_limit=team_member_rpm_limit, team_member_tpm_limit=team_member_tpm_limit, + team_member_budget_duration=team_member_budget_duration, ) # Remove team member fields from updated_kv @@ -233,6 +242,7 @@ class TeamMemberBudgetHandler: def _clean_team_member_fields(data_dict: dict) -> None: """Remove team member fields from data dictionary""" data_dict.pop("team_member_budget", None) + data_dict.pop("team_member_budget_duration", None) data_dict.pop("team_member_rpm_limit", None) data_dict.pop("team_member_tpm_limit", None) @@ -1214,6 +1224,7 @@ async def update_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" @@ -1349,6 +1360,7 @@ async def update_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, ): updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=existing_team_row, @@ -1357,6 +1369,7 @@ async def update_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, ) else: TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) 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 6cf8f745e0..57064586af 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1279,7 +1279,7 @@ async def test_update_team_team_member_budget_not_passed_to_db(): # Mock budget upsert to return updated_kv without team_member_budget def mock_upsert_side_effect( - team_table, user_api_key_dict, updated_kv, team_member_budget=None, team_member_rpm_limit=None, team_member_tpm_limit=None + team_table, user_api_key_dict, updated_kv, team_member_budget=None, team_member_rpm_limit=None, team_member_tpm_limit=None, team_member_budget_duration=None ): # Remove team_member_budget from updated_kv as the real function does result_kv = updated_kv.copy() @@ -1376,6 +1376,370 @@ async def test_update_team_team_member_budget_not_passed_to_db(): ) +def test_clean_team_member_fields(): + """ + Test that _clean_team_member_fields removes all team member fields from a dictionary. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + data_dict = { + "team_id": "test_team", + "team_alias": "Test Team", + "team_member_budget": 100.0, + "team_member_budget_duration": "30d", + "team_member_rpm_limit": 50, + "team_member_tpm_limit": 1000, + "other_field": "should_remain", + } + + TeamMemberBudgetHandler._clean_team_member_fields(data_dict) + + assert "team_member_budget" not in data_dict + assert "team_member_budget_duration" not in data_dict + assert "team_member_rpm_limit" not in data_dict + assert "team_member_tpm_limit" not in data_dict + assert data_dict["team_id"] == "test_team" + assert data_dict["team_alias"] == "Test Team" + assert data_dict["other_field"] == "should_remain" + + +def test_clean_team_member_fields_with_missing_fields(): + """ + Test that _clean_team_member_fields handles dictionaries without team member fields gracefully. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + data_dict = { + "team_id": "test_team", + "team_alias": "Test Team", + } + + TeamMemberBudgetHandler._clean_team_member_fields(data_dict) + + assert data_dict["team_id"] == "test_team" + assert data_dict["team_alias"] == "Test Team" + + +@pytest.mark.asyncio +async def test_create_team_member_budget_table(): + """ + Test that create_team_member_budget_table creates a budget and adds it to metadata. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LitellmUserRoles, NewTeamRequest, UserAPIKeyAuth + 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" + ) + + data = NewTeamRequest( + team_id="test_team_id", + team_alias="Test Team", + budget_duration="1mo", + ) + new_team_data_json = { + "team_id": "test_team_id", + "team_alias": "Test Team", + "team_member_budget": 100.0, + "team_member_budget_duration": "30d", + "team_member_rpm_limit": 50, + "team_member_tpm_limit": 1000, + } + + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget_123" + + 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 + + result = await TeamMemberBudgetHandler.create_team_member_budget_table( + data=data, + new_team_data_json=new_team_data_json, + user_api_key_dict=mock_user_api_key_dict, + team_member_budget=100.0, + team_member_rpm_limit=50, + team_member_tpm_limit=1000, + team_member_budget_duration="30d", + ) + + assert mock_new_budget.called + call_args = mock_new_budget.call_args + budget_request = call_args[1]["budget_obj"] + + assert budget_request.max_budget == 100.0 + assert budget_request.rpm_limit == 50 + assert budget_request.tpm_limit == 1000 + assert budget_request.budget_duration == "30d" + assert budget_request.budget_id is not None + assert "team-" in budget_request.budget_id + + assert "team_member_budget_id" in result["metadata"] + assert result["metadata"]["team_member_budget_id"] == "budget_123" + + assert "team_member_budget" not in result + assert "team_member_budget_duration" not in result + assert "team_member_rpm_limit" not in result + assert "team_member_tpm_limit" not in result + + +@pytest.mark.asyncio +async def test_create_team_member_budget_table_without_team_alias(): + """ + Test that create_team_member_budget_table generates budget_id correctly when team_alias is None. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LitellmUserRoles, NewTeamRequest, UserAPIKeyAuth + 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" + ) + + data = NewTeamRequest(team_id="test_team_id") + new_team_data_json = { + "team_id": "test_team_id", + "team_member_budget": 100.0, + } + + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget_123" + + 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 + + result = await TeamMemberBudgetHandler.create_team_member_budget_table( + data=data, + new_team_data_json=new_team_data_json, + user_api_key_dict=mock_user_api_key_dict, + team_member_budget=100.0, + ) + + assert mock_new_budget.called + call_args = mock_new_budget.call_args + budget_request = call_args[1]["budget_obj"] + + assert budget_request.budget_id is not None + assert budget_request.budget_id.startswith("team-budget-") + + +@pytest.mark.asyncio +async def test_upsert_team_member_budget_table_existing_budget(): + """ + Test that upsert_team_member_budget_table updates an existing budget when team_member_budget_id exists. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth + 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"} + + updated_kv = { + "team_id": "test_team_id", + "team_member_budget": 200.0, + "team_member_budget_duration": "60d", + "team_member_rpm_limit": 100, + } + + 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 + + result = await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + team_member_budget=200.0, + team_member_budget_duration="60d", + team_member_rpm_limit=100, + ) + + assert mock_update_budget.called + call_args = mock_update_budget.call_args + budget_request = call_args[1]["budget_obj"] + + assert budget_request.budget_id == "existing_budget_123" + assert budget_request.max_budget == 200.0 + assert budget_request.budget_duration == "60d" + assert budget_request.rpm_limit == 100 + + assert "team_member_budget_id" in result["metadata"] + 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 + assert "team_member_rpm_limit" not in result + + +@pytest.mark.asyncio +async def test_upsert_team_member_budget_table_no_existing_budget(): + """ + Test that upsert_team_member_budget_table creates a new budget when team_member_budget_id does not exist. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth + 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 = None + + updated_kv = { + "team_id": "test_team_id", + "team_member_budget": 150.0, + "team_member_budget_duration": "45d", + } + + 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 + + result = await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + team_member_budget=150.0, + team_member_budget_duration="45d", + ) + + assert mock_new_budget.called + assert "team_member_budget_id" in result["metadata"] + assert result["metadata"]["team_member_budget_id"] == "new_budget_456" + + assert "team_member_budget" not in result + assert "team_member_budget_duration" not in result + + +@pytest.mark.asyncio +async def test_update_team_with_team_member_budget_duration(): + """ + Test that team/update endpoint properly handles team_member_budget_duration. + """ + from unittest.mock import AsyncMock, MagicMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( + "litellm.proxy.proxy_server.llm_router" + ) as mock_llm_router, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.auth.auth_checks._cache_team_object" + ) as mock_cache_team, patch( + "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" + ) as mock_upsert_budget: + + mock_existing_team = MagicMock() + mock_existing_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "metadata": {"team_member_budget_id": "budget_123"}, + } + mock_existing_team.metadata = {"team_member_budget_id": "budget_123"} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "test_team_id" + mock_updated_team.model_dump.return_value = {"team_id": "test_team_id"} + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + mock_prisma_client.jsonify_team_object = MagicMock( + side_effect=lambda db_data: db_data + ) + + def mock_upsert_side_effect( + team_table, user_api_key_dict, updated_kv, team_member_budget=None, team_member_rpm_limit=None, team_member_tpm_limit=None, team_member_budget_duration=None + ): + result_kv = updated_kv.copy() + result_kv.pop("team_member_budget", None) + result_kv.pop("team_member_budget_duration", None) + return result_kv + + mock_upsert_budget.side_effect = mock_upsert_side_effect + + update_request = UpdateTeamRequest( + team_id="test_team_id", + team_alias="updated_alias", + team_member_budget=100.0, + team_member_budget_duration="30d", + ) + + result = await update_team( + data=update_request, + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert mock_upsert_budget.called + call_args = mock_upsert_budget.call_args + assert call_args[1]["team_member_budget"] == 100.0 + assert call_args[1]["team_member_budget_duration"] == "30d" + + assert mock_prisma_client.db.litellm_teamtable.update.called + update_call_args = mock_prisma_client.db.litellm_teamtable.update.call_args + update_data = update_call_args[1]["data"] + + assert "team_member_budget" not in update_data + assert "team_member_budget_duration" not in update_data + + @pytest.mark.asyncio async def test_bulk_team_member_add_success(): """