add team member budget duration in team/update

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

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)

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():
     """