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/_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): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9170305dff4..7679d873264 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) @@ -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"], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index da2857075ff..2164358ccc4 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? + user_role String? max_budget Float? spend Float @default(0.0) user_email String? 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 diff --git a/schema.prisma b/schema.prisma index da2857075ff..2164358ccc4 100644 --- a/schema.prisma +++ b/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?