From aee477f4aa5a4a862204c1f35877ab14f79dbd85 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:09:27 -0800 Subject: [PATCH 1/7] (feat) add user_rol to schema.prisma --- litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index da2857075ff..ab6ff638bea 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -11,6 +11,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique team_id String? + role String? max_budget Float? spend Float @default(0.0) user_email String? diff --git a/schema.prisma b/schema.prisma index da2857075ff..ab6ff638bea 100644 --- a/schema.prisma +++ b/schema.prisma @@ -11,6 +11,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique team_id String? + role String? max_budget Float? spend Float @default(0.0) user_email String? From d41f1b4997ffe8dc0cb124797d4e2c9dc2789fd7 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:11:42 -0800 Subject: [PATCH 2/7] (feat) set user_role for Users on proxy --- litellm/proxy/proxy_server.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9170305dff4..97a740072e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/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) From 93c7e1d457113fa0f5c22bd2843b44f137e3b7f3 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:15:03 -0800 Subject: [PATCH 3/7] (feat) add user_role to _types --- litellm/proxy/_types.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f7d1e8ad0c4..e8e01602097 100644 --- a/litellm/proxy/_types.py +++ b/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): From 29885bbdc1901634bdcb300e022b7a16196819d1 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:23:50 -0800 Subject: [PATCH 4/7] (feat) set user_role for Users --- litellm/proxy/schema.prisma | 2 +- schema.prisma | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index ab6ff638bea..2164358ccc4 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -11,7 +11,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique team_id String? - role String? + user_role String? max_budget Float? spend Float @default(0.0) user_email String? diff --git a/schema.prisma b/schema.prisma index ab6ff638bea..2164358ccc4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -11,7 +11,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique team_id String? - role String? + user_role String? max_budget Float? spend Float @default(0.0) user_email String? From c9958eb329ff0f80cd8d5b9abdb91a22b216fcbf Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:40:59 -0800 Subject: [PATCH 5/7] (feat) raise error when user_role not in known roles --- litellm/proxy/proxy_server.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97a740072e2..919499bdfb6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3005,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 ["proxy_admin", "app_owner", "app_user"]: + raise ProxyException( + message=f"Invalid user role, passed in {user_role}. Must be one of 'proxy_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"], From 3f2f864cc8087edf0e3fad6313bc8a9b51c2b5fc Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 10:59:06 -0800 Subject: [PATCH 6/7] (docs) /user/new --- docs/my-website/docs/proxy/virtual_keys.md | 76 +++++++++++++++++++++- litellm/proxy/proxy_server.py | 4 +- 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 11ca903cda1..dd5edc6da81 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 919499bdfb6..7679d873264 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3008,9 +3008,9 @@ async def new_user(data: NewUserRequest): if "user_role" in data_json: user_role = data_json["user_role"] if user_role is not None: - if user_role not in ["proxy_admin", "app_owner", "app_user"]: + 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 'proxy_admin', 'app_owner', 'app_user'", + 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, From 8455aa4bd4e0bacc05541fa2b6ecb77af14fc359 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 11:06:51 -0800 Subject: [PATCH 7/7] (test) setting user_role --- litellm/tests/test_key_generate_prisma.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 140d638e880..734a0b11487 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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