Merge pull request #14514 from iabhi4/fix-13581

feat(proxy): Assign default budget to auto-generated JWT teams

Author: krrishdholakia

docs/my-website/docs/proxy/team_budgets.md

@@ -10,8 +10,30 @@ import TabItem from '@theme/TabItem';
 - You must set up a Postgres database (e.g. Supabase, Neon, etc.)
 - To enable team member rate limits, set the environment variable `EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING=true` **before starting the proxy server**. Without this, team member rate limits will not be enforced.
 
+
+## Default Budget for Auto-Generated JWT Teams
+
+When using JWT authentication with `team_id_upsert: true`, you can automatically assign a default budget to any newly created team.
+
+This is configured in `default_team_settings` in your `config.yaml`.
+
+**Example:**
+```yaml
+# in your config.yaml
+
+litellm_jwtauth:
+  team_id_upsert: true
+  team_id_jwt_field: "team_id"
+  # ... other jwt settings
+
+litellm_settings:
+  default_team_settings: 
+    - team_id: "default-settings"
+      max_budget: 100.0
+```
 Track spend, set budgets for your Internal Team
 
+
 ## Setting Monthly Team Budgets
 
 ### 1. Create a team 

litellm/proxy/auth/auth_checks.py

@@ -46,6 +46,7 @@
     RoleBasedPermissions,
     SpecialModelNames,
     UserAPIKeyAuth,
+    NewTeamRequest,
 )
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.route_llm_request import route_request
@@ -889,10 +890,17 @@ async def _get_team_db_check(
     )
 
     if response is None and team_id_upsert:
-        response = await prisma_client.db.litellm_teamtable.create(
-            data={"team_id": team_id}
-        )
+        from litellm.proxy.management_endpoints.team_endpoints import new_team
+
+        new_team_data = NewTeamRequest(team_id=team_id)
 
+        mock_request = Request(scope={"type": "http"})
+        system_admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+
+        created_team_dict = await new_team(
+            data=new_team_data, http_request=mock_request, user_api_key_dict=system_admin_user
+        )
+        response = LiteLLM_TeamTable(**created_team_dict)
     return response
 
 

litellm/proxy/management_endpoints/team_endpoints.py

@@ -383,6 +383,19 @@ async def new_team(  # noqa: PLR0915
                         "error": f"Team id = {data.team_id} already exists. Please use a different team id."
                     },
                 )
+            
+        # If max_budget is not explicitly provided in the request,
+        # check for a default value in the proxy configuration.
+        if data.max_budget is None:
+            if (
+                isinstance(litellm.default_team_settings, list)
+                and len(litellm.default_team_settings) > 0
+                and isinstance(litellm.default_team_settings[0], dict)
+            ):
+                default_settings = litellm.default_team_settings[0]
+                default_budget = default_settings.get("max_budget")
+                if default_budget is not None:
+                    data.max_budget = default_budget
 
         if (
             user_api_key_dict.user_role is None

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -28,6 +28,7 @@
     _can_object_call_vector_stores,
     get_user_object,
     vector_store_access_check,
+    _get_team_db_check,
 )
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
 from litellm.utils import get_utc_datetime
@@ -192,6 +193,64 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch):
     assert creation_args["user_role"] == "internal_user"
 
 
+@pytest.mark.asyncio
+@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock)
+async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeypatch):
+    """
+    Test that _get_team_db_check correctly calls the `new_team` function
+    when a team does not exist and upsert is enabled.
+    """
+    mock_prisma_client = MagicMock()
+    mock_db = AsyncMock()
+    mock_prisma_client.db = mock_db
+    mock_prisma_client.db.litellm_teamtable.find_unique.return_value = None
+
+    # Define what our mocked `new_team` function should return
+    team_id_to_create = "new-jwt-team"
+    mock_new_team.return_value = {"team_id": team_id_to_create, "max_budget": 123.45}
+
+    await _get_team_db_check(
+        team_id=team_id_to_create,
+        prisma_client=mock_prisma_client,
+        team_id_upsert=True,
+    )
+
+    # Verify that our mocked `new_team` function was called exactly once
+    mock_new_team.assert_called_once()
+
+    call_args = mock_new_team.call_args[1]
+    data_arg = call_args["data"]
+
+    # Verify that `new_team` was called with the correct team_id and that
+    # `max_budget` was None, as our function's job is to delegate, not to set defaults.
+    assert data_arg.team_id == team_id_to_create
+    assert data_arg.max_budget is None
+
+
+@pytest.mark.asyncio
+@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock)
+async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team, monkeypatch):
+    """
+    Test that _get_team_db_check does NOT call the `new_team` function
+    if the team already exists in the database.
+    """
+    mock_prisma_client = MagicMock()
+    mock_db = AsyncMock()
+    mock_prisma_client.db = mock_db
+    mock_prisma_client.db.litellm_teamtable.find_unique.return_value = MagicMock()
+
+    team_id_to_find = "existing-jwt-team"
+
+    await _get_team_db_check(
+        team_id=team_id_to_find,
+        prisma_client=mock_prisma_client,
+        team_id_upsert=True,
+    )
+
+    # Verify that `new_team` was NEVER called, because the team was found.
+    mock_new_team.assert_not_called()
+
+
 # Vector Store Auth Check Tests