Merge pull request #1795 from BerriAI/litellm_allow_setting_user_roles

[Feat] Allow setting user roles for UserTable

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

@@ -1,4 +1,4 @@
-# Virtual Keys
+# Virtual Keys, Users
 Track Spend, Set budgets and create virtual keys for the proxy
 
 Grant other's temporary access to your proxy, with keys that expire after a set duration.
@@ -278,6 +278,80 @@ Request Params:
 }
 ```
 
+## /user/new
+
+### Request
+
+All [key/generate params supported](#keygenerate) for creating a user
+```shell
+curl 'http://0.0.0.0:4000/user/new' \
+--header 'Authorization: Bearer sk-1234' \
+--header 'Content-Type: application/json' \
+--data-raw '{
+  "user_id": "ishaan1",
+  "user_email": "ishaan@litellm.ai",
+  "user_role": "admin",
+  "team_id": "cto-team",
+  "max_budget": 20,
+  "budget_duration": "1h"
+
+}'
+```
+
+Request Params:
+
+- user_id: str (optional - defaults to uuid)  - The unique identifier for the user.
+- user_email: str (optional - defaults to "")  - The email address associated with the user.
+- user_role: str (optional - defaults to "app_user") - The role assigned to the user. Can be "admin", "app_owner", "app_user"
+
+**Possible `user_role` values**
+```
+"admin" - Maintaining the proxy and owning the overall budget
+"app_owner" - employees maintaining the apps, each owner may own more than one app
+"app_user" - users who know nothing about the proxy. These users get created when you pass `user` to /chat/completions
+```
+- team_id: str (optional - defaults to "") - The identifier for the team to which the user belongs.
+- max_budget: float (optional - defaults to `null`) - The maximum budget allocated for the user. No budget checks done if `max_budget==null`
+- budget_duration: str (optional - defaults to `null`) - The duration for which the budget is valid, e.g., "1h", "1d"
+
+### Response
+A key will be generated for the new user created
+
+```shell
+{
+  "models": [],
+  "spend": 0.0,
+  "max_budget": null,
+  "user_id": "ishaan1",
+  "team_id": null,
+  "max_parallel_requests": null,
+  "metadata": {},
+  "tpm_limit": null,
+  "rpm_limit": null,
+  "budget_duration": null,
+  "allowed_cache_controls": [],
+  "key_alias": null,
+  "duration": null,
+  "aliases": {},
+  "config": {},
+  "key": "sk-JflB33ucTqc2NYvNAgiBCA",
+  "key_name": null,
+  "expires": null
+}
+
+```
+
+Request Params:
+- keys: List[str] - List of keys to delete
+
+### Response
+
+```json
+{
+  "deleted_keys": ["sk-kdEXbIqZRwEeEiHwdg7sFA"]
+}
+```
+
 ## Default /key/generate params
 Use this, if you need to control the default `max_budget` or any `key/generate` param per key. 
 

litellm/proxy/_types.py

@@ -194,6 +194,7 @@ class DeleteKeyRequest(LiteLLMBase):
 class NewUserRequest(GenerateKeyRequest):
     max_budget: Optional[float] = None
     user_email: Optional[str] = None
+    user_role: Optional[str] = None
 
 
 class NewUserResponse(GenerateKeyResponse):
@@ -206,6 +207,7 @@ class UpdateUserRequest(GenerateRequestBase):
     user_id: str
     spend: Optional[float] = None
     metadata: Optional[dict] = None
+    user_role: Optional[str] = None
 
 
 class KeyManagementSystem(enum.Enum):

litellm/proxy/proxy_server.py

@@ -1336,6 +1336,7 @@ async def generate_key_helper_fn(
     user_id: Optional[str] = None,
     team_id: Optional[str] = None,
     user_email: Optional[str] = None,
+    user_role: Optional[str] = None,
     max_parallel_requests: Optional[int] = None,
     metadata: Optional[dict] = {},
     tpm_limit: Optional[int] = None,
@@ -1396,6 +1397,7 @@ async def generate_key_helper_fn(
     config_json = json.dumps(config)
     metadata_json = json.dumps(metadata)
     user_id = user_id or str(uuid.uuid4())
+    user_role = user_role or "app_user"
     tpm_limit = tpm_limit
     rpm_limit = rpm_limit
     allowed_cache_controls = allowed_cache_controls
@@ -1408,6 +1410,7 @@ async def generate_key_helper_fn(
             "user_email": user_email,
             "user_id": user_id,
             "team_id": team_id,
+            "user_role": user_role,
             "spend": spend,
             "models": models,
             "max_parallel_requests": max_parallel_requests,
@@ -2985,6 +2988,7 @@ async def new_user(data: NewUserRequest):
     Parameters:
     - user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated.
     - user_email: Optional[str] - Specify a user email.
+    - user_role: Optional[str] - Specify a user role - "proxy_admin", "app_owner", "app_user"
     - max_budget: Optional[float] - Specify max budget for a given user.
     - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). **(Default is set to 1 hour.)**
     - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
@@ -3001,6 +3005,16 @@ async def new_user(data: NewUserRequest):
     - max_budget: (float|None) Max budget for given user.
     """
     data_json = data.json()  # type: ignore
+    if "user_role" in data_json:
+        user_role = data_json["user_role"]
+        if user_role is not None:
+            if user_role not in ["admin", "app_owner", "app_user"]:
+                raise ProxyException(
+                    message=f"Invalid user role, passed in {user_role}. Must be one of 'admin', 'app_owner', 'app_user'",
+                    type="invalid_user_role",
+                    param="user_role",
+                    code=status.HTTP_400_BAD_REQUEST,
+                )
     response = await generate_key_helper_fn(**data_json)
     return NewUserResponse(
         key=response["token"],

litellm/proxy/schema.prisma

@@ -11,6 +11,7 @@ generator client {
 model LiteLLM_UserTable {
 		user_id    String @unique
     team_id    String?
+    user_role  String?
 		max_budget Float?
     spend      Float    @default(0.0)
     user_email    String?

litellm/tests/test_key_generate_prisma.py

@@ -47,6 +47,7 @@ from litellm.proxy.proxy_server import (
     spend_user_fn,
     spend_key_fn,
     view_spend_logs,
+    user_info,
 )
 from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
 from litellm._logging import verbose_proxy_logger
@@ -106,9 +107,17 @@ def test_generate_and_call_with_valid_key(prisma_client):
             await litellm.proxy.proxy_server.prisma_client.connect()
             from litellm.proxy.proxy_server import user_api_key_cache
 
-            request = NewUserRequest()
+            request = NewUserRequest(user_role="app_owner")
             key = await new_user(request)
             print(key)
+            user_id = key.user_id
+
+            # check /user/info to verify user_role was set correctly
+            new_user_info = await user_info(user_id=user_id)
+            new_user_info = new_user_info["user_info"]
+            print("new_user_info=", new_user_info)
+            assert new_user_info.user_role == "app_owner"
+            assert new_user_info.user_id == user_id
 
             generated_key = key.key
             bearer_token = "Bearer " + generated_key

schema.prisma

@@ -11,6 +11,7 @@ generator client {
 model LiteLLM_UserTable {
 		user_id    String @unique
     team_id    String?
+    user_role  String?
 		max_budget Float?
     spend      Float    @default(0.0)
     user_email    String?