From db666b01e5ebd632568bc904d62033158a63def0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 8 May 2024 18:50:36 -0700 Subject: [PATCH] feat(proxy_server.py): add CRUD endpoints for 'end_user' management allow admin to specify region + default models for end users --- .gitignore | 3 + litellm/proxy/_types.py | 31 ++++++ litellm/proxy/proxy_server.py | 180 ++++++++++++++++++++++++++++++---- litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 2 +- schema.prisma | 2 + 6 files changed, 201 insertions(+), 19 deletions(-) diff --git a/.gitignore b/.gitignore index abc4ecb0c..1f827e463 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .venv .env +litellm/proxy/myenv/* litellm_uuid.txt __pycache__/ *.pyc @@ -52,3 +53,5 @@ litellm/proxy/_new_secret_config.yaml litellm/proxy/_new_secret_config.yaml litellm/proxy/_super_secret_config.yaml litellm/proxy/_super_secret_config.yaml +litellm/proxy/myenv/bin/activate +litellm/proxy/myenv/bin/Activate.ps1 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b7b2d0ab6..37e00e27e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -458,6 +458,37 @@ class UpdateUserRequest(GenerateRequestBase): return values +class NewEndUserRequest(LiteLLMBase): + user_id: str + alias: Optional[str] = None # human-friendly alias + blocked: bool = False # allow/disallow requests for this end-user + max_budget: Optional[float] = None + budget_id: Optional[str] = None # give either a budget_id or max_budget + allowed_model_region: Optional[Literal["eu"]] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) + + @root_validator(pre=True) + def check_user_info(cls, values): + if values.get("max_budget") is not None and values.get("budget_id") is not None: + raise ValueError("Set either 'max_budget' or 'budget_id', not both.") + + if ( + values.get("allowed_model_region") is not None + and values.get("default_model") is None + ) or ( + values.get("allowed_model_region") is None + and values.get("default_model") is not None + ): + raise ValueError( + "If 'allowed_model_region' is set, then 'default_model' must be set." + ) + return values + + class Member(LiteLLMBase): role: Literal["admin", "user"] user_id: Optional[str] = None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ce62a7609..009ea279f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -231,6 +231,11 @@ class SpecialModelNames(enum.Enum): all_team_models = "all-team-models" +class CommonProxyErrors(enum.Enum): + db_not_connected_error = "DB not connected" + no_llm_router = "No models configured on proxy" + + @app.exception_handler(ProxyException) async def openai_exception_handler(request: Request, exc: ProxyException): # NOTE: DO NOT MODIFY THIS, its crucial to map to Openai exceptions @@ -5883,7 +5888,7 @@ async def global_predict_spend_logs(request: Request): return _forecast_daily_cost(data) -#### USER MANAGEMENT #### +#### INTERNAL USER MANAGEMENT #### @router.post( "/user/new", tags=["user management"], @@ -6376,6 +6381,43 @@ async def user_get_requests(): ) +@router.get( + "/user/get_users", + tags=["user management"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_users( + role: str = fastapi.Query( + default=None, + description="Either 'proxy_admin', 'proxy_viewer', 'app_owner', 'app_user'", + ) +): + """ + [BETA] This could change without notice. Give feedback - https://github.com/BerriAI/litellm/issues + + Get all users who are a specific `user_role`. + + Used by the UI to populate the user lists. + + Currently - admin-only endpoint. + """ + global prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": f"No db connected. prisma client={prisma_client}"}, + ) + all_users = await prisma_client.get_data( + table_name="user", query_type="find_all", key_val={"user_role": role} + ) + + return all_users + + +#### END-USER MANAGEMENT #### + + @router.post( "/end_user/block", tags=["End User Management"], @@ -6466,38 +6508,140 @@ async def unblock_user(data: BlockUsers): return {"blocked_users": litellm.blocked_user_list} -@router.get( - "/user/get_users", - tags=["user management"], +@router.post( + "/end_user/new", + tags=["End User Management"], dependencies=[Depends(user_api_key_auth)], ) -async def get_users( - role: str = fastapi.Query( - default=None, - description="Either 'proxy_admin', 'proxy_viewer', 'app_owner', 'app_user'", - ) +async def new_end_user( + data: NewEndUserRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - [BETA] This could change without notice. Give feedback - https://github.com/BerriAI/litellm/issues + [TODO] Needs to be implemented. - Get all users who are a specific `user_role`. + Allow creating a new end-user - Used by the UI to populate the user lists. + - Allow specifying allowed regions + - Allow specifying default model - Currently - admin-only endpoint. + Example curl: + ``` + curl --location 'http://0.0.0.0:4000/end_user/new' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "end_user_id" : "ishaan-jaff-3", <- specific customer + + "allowed_region": "eu" <- set region for models + + + + + "default_model": "azure/gpt-3.5-turbo-eu" <- all calls from this user, use this model? + + }' + + # return end-user object + ``` """ - global prisma_client + global prisma_client, llm_router + """ + Validation: + - check if default model exists + - create budget object if not already created + + - Add user to end user table + Return + - end-user object + - currently allowed models + """ if prisma_client is None: raise HTTPException( status_code=500, - detail={"error": f"No db connected. prisma client={prisma_client}"}, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - all_users = await prisma_client.get_data( - table_name="user", query_type="find_all", key_val={"user_role": role} + + ## VALIDATION ## + if data.default_model is not None: + if llm_router is None: + raise HTTPException( + status_code=422, detail={"error": CommonProxyErrors.no_llm_router.value} + ) + elif data.default_model not in llm_router.get_model_names(): + raise HTTPException( + status_code=422, + detail={ + "error": "Default Model not on proxy. Configure via `/model/new` or config.yaml. Default_model={}, proxy_model_names={}".format( + data.default_model, set(llm_router.get_model_names()) + ) + }, + ) + + new_end_user_obj: Dict = {} + + ## CREATE BUDGET ## if set + if data.max_budget is not None: + budget_record = await prisma_client.db.litellm_budgettable.create( + data={ + "max_budget": data.max_budget, + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + ) + + new_end_user_obj["budget_id"] = budget_record.budget_id + elif data.budget_id is not None: + new_end_user_obj["budget_id"] = data.budget_id + + _user_data = data.dict(exclude_none=True) + + for k, v in _user_data.items(): + if k != "max_budget" and k != "budget_id": + new_end_user_obj[k] = v + + ## WRITE TO DB ## + end_user_record = await prisma_client.db.litellm_endusertable.create( + data=new_end_user_obj # type: ignore ) - return all_users + return end_user_record + + +@router.post( + "/end_user/info", + tags=["End User Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def end_user_info(): + """ + [TODO] Needs to be implemented. + """ + pass + + +@router.post( + "/end_user/update", + tags=["End User Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_end_user(): + """ + [TODO] Needs to be implemented. + """ + pass + + +@router.post( + "/end_user/delete", + tags=["End User Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_end_user(): + """ + [TODO] Needs to be implemented. + """ + pass #### TEAM MANAGEMENT #### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f078dbcf4..e24a9d5f3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -150,6 +150,8 @@ model LiteLLM_EndUserTable { user_id String @id alias String? // admin-facing alias spend Float @default(0.0) + allowed_model_region String? // require all user requests to use models in this specific region + default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model. budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) blocked Boolean @default(false) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0379d5152..4f577a84c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -526,7 +526,7 @@ class PrismaClient: finally: os.chdir(original_dir) # Now you can import the Prisma Client - from prisma import Prisma # type: ignore + from prisma import Prisma self.db = Prisma() # Client to connect to Prisma db diff --git a/schema.prisma b/schema.prisma index f078dbcf4..e24a9d5f3 100644 --- a/schema.prisma +++ b/schema.prisma @@ -150,6 +150,8 @@ model LiteLLM_EndUserTable { user_id String @id alias String? // admin-facing alias spend Float @default(0.0) + allowed_model_region String? // require all user requests to use models in this specific region + default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model. budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) blocked Boolean @default(false)