Repository: BerriAI/litellm
Pull request: fix: default role for JWT authentication  
## Fix: default role for JWT authentication)

Fixes https://github.com/BerriAI/litellm/issues/10286 

The PR introduces support for applying default internal user parameters when upserting a user via JWT authentication and adds a corresponding unit test to verify this behavior.

Merge default_internal_user_params into new_user_params during user creation in get_user_object
Add a unit test to ensure those defaults are applied on an upsert

<!-- e.g. "Implement user authentication feature" -->

## Relevant issues

<!-- e.g. "Fixes #000" -->

## Pre-Submission checklist

**Please complete all items before asking a LiteLLM maintainer to review your PR**

- [x] I have Added testing in the [`tests/litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code)
- [x] I have added a screenshot of my new test passing locally 
- [x] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code)
- [x] My PR's scope is as isolated as possible, it only solves 1 specific problem


## Type

<!-- Select the type of Pull Request -->
<!-- Keep only the necessary ones -->

🐛 Bug Fix
✅ Test

## Changes




Problem statement:
[Feature]: default role for JWT authentication
### The Feature

Allow to specify the default role to be assigned to a new user using the JWT authentication

### Motivation, pitch

To ensure every new user has constraints in place through the default budget, rpm and budget duration

### Are you a ML Ops Team?

Yes

### Twitter / LinkedIn details

_No response_

Hints from issue discussion:
cc @S1LV3RJ1NX 
any ETA on this @ishaan-jaff @S1LV3RJ1NX ?
hi @abarahonar doing this tomorrow. 
@abarahonar added here: https://github.com/BerriAI/litellm/pull/10995

Commit message: test: test_default_internal_user_params_with_get_user_object

Patch to review:
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index e18b358fa63f..f6cf9e5c87c3 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -697,8 +697,14 @@ async def get_user_object(
 
         if response is None:
             if user_id_upsert:
+                new_user_params: Dict[str, Any] = {
+                    "user_id": user_id,
+                }
+                if litellm.default_internal_user_params:
+                    new_user_params.update(litellm.default_internal_user_params)
+
                 response = await prisma_client.db.litellm_usertable.create(
-                    data={"user_id": user_id},
+                    data=new_user_params,
                     include={"organization_memberships": True},
                 )
             else:
diff --git a/tests/litellm/proxy/auth/test_auth_checks.py b/tests/litellm/proxy/auth/test_auth_checks.py
index 24e8506e9719..97c952b63f5c 100644
--- a/tests/litellm/proxy/auth/test_auth_checks.py
+++ b/tests/litellm/proxy/auth/test_auth_checks.py
@@ -18,7 +18,7 @@
     LitellmUserRoles,
     SSOUserDefinedValues,
 )
-from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
+from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
 from litellm.utils import get_utc_datetime
 
@@ -59,6 +59,8 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v
 
     # Decrypt and verify token contents
     decrypted_token = decrypt_value_helper(token, exception_type="debug")
+    # Check that decrypted_token is not None before using json.loads
+    assert decrypted_token is not None
     token_data = json.loads(decrypted_token)
 
     assert token_data["user_id"] == "test_user"
@@ -110,3 +112,69 @@ def test_get_key_object_from_ui_hash_key_invalid():
     # Test with invalid token
     key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key("invalid_token")
     assert key_object is None
+
+
+@pytest.mark.asyncio
+async def test_default_internal_user_params_with_get_user_object(monkeypatch):
+    """Test that default_internal_user_params is used when creating a new user via get_user_object"""
+    # Set up default_internal_user_params
+    default_params = {
+        "models": ["gpt-4", "claude-3-opus"],
+        "max_budget": 200.0,
+        "user_role": "internal_user",
+    }
+    monkeypatch.setattr(litellm, "default_internal_user_params", default_params)
+
+    # Mock the necessary dependencies
+    mock_prisma_client = MagicMock()
+    mock_db = AsyncMock()
+    mock_prisma_client.db = mock_db
+
+    # Set up the user creation mock - create a complete user model that can be converted to a dict
+    mock_user = MagicMock()
+    mock_user.user_id = "new_test_user"
+    mock_user.models = ["gpt-4", "claude-3-opus"]
+    mock_user.max_budget = 200.0
+    mock_user.user_role = "internal_user"
+    mock_user.organization_memberships = []
+
+    # Make the mock model_dump or dict method return appropriate data
+    mock_user.dict = lambda: {
+        "user_id": "new_test_user",
+        "models": ["gpt-4", "claude-3-opus"],
+        "max_budget": 200.0,
+        "user_role": "internal_user",
+        "organization_memberships": [],
+    }
+
+    # Setup the mock returns
+    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_usertable.create = AsyncMock(return_value=mock_user)
+
+    # Create a mock cache - use AsyncMock for async methods
+    mock_cache = MagicMock()
+    mock_cache.async_get_cache = AsyncMock(return_value=None)
+    mock_cache.async_set_cache = AsyncMock()
+
+    # Call get_user_object with user_id_upsert=True to trigger user creation
+    try:
+        user_obj = await get_user_object(
+            user_id="new_test_user",
+            prisma_client=mock_prisma_client,
+            user_api_key_cache=mock_cache,
+            user_id_upsert=True,
+            proxy_logging_obj=None,
+        )
+    except Exception as e:
+        # this fails since the mock object is a MagicMock and not a LiteLLM_UserTable
+        print(e)
+
+    # Verify the user was created with the default params
+    mock_prisma_client.db.litellm_usertable.create.assert_called_once()
+    creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"]
+
+    # Verify defaults were applied to the creation args
+    assert "models" in creation_args
+    assert creation_args["models"] == ["gpt-4", "claude-3-opus"]
+    assert creation_args["max_budget"] == 200.0
+    assert creation_args["user_role"] == "internal_user"


Code review comments:
diff_hunk: @@ -697,8 +697,14 @@ async def get_user_object(
 
         if response is None:
             if user_id_upsert:
+                new_user_params: Dict[str, Any] = {
+                    "user_id": user_id,
+                }
+                if litellm.default_internal_user_params:
path: litellm/proxy/auth/auth_checks.py
text: @user1:
[nitpick] Use an explicit `is not None` check instead of truthiness to ensure an empty dict still applies the default parameters.
```suggestion
                if litellm.default_internal_user_params is not None:
```