diff --git a/docs/my-website/docs/proxy/team_budgets.md b/docs/my-website/docs/proxy/team_budgets.md index 66ba679c65..3847406641 100644 --- a/docs/my-website/docs/proxy/team_budgets.md +++ b/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 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f1242a1f34..51092f0897 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( 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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 7b0df5a562..9766386344 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index eb26eb776f..9a50986a1b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.auth_checks import ( _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