diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py new file mode 100644 index 000000000..cb57619b9 --- /dev/null +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -0,0 +1,581 @@ +#### END-USER/CUSTOMER MANAGEMENT #### +import asyncio +import copy +import json +import re +import secrets +import time +import traceback +import uuid +from datetime import datetime, timedelta, timezone +from typing import List, Optional + +import fastapi +from fastapi import APIRouter, Depends, Header, HTTPException, Request, status + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import handle_exception_on_proxy + +router = APIRouter() + + +@router.post( + "/end_user/block", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + include_in_schema=False, +) +@router.post( + "/customer/block", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def block_user(data: BlockUsers): + """ + [BETA] Reject calls with this end-user id + + Parameters: + - user_ids (List[str], required): The unique `user_id`s for the users to block + + (any /chat/completion call with this user={end-user-id} param, will be rejected.) + + ``` + curl -X POST "http://0.0.0.0:8000/user/block" + -H "Authorization: Bearer sk-1234" + -D '{ + "user_ids": [, ...] + }' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + records = [] + if prisma_client is not None: + for id in data.user_ids: + record = await prisma_client.db.litellm_endusertable.upsert( + where={"user_id": id}, # type: ignore + data={ + "create": {"user_id": id, "blocked": True}, # type: ignore + "update": {"blocked": True}, + }, + ) + records.append(record) + else: + raise HTTPException( + status_code=500, + detail={"error": "Postgres DB Not connected"}, + ) + + return {"blocked_users": records} + except Exception as e: + verbose_proxy_logger.error(f"An error occurred - {str(e)}") + raise HTTPException(status_code=500, detail={"error": str(e)}) + + +@router.post( + "/end_user/unblock", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + include_in_schema=False, +) +@router.post( + "/customer/unblock", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def unblock_user(data: BlockUsers): + """ + [BETA] Unblock calls with this user id + + Example + ``` + curl -X POST "http://0.0.0.0:8000/user/unblock" + -H "Authorization: Bearer sk-1234" + -D '{ + "user_ids": [, ...] + }' + ``` + """ + from enterprise.enterprise_hooks.blocked_user_list import ( + _ENTERPRISE_BlockedUserList, + ) + + if ( + not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks) + or litellm.blocked_user_list is None + ): + raise HTTPException( + status_code=400, + detail={ + "error": "Blocked user check was never set. This call has no effect." + }, + ) + + if isinstance(litellm.blocked_user_list, list): + for id in data.user_ids: + litellm.blocked_user_list.remove(id) + else: + raise HTTPException( + status_code=500, + detail={ + "error": "`blocked_user_list` must be set as a list. Filepaths can't be updated." + }, + ) + + return {"blocked_users": litellm.blocked_user_list} + + +@router.post( + "/end_user/new", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/customer/new", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_end_user( + data: NewCustomerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Allow creating a new Customer + + + Parameters: + - user_id: str - The unique identifier for the user. + - alias: Optional[str] - A human-friendly alias for the user. + - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. + - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. + - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. + - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} + + + - Allow specifying allowed regions + - Allow specifying default model + + Example curl: + ``` + curl --location 'http://0.0.0.0:4000/customer/new' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "user_id" : "ishaan-jaff-3", + "allowed_region": "eu", + "budget_id": "free_tier", + "default_model": "azure/gpt-3.5-turbo-eu" <- all calls from this user, use this model? + }' + + # return end-user object + ``` + + NOTE: This used to be called `/end_user/new`, we will still be maintaining compatibility for /end_user/XXX for these endpoints + """ + """ + 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 + """ + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + llm_router, + prisma_client, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + try: + + ## 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 end_user_record + except Exception as e: + if "Unique constraint failed on the fields: (`user_id`)" in str(e): + raise ProxyException( + message=f"Customer already exists, passed user_id={data.user_id}. Please pass a new user_id.", + type="bad_request", + code=400, + param="user_id", + ) + + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Internal Server Error({str(e)})"), + type="internal_error", + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Internal Server Error, " + str(e), + type="internal_error", + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@router.get( + "/customer/info", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_EndUserTable, +) +@router.get( + "/end_user/info", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def end_user_info( + end_user_id: str = fastapi.Query( + description="End User ID in the request parameters" + ), +): + """ + Get information about an end-user. An `end_user` is a customer (external user) of the proxy. + + Parameters: + - end_user_id (str, required): The unique identifier for the end-user + + Example curl: + ``` + curl -X GET 'http://localhost:4000/customer/info?end_user_id=test-litellm-user-4' \ + -H 'Authorization: Bearer sk-1234' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + user_info = await prisma_client.db.litellm_endusertable.find_first( + where={"user_id": end_user_id}, include={"litellm_budget_table": True} + ) + + if user_info is None: + raise HTTPException( + status_code=400, + detail={"error": "End User Id={} does not exist in db".format(end_user_id)}, + ) + return user_info.model_dump(exclude_none=True) + + +@router.post( + "/customer/update", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/end_user/update", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def update_end_user( + data: UpdateCustomerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Example curl + + Parameters: + - 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[AllowedModelRegion] = ( + 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 + ) + + Example curl: + ``` + curl --location 'http://0.0.0.0:4000/customer/update' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "user_id": "test-litellm-user-4", + "budget_id": "paid_tier" + }' + + See below for all params + ``` + """ + + from litellm.proxy.proxy_server import prisma_client + + try: + data_json: dict = data.json() + # get the row from db + if prisma_client is None: + raise Exception("Not connected to DB!") + + # get non default values for key + non_default_values = {} + for k, v in data_json.items(): + if v is not None and v not in ( + [], + {}, + 0, + ): # models default to [], spend defaults to 0, we should not reset these values + non_default_values[k] = v + + ## ADD USER, IF NEW ## + verbose_proxy_logger.debug("/customer/update: Received data = %s", data) + if data.user_id is not None and len(data.user_id) > 0: + non_default_values["user_id"] = data.user_id # type: ignore + verbose_proxy_logger.debug("In update customer, user_id condition block.") + response = await prisma_client.db.litellm_endusertable.update( + where={"user_id": data.user_id}, data=non_default_values # type: ignore + ) + if response is None: + raise ValueError( + f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}" + ) + verbose_proxy_logger.debug( + f"received response from updating prisma client. response={response}" + ) + return response + else: + raise ValueError(f"user_id is required, passed user_id = {data.user_id}") + + # update based on remaining passed in values + except Exception as e: + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.update_end_user(): Exception occured - {}".format( + str(e) + ) + ) + verbose_proxy_logger.debug(traceback.format_exc()) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Internal Server Error({str(e)})"), + type="internal_error", + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Internal Server Error, " + str(e), + type="internal_error", + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + pass + + +@router.post( + "/customer/delete", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/end_user/delete", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def delete_end_user( + data: DeleteCustomerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete multiple end-users. + + Parameters: + - user_ids (List[str], required): The unique `user_id`s for the users to delete + + Example curl: + ``` + curl --location 'http://0.0.0.0:4000/customer/delete' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "user_ids" :["ishaan-jaff-5"] + }' + + See below for all params + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client is None: + raise Exception("Not connected to DB!") + + verbose_proxy_logger.debug("/customer/delete: Received data = %s", data) + if ( + data.user_ids is not None + and isinstance(data.user_ids, list) + and len(data.user_ids) > 0 + ): + response = await prisma_client.db.litellm_endusertable.delete_many( + where={"user_id": {"in": data.user_ids}} + ) + if response is None: + raise ValueError( + f"Failed deleting customer data. User ID does not exist passed user_id={data.user_ids}" + ) + if response != len(data.user_ids): + raise ValueError( + f"Failed deleting all customer data. User ID does not exist passed user_id={data.user_ids}. Deleted {response} customers, passed {len(data.user_ids)} customers" + ) + verbose_proxy_logger.debug( + f"received response from updating prisma client. response={response}" + ) + return { + "deleted_customers": response, + "message": "Successfully deleted customers with ids: " + + str(data.user_ids), + } + else: + raise ValueError(f"user_id is required, passed user_id = {data.user_ids}") + + # update based on remaining passed in values + except Exception as e: + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.delete_end_user(): Exception occured - {}".format( + str(e) + ) + ) + verbose_proxy_logger.debug(traceback.format_exc()) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Internal Server Error({str(e)})"), + type="internal_error", + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Internal Server Error, " + str(e), + type="internal_error", + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + pass + + +@router.get( + "/customer/list", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=List[LiteLLM_EndUserTable], +) +@router.get( + "/end_user/list", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def list_end_user( + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + [Admin-only] List all available customers + + Example curl: + ``` + curl --location --request GET 'http://0.0.0.0:4000/customer/list' \ + --header 'Authorization: Bearer sk-1234' + ``` + + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + if ( + user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ): + raise HTTPException( + status_code=401, + detail={ + "error": "Admin-only endpoint. Your user role={}".format( + user_api_key_dict.user_role + ) + }, + ) + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + response = await prisma_client.db.litellm_endusertable.find_many( + include={"litellm_budget_table": True} + ) + + returned_response: List[LiteLLM_EndUserTable] = [] + for item in response: + returned_response.append(LiteLLM_EndUserTable(**item.model_dump())) + return returned_response diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 23c1803ca..49ef25149 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) +from litellm.proxy.utils import handle_exception_on_proxy router = APIRouter() @@ -197,76 +198,6 @@ async def new_user( ) -@router.post( - "/user/auth", - tags=["Internal User management"], - dependencies=[Depends(user_api_key_auth)], -) -async def user_auth(request: Request): - """ - Allows UI ("https://dashboard.litellm.ai/", or self-hosted - os.getenv("LITELLM_HOSTED_UI")) to request a magic link to be sent to user email, for auth to proxy. - - Only allows emails from accepted email subdomains. - - Rate limit: 1 request every 60s. - - Only works, if you enable 'allow_user_auth' in general settings: - e.g.: - ```yaml - general_settings: - allow_user_auth: true - ``` - - Requirements: - SMTP server details saved in .env: - - os.environ["SMTP_HOST"] - - os.environ["SMTP_PORT"] - - os.environ["SMTP_USERNAME"] - - os.environ["SMTP_PASSWORD"] - - os.environ["SMTP_SENDER_EMAIL"] - """ - from litellm.proxy.proxy_server import prisma_client - from litellm.proxy.utils import send_email - - data = await request.json() # type: ignore - user_email = data["user_email"] - page_params = data["page"] - if user_email is None: - raise HTTPException(status_code=400, detail="User email is none") - - if prisma_client is None: # if no db connected, raise an error - raise Exception("No connected db.") - - ### Check if user email in user table - response = await prisma_client.get_generic_data( - key="user_email", value=user_email, table_name="users" - ) - ### if so - generate a 24 hr key with that user id - if response is not None: - user_id = response.user_id # type: ignore - response = await generate_key_helper_fn( - request_type="key", - **{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id}, # type: ignore - ) - else: ### else - create new user - response = await generate_key_helper_fn( - request_type="key", - **{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_email": user_email}, # type: ignore - ) - - base_url = os.getenv("LITELLM_HOSTED_UI", "https://dashboard.litellm.ai/") - - params = { - "sender_name": "LiteLLM Proxy", - "receiver_email": user_email, - "subject": "Your Magic Link", - "html": f" Follow this link, to login:\n\n{base_url}user/?token={response['token']}&user_id={response['user_id']}&page={page_params}", - } - - await send_email(**params) - return "Email sent!" - - @router.get( "/user/available_roles", tags=["Internal User management"], @@ -338,7 +269,7 @@ async def user_info( # noqa: PLR0915 Example request ``` - curl -X GET 'http://localhost:8000/user/info?user_id=krrish7%40berri.ai' \ + curl -X GET 'http://localhost:4000/user/info?user_id=krrish7%40berri.ai' \ --header 'Authorization: Bearer sk-1234' ``` """ @@ -488,21 +419,7 @@ async def user_info( # noqa: PLR0915 str(e) ) ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + raise handle_exception_on_proxy(e) @router.post( @@ -527,7 +444,55 @@ async def user_update( "user_role": "proxy_admin_viewer" }' - See below for all params + Parameters: + user_id: Optional[str] + Unique identifier for the user to update + + user_email: Optional[str] + Email address for the user + + password: Optional[str] + Password for the user + + user_role: Optional[Literal["proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer"]] + Role assigned to the user. Can be one of: + - proxy_admin: Full admin access + - proxy_admin_viewer: Read-only admin access + - internal_user: Standard internal user + - internal_user_viewer: Read-only internal user + + models: Optional[list] + List of model names the user is allowed to access + + spend: Optional[float] + Current spend amount for the user + + max_budget: Optional[float] + Maximum budget allowed for the user + + team_id: Optional[str] + ID of the team the user belongs to + + max_parallel_requests: Optional[int] + Maximum number of concurrent requests allowed + + metadata: Optional[dict] + Additional metadata associated with the user + + tpm_limit: Optional[int] + Maximum tokens per minute allowed + + rpm_limit: Optional[int] + Maximum requests per minute allowed + + budget_duration: Optional[str] + Duration for budget renewal (e.g., "30d" for 30 days) + + allowed_cache_controls: Optional[list] + List of allowed cache control options + + soft_budget: Optional[float] + Soft budget limit for alerting purposes ``` """ from litellm.proxy.proxy_server import prisma_client @@ -643,113 +608,6 @@ async def user_update( ) -@router.post( - "/user/request_model", - tags=["Internal User management"], - dependencies=[Depends(user_api_key_auth)], -) -async def user_request_model(request: Request): - """ - Allow a user to create a request to access a model - """ - from litellm.proxy.proxy_server import prisma_client - - try: - data_json = await request.json() - - # get the row from db - if prisma_client is None: - raise Exception("Not connected to DB!") - - non_default_values = {k: v for k, v in data_json.items() if v is not None} - new_models = non_default_values.get("models", None) - user_id = non_default_values.get("user_id", None) - justification = non_default_values.get("justification", None) - - await prisma_client.insert_data( - data={ - "models": new_models, - "justification": justification, - "user_id": user_id, - "status": "pending", - "request_id": str(uuid.uuid4()), - }, - table_name="user_notification", - ) - return {"status": "success"} - # update based on remaining passed in values - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.user_request_model(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - -@router.get( - "/user/get_requests", - tags=["Internal User management"], - dependencies=[Depends(user_api_key_auth)], -) -async def user_get_requests(): - """ - Get all "Access" requests made by proxy users, access requests are requests for accessing models - """ - from litellm.proxy.proxy_server import prisma_client - - try: - - # get the row from db - if prisma_client is None: - raise Exception("Not connected to DB!") - - # TODO: Optimize this so we don't read all the data here, eventually move to pagination - response = await prisma_client.get_data( - query_type="find_all", - table_name="user_notification", - ) - return {"requests": response} - # update based on remaining passed in values - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.user_get_requests(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - @router.get( "/user/get_users", tags=["Internal User management"], @@ -774,6 +632,18 @@ async def get_users( Used by the UI to populate the user lists. + Parameters: + role: Optional[str] + Filter users by role. Can be one of: + - proxy_admin + - proxy_admin_viewer + - internal_user + - internal_user_viewer + page: int + The page number to return + page_size: int + The number of items per page + Currently - admin-only endpoint. """ from litellm.proxy.proxy_server import prisma_client @@ -842,7 +712,7 @@ async def delete_user( delete user and associated user keys ``` - curl --location 'http://0.0.0.0:8000/user/delete' \ + curl --location 'http://0.0.0.0:4000/user/delete' \ --header 'Authorization: Bearer sk-1234' \ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2fdc44752..c2de82ce7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -33,7 +33,11 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_helpers.utils import management_endpoint_wrapper -from litellm.proxy.utils import _duration_in_seconds, _hash_token_if_needed +from litellm.proxy.utils import ( + _duration_in_seconds, + _hash_token_if_needed, + handle_exception_on_proxy, +) from litellm.secret_managers.main import get_secret router = APIRouter() @@ -84,7 +88,7 @@ async def generate_key_fn( # noqa: PLR0915 1. Allow users to turn on/off pii masking ```bash - curl --location 'http://0.0.0.0:8000/key/generate' \ + curl --location 'http://0.0.0.0:4000/key/generate' \ --header 'Authorization: Bearer sk-1234' \ --header 'Content-Type: application/json' \ --data '{ @@ -251,21 +255,7 @@ async def generate_key_fn( # noqa: PLR0915 str(e) ) ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + raise handle_exception_on_proxy(e) def prepare_key_update_data( @@ -362,7 +352,7 @@ async def update_key_fn( Example: ```bash - curl --location 'http://0.0.0.0:8000/key/update' \ + curl --location 'http://0.0.0.0:4000/key/update' \ --header 'Authorization: Bearer sk-1234' \ --header 'Content-Type: application/json' \ --data '{ @@ -477,7 +467,7 @@ async def delete_key_fn( Example: ```bash - curl --location 'http://0.0.0.0:8000/key/delete' \ + curl --location 'http://0.0.0.0:4000/key/delete' \ --header 'Authorization: Bearer sk-1234' \ --header 'Content-Type: application/json' \ --data '{ @@ -568,21 +558,7 @@ async def delete_key_fn( return {"deleted_keys": keys} except Exception as e: - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + raise handle_exception_on_proxy(e) @router.post( @@ -607,7 +583,7 @@ async def info_key_fn_v2( Example Curl: ``` - curl -X GET "http://0.0.0.0:8000/key/info" \ + curl -X GET "http://0.0.0.0:4000/key/info" \ -H "Authorization: Bearer sk-1234" \ -d {"keys": ["sk-1", "sk-2", "sk-3"]} ``` @@ -651,21 +627,7 @@ async def info_key_fn_v2( return {"key": data.keys, "info": filtered_key_info} except Exception as e: - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + raise handle_exception_on_proxy(e) @router.get( @@ -687,13 +649,13 @@ async def info_key_fn( Example Curl: ``` - curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \ + curl -X GET "http://0.0.0.0:4000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \ -H "Authorization: Bearer sk-1234" ``` Example Curl - if no key is passed, it will use the Key Passed in Authorization Header ``` - curl -X GET "http://0.0.0.0:8000/key/info" \ + curl -X GET "http://0.0.0.0:4000/key/info" \ -H "Authorization: Bearer sk-02Wr4IAlN3NvPXvL5JVvDA" ``` """ @@ -752,21 +714,7 @@ async def info_key_fn( key_info.pop("token") return {"key": key, "info": key_info} except Exception as e: - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + raise handle_exception_on_proxy(e) async def generate_key_helper_fn( # noqa: PLR0915 @@ -1082,105 +1030,155 @@ async def regenerate_key_fn( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), -) -> GenerateKeyResponse: - from litellm.proxy.proxy_server import ( - hash_token, - premium_user, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - +) -> Optional[GenerateKeyResponse]: """ - Endpoint for regenerating a key + Regenerate an existing API key while optionally updating its parameters. + + Parameters: + - key: str (path parameter) - The key to regenerate + - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update + - key_alias: Optional[str] - User-friendly key alias + - user_id: Optional[str] - User ID associated with key + - team_id: Optional[str] - Team ID associated with key + - models: Optional[list] - Model_name's a user is allowed to call + - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) + - spend: Optional[float] - Amount spent by key + - max_budget: Optional[float] - Max budget for key + - model_max_budget: Optional[dict] - Model-specific budgets {"gpt-4": 0.5, "claude-v1": 1.0} + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + - max_parallel_requests: Optional[int] - Rate limit for parallel requests + - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} + - tpm_limit: Optional[int] - Tokens per minute limit + - rpm_limit: Optional[int] - Requests per minute limit + - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} + - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} + - allowed_cache_controls: Optional[list] - List of allowed cache control values + - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) + - permissions: Optional[dict] - Key-specific permissions + - guardrails: Optional[List[str]] - List of active guardrails for the key + - blocked: Optional[bool] - Whether the key is blocked + + + Returns: + - GenerateKeyResponse containing the new key and its updated parameters + + Example: + ```bash + curl --location --request POST 'http://localhost:4000/key/sk-1234/regenerate' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data-raw '{ + "max_budget": 100, + "metadata": {"team": "core-infra"}, + "models": ["gpt-4", "gpt-3.5-turbo"], + "model_max_budget": {"gpt-4": 50, "gpt-3.5-turbo": 50} + }' + ``` + + Note: This is an Enterprise feature. It requires a premium license to use. """ + try: - if premium_user is not True: - raise ValueError( - f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}" + from litellm.proxy.proxy_server import ( + hash_token, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, ) - # Check if key exists, raise exception if key is not in the DB + if premium_user is not True: + raise ValueError( + f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}" + ) - ### 1. Create New copy that is duplicate of existing key - ###################################################################### + # Check if key exists, raise exception if key is not in the DB - # create duplicate of existing key - # set token = new token generated - # insert new token in DB + ### 1. Create New copy that is duplicate of existing key + ###################################################################### - # create hash of token - if prisma_client is None: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": "DB not connected. prisma_client is None"}, + # create duplicate of existing key + # set token = new token generated + # insert new token in DB + + # create hash of token + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "DB not connected. prisma_client is None"}, + ) + + if "sk" not in key: + hashed_api_key = key + else: + hashed_api_key = hash_token(key) + + _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_api_key}, + ) + if _key_in_db is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Key {key} not found."}, + ) + + verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) + + new_token = f"sk-{secrets.token_urlsafe(16)}" + new_token_hash = hash_token(new_token) + new_token_key_name = f"sk-...{new_token[-4:]}" + + # Prepare the update data + update_data = { + "token": new_token_hash, + "key_name": new_token_key_name, + } + + non_default_values = {} + if data is not None: + # Update with any provided parameters from GenerateKeyRequest + non_default_values = prepare_key_update_data( + data=data, existing_key_row=_key_in_db + ) + verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + + update_data.update(non_default_values) + update_data = prisma_client.jsonify_object(data=update_data) + # Update the token in the database + updated_token = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_api_key}, + data=update_data, # type: ignore ) - if "sk" not in key: - hashed_api_key = key - else: - hashed_api_key = hash_token(key) + updated_token_dict = {} + if updated_token is not None: + updated_token_dict = dict(updated_token) - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_api_key}, - ) - if _key_in_db is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"Key {key} not found."}, + updated_token_dict["key"] = new_token + updated_token_dict.pop("token") + + ### 3. remove existing key entry from cache + ###################################################################### + if key: + await _delete_cache_key_object( + hashed_token=hash_token(key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if hashed_api_key: + await _delete_cache_key_object( + hashed_token=hash_token(key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return GenerateKeyResponse( + **updated_token_dict, ) - - verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) - - new_token = f"sk-{secrets.token_urlsafe(16)}" - new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" - - # Prepare the update data - update_data = { - "token": new_token_hash, - "key_name": new_token_key_name, - } - - non_default_values = {} - if data is not None: - # Update with any provided parameters from GenerateKeyRequest - non_default_values = prepare_key_update_data( - data=data, existing_key_row=_key_in_db - ) - - update_data.update(non_default_values) - # Update the token in the database - updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, - data=update_data, # type: ignore - ) - - updated_token_dict = {} - if updated_token is not None: - updated_token_dict = dict(updated_token) - - updated_token_dict["token"] = new_token - - ### 3. remove existing key entry from cache - ###################################################################### - if key: - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - if hashed_api_key: - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - return GenerateKeyResponse( - **updated_token_dict, - ) + except Exception as e: + raise handle_exception_on_proxy(e) @router.get( @@ -1303,9 +1301,24 @@ async def block_key( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), -): +) -> Optional[LiteLLM_VerificationToken]: """ - Blocks all calls from keys with this team id. + Block an Virtual key from making any requests. + + Parameters: + - key: str - The key to block. Can be either the unhashed key (sk-...) or the hashed key value + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/block' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + }' + ``` + + Note: This is an admin-only endpoint. Only proxy admins can block keys. """ from litellm.proxy.proxy_server import ( create_audit_log_for_update, @@ -1397,7 +1410,22 @@ async def unblock_key( ), ): """ - Unblocks all calls from this key. + Unblock a Virtual key to allow it to make requests again. + + Parameters: + - key: str - The key to unblock. Can be either the unhashed key (sk-...) or the hashed key value + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/unblock' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + }' + ``` + + Note: This is an admin-only endpoint. Only proxy admins can unblock keys. """ from litellm.proxy.proxy_server import ( create_audit_log_for_update, diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index d51ca9ea1..b60ea3acf 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -55,6 +55,23 @@ async def add_team_callbacks( Use this if if you want different teams to have different success/failure callbacks + Parameters: + - callback_name (Literal["langfuse", "langsmith", "gcs"], required): The name of the callback to add + - callback_type (Literal["success", "failure", "success_and_failure"], required): The type of callback to add. One of: + - "success": Callback for successful LLM calls + - "failure": Callback for failed LLM calls + - "success_and_failure": Callback for both successful and failed LLM calls + - callback_vars (StandardCallbackDynamicParams, required): A dictionary of variables to pass to the callback + - langfuse_public_key: The public key for the Langfuse callback + - langfuse_secret_key: The secret key for the Langfuse callback + - langfuse_secret: The secret for the Langfuse callback + - langfuse_host: The host for the Langfuse callback + - gcs_bucket_name: The name of the GCS bucket + - gcs_path_service_account: The path to the GCS service account + - langsmith_api_key: The API key for the Langsmith callback + - langsmith_project: The project for the Langsmith callback + - langsmith_base_url: The base URL for the Langsmith callback + Example curl: ``` curl -X POST 'http:/localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \ @@ -201,6 +218,20 @@ async def disable_team_logging( team_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): + """ + Disable all logging callbacks for a team + + Parameters: + - team_id (str, required): The unique identifier for the team + + Example curl: + ``` + curl -X POST 'http://localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/disable_logging' \ + -H 'Authorization: Bearer sk-1234' + ``` + + + """ try: from litellm.proxy.proxy_server import prisma_client @@ -289,6 +320,9 @@ async def get_team_callbacks( """ Get the success/failure callbacks and variables for a team + Parameters: + - team_id (str, required): The unique identifier for the team + Example curl: ``` curl -X GET 'http://localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8dcd0c7eb..ec6949936 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -932,6 +932,9 @@ async def delete_team( """ delete team and associated team keys + Parameters: + - team_ids: List[str] - Required. List of team IDs to delete. Example: ["team-1234", "team-5678"] + ``` curl --location 'http://0.0.0.0:4000/team/delete' \ --header 'Authorization: Bearer sk-1234' \ @@ -1022,6 +1025,9 @@ async def team_info( """ get info on team + related keys + Parameters: + - team_id: str - Required. The unique identifier of the team to get info on. + ``` curl --location 'http://localhost:4000/team/info?team_id=your_team_id_here' \ --header 'Authorization: Bearer your_api_key_here' @@ -1156,6 +1162,25 @@ async def block_team( ): """ Blocks all calls from keys with this team id. + + Parameters: + - team_id: str - Required. The unique identifier of the team to block. + + Example: + ``` + curl --location 'http://0.0.0.0:4000/team/block' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "team_id": "team-1234" + }' + ``` + + Returns: + - The updated team record with blocked=True + + + """ from litellm.proxy.proxy_server import ( _duration_in_seconds, @@ -1171,6 +1196,12 @@ async def block_team( where={"team_id": data.team_id}, data={"blocked": True} # type: ignore ) + if record is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) + return record @@ -1185,6 +1216,19 @@ async def unblock_team( ): """ Blocks all calls from keys with this team id. + + Parameters: + - team_id: str - Required. The unique identifier of the team to unblock. + + Example: + ``` + curl --location 'http://0.0.0.0:4000/team/unblock' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "team_id": "team-1234" + }' + ``` """ from litellm.proxy.proxy_server import ( _duration_in_seconds, @@ -1200,6 +1244,12 @@ async def unblock_team( where={"team_id": data.team_id}, data={"blocked": False} # type: ignore ) + if record is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) + return record @@ -1219,6 +1269,9 @@ async def list_team( curl --location --request GET 'http://0.0.0.0:4000/team/list' \ --header 'Authorization: Bearer sk-1234' ``` + + Parameters: + - user_id: str - Optional. If passed will only return teams that the user_id is a member of. """ from litellm.proxy.proxy_server import ( _duration_in_seconds, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8577181ce..548d07689 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -665,7 +665,6 @@ async def initialize_pass_through_endpoints(pass_through_endpoints: list): @router.get( "/config/pass_through_endpoint", - tags=["Internal User management"], dependencies=[Depends(user_api_key_auth)], response_model=PassThroughEndpointResponse, ) @@ -715,7 +714,6 @@ async def get_pass_through_endpoints( @router.post( "/config/pass_through_endpoint/{endpoint_id}", - tags=["Internal User management"], dependencies=[Depends(user_api_key_auth)], ) async def update_pass_through_endpoints(request: Request, endpoint_id: str): @@ -727,7 +725,6 @@ async def update_pass_through_endpoints(request: Request, endpoint_id: str): @router.post( "/config/pass_through_endpoint", - tags=["Internal User management"], dependencies=[Depends(user_api_key_auth)], ) async def create_pass_through_endpoints( @@ -773,7 +770,6 @@ async def create_pass_through_endpoints( @router.delete( "/config/pass_through_endpoint", - tags=["Internal User management"], dependencies=[Depends(user_api_key_auth)], response_model=PassThroughEndpointResponse, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 92ca32e52..2ece9705a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -174,6 +174,9 @@ from litellm.proxy.hooks.prompt_injection_detection import ( _OPTIONAL_PromptInjectionDetection, ) from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.proxy.management_endpoints.customer_endpoints import ( + router as customer_router, +) from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) @@ -5954,525 +5957,6 @@ async def supported_openai_params(model: str): ) -#### END-USER MANAGEMENT #### - - -@router.post( - "/end_user/block", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, -) -@router.post( - "/customer/block", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], -) -async def block_user(data: BlockUsers): - """ - [BETA] Reject calls with this end-user id - - (any /chat/completion call with this user={end-user-id} param, will be rejected.) - - ``` - curl -X POST "http://0.0.0.0:8000/user/block" - -H "Authorization: Bearer sk-1234" - -D '{ - "user_ids": [, ...] - }' - ``` - """ - try: - records = [] - if prisma_client is not None: - for id in data.user_ids: - record = await prisma_client.db.litellm_endusertable.upsert( - where={"user_id": id}, # type: ignore - data={ - "create": {"user_id": id, "blocked": True}, # type: ignore - "update": {"blocked": True}, - }, - ) - records.append(record) - else: - raise HTTPException( - status_code=500, - detail={"error": "Postgres DB Not connected"}, - ) - - return {"blocked_users": records} - except Exception as e: - verbose_proxy_logger.error(f"An error occurred - {str(e)}") - raise HTTPException(status_code=500, detail={"error": str(e)}) - - -@router.post( - "/end_user/unblock", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, -) -@router.post( - "/customer/unblock", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], -) -async def unblock_user(data: BlockUsers): - """ - [BETA] Unblock calls with this user id - - Example - ``` - curl -X POST "http://0.0.0.0:8000/user/unblock" - -H "Authorization: Bearer sk-1234" - -D '{ - "user_ids": [, ...] - }' - ``` - """ - from enterprise.enterprise_hooks.blocked_user_list import ( - _ENTERPRISE_BlockedUserList, - ) - - if ( - not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks) - or litellm.blocked_user_list is None - ): - raise HTTPException( - status_code=400, - detail={ - "error": "Blocked user check was never set. This call has no effect." - }, - ) - - if isinstance(litellm.blocked_user_list, list): - for id in data.user_ids: - litellm.blocked_user_list.remove(id) - else: - raise HTTPException( - status_code=500, - detail={ - "error": "`blocked_user_list` must be set as a list. Filepaths can't be updated." - }, - ) - - return {"blocked_users": litellm.blocked_user_list} - - -@router.post( - "/end_user/new", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -@router.post( - "/customer/new", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], -) -async def new_end_user( - data: NewCustomerRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Allow creating a new Customer - - - Parameters: - - user_id: str - The unique identifier for the user. - - alias: Optional[str] - A human-friendly alias for the user. - - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. - - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. - - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. - - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. - - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. - - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} - - - - Allow specifying allowed regions - - Allow specifying default model - - Example curl: - ``` - curl --location 'http://0.0.0.0:4000/customer/new' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "user_id" : "ishaan-jaff-3", - "allowed_region": "eu", - "budget_id": "free_tier", - "default_model": "azure/gpt-3.5-turbo-eu" <- all calls from this user, use this model? - }' - - # return end-user object - ``` - - NOTE: This used to be called `/end_user/new`, we will still be maintaining compatibility for /end_user/XXX for these endpoints - """ - 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": CommonProxyErrors.db_not_connected_error.value}, - ) - try: - - ## 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 end_user_record - except Exception as e: - if "Unique constraint failed on the fields: (`user_id`)" in str(e): - raise ProxyException( - message=f"Customer already exists, passed user_id={data.user_id}. Please pass a new user_id.", - type="bad_request", - code=400, - param="user_id", - ) - - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Internal Server Error({str(e)})"), - type="internal_error", - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Internal Server Error, " + str(e), - type="internal_error", - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - -@router.get( - "/customer/info", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], - response_model=LiteLLM_EndUserTable, -) -@router.get( - "/end_user/info", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def end_user_info( - end_user_id: str = fastapi.Query( - description="End User ID in the request parameters" - ), -): - global prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - user_info = await prisma_client.db.litellm_endusertable.find_first( - where={"user_id": end_user_id}, include={"litellm_budget_table": True} - ) - - if user_info is None: - raise HTTPException( - status_code=400, - detail={"error": "End User Id={} does not exist in db".format(end_user_id)}, - ) - return user_info.model_dump(exclude_none=True) - - -@router.post( - "/customer/update", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], -) -@router.post( - "/end_user/update", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def update_end_user( - data: UpdateCustomerRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Example curl - - ``` - curl --location 'http://0.0.0.0:4000/customer/update' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "user_id": "test-litellm-user-4", - "budget_id": "paid_tier" - }' - - See below for all params - ``` - """ - - global prisma_client - try: - data_json: dict = data.json() - # get the row from db - if prisma_client is None: - raise Exception("Not connected to DB!") - - # get non default values for key - non_default_values = {} - for k, v in data_json.items(): - if v is not None and v not in ( - [], - {}, - 0, - ): # models default to [], spend defaults to 0, we should not reset these values - non_default_values[k] = v - - ## ADD USER, IF NEW ## - verbose_proxy_logger.debug("/customer/update: Received data = %s", data) - if data.user_id is not None and len(data.user_id) > 0: - non_default_values["user_id"] = data.user_id # type: ignore - verbose_proxy_logger.debug("In update customer, user_id condition block.") - response = await prisma_client.db.litellm_endusertable.update( - where={"user_id": data.user_id}, data=non_default_values # type: ignore - ) - if response is None: - raise ValueError( - f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}" - ) - verbose_proxy_logger.debug( - f"received response from updating prisma client. response={response}" - ) - return response - else: - raise ValueError(f"user_id is required, passed user_id = {data.user_id}") - - # update based on remaining passed in values - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.update_end_user(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Internal Server Error({str(e)})"), - type="internal_error", - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Internal Server Error, " + str(e), - type="internal_error", - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - pass - - -@router.post( - "/customer/delete", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], -) -@router.post( - "/end_user/delete", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def delete_end_user( - data: DeleteCustomerRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Example curl - - ``` - curl --location 'http://0.0.0.0:4000/customer/delete' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "user_ids" :["ishaan-jaff-5"] - }' - - See below for all params - ``` - """ - global prisma_client - - try: - if prisma_client is None: - raise Exception("Not connected to DB!") - - verbose_proxy_logger.debug("/customer/delete: Received data = %s", data) - if ( - data.user_ids is not None - and isinstance(data.user_ids, list) - and len(data.user_ids) > 0 - ): - response = await prisma_client.db.litellm_endusertable.delete_many( - where={"user_id": {"in": data.user_ids}} - ) - if response is None: - raise ValueError( - f"Failed deleting customer data. User ID does not exist passed user_id={data.user_ids}" - ) - if response != len(data.user_ids): - raise ValueError( - f"Failed deleting all customer data. User ID does not exist passed user_id={data.user_ids}. Deleted {response} customers, passed {len(data.user_ids)} customers" - ) - verbose_proxy_logger.debug( - f"received response from updating prisma client. response={response}" - ) - return { - "deleted_customers": response, - "message": "Successfully deleted customers with ids: " - + str(data.user_ids), - } - else: - raise ValueError(f"user_id is required, passed user_id = {data.user_ids}") - - # update based on remaining passed in values - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.delete_end_user(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Internal Server Error({str(e)})"), - type="internal_error", - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Internal Server Error, " + str(e), - type="internal_error", - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - pass - - -@router.get( - "/customer/list", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_EndUserTable], -) -@router.get( - "/end_user/list", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def list_end_user( - http_request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - [Admin-only] List all available customers - - ``` - curl --location --request GET 'http://0.0.0.0:4000/customer/list' \ - --header 'Authorization: Bearer sk-1234' - ``` - """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client - - if ( - user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ): - raise HTTPException( - status_code=401, - detail={ - "error": "Admin-only endpoint. Your user role={}".format( - user_api_key_dict.user_role - ) - }, - ) - - if prisma_client is None: - raise HTTPException( - status_code=400, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - response = await prisma_client.db.litellm_endusertable.find_many( - include={"litellm_budget_table": True} - ) - - returned_response: List[LiteLLM_EndUserTable] = [] - for item in response: - returned_response.append(LiteLLM_EndUserTable(**item.model_dump())) - return returned_response - - #### BUDGET TABLE MANAGEMENT #### @@ -9651,6 +9135,7 @@ app.include_router(internal_user_router) app.include_router(team_router) app.include_router(ui_sso_router) app.include_router(organization_router) +app.include_router(customer_router) app.include_router(spend_management_router) app.include_router(caching_router) app.include_router(analytics_router) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c143d30e4..bcd3ce16c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -26,6 +26,8 @@ from typing import ( overload, ) +from litellm.proxy._types import ProxyErrorTypes, ProxyException + try: import backoff except ImportError: @@ -3095,3 +3097,26 @@ def get_error_message_str(e: Exception) -> str: else: error_message = str(e) return error_message + + +def handle_exception_on_proxy(e: Exception) -> ProxyException: + """ + Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible + """ + from fastapi import status + + if isinstance(e, HTTPException): + return ProxyException( + message=getattr(e, "detail", f"error({str(e)})"), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + return e + return ProxyException( + message="Internal Server Error, " + str(e), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) diff --git a/tests/local_testing/test_blocked_user_list.py b/tests/local_testing/test_blocked_user_list.py index 10635befd..172d6e85e 100644 --- a/tests/local_testing/test_blocked_user_list.py +++ b/tests/local_testing/test_blocked_user_list.py @@ -44,7 +44,8 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( info_key_fn, update_key_fn, ) -from litellm.proxy.proxy_server import block_user, user_api_key_auth +from litellm.proxy.proxy_server import user_api_key_auth +from litellm.proxy.management_endpoints.customer_endpoints import block_user from litellm.proxy.spend_tracking.spend_management_endpoints import ( spend_key_fn, spend_user_fn, diff --git a/tests/local_testing/test_update_spend.py b/tests/local_testing/test_update_spend.py index 3eb9f1ab4..6aeae851a 100644 --- a/tests/local_testing/test_update_spend.py +++ b/tests/local_testing/test_update_spend.py @@ -41,7 +41,8 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( info_key_fn, update_key_fn, ) -from litellm.proxy.proxy_server import block_user, user_api_key_auth +from litellm.proxy.proxy_server import user_api_key_auth +from litellm.proxy.management_endpoints.customer_endpoints import block_user from litellm.proxy.spend_tracking.spend_management_endpoints import ( spend_key_fn, spend_user_fn, diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index bc7371843..b039a101b 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -55,9 +55,11 @@ from litellm.proxy.proxy_server import ( image_generation, model_list, moderations, - new_end_user, user_api_key_auth, ) +from litellm.proxy.management_endpoints.customer_endpoints import ( + new_end_user, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( global_spend, global_spend_logs, diff --git a/tests/proxy_admin_ui_tests/test_role_based_access.py b/tests/proxy_admin_ui_tests/test_role_based_access.py index 6f59fd6f5..609a3598d 100644 --- a/tests/proxy_admin_ui_tests/test_role_based_access.py +++ b/tests/proxy_admin_ui_tests/test_role_based_access.py @@ -58,9 +58,11 @@ from litellm.proxy.proxy_server import ( image_generation, model_list, moderations, - new_end_user, user_api_key_auth, ) +from litellm.proxy.management_endpoints.customer_endpoints import ( + new_end_user, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( global_spend, global_spend_logs, diff --git a/tests/proxy_admin_ui_tests/test_usage_endpoints.py b/tests/proxy_admin_ui_tests/test_usage_endpoints.py index 4a9ba9588..cd704e49c 100644 --- a/tests/proxy_admin_ui_tests/test_usage_endpoints.py +++ b/tests/proxy_admin_ui_tests/test_usage_endpoints.py @@ -67,9 +67,11 @@ from litellm.proxy.proxy_server import ( image_generation, model_list, moderations, - new_end_user, user_api_key_auth, ) +from litellm.proxy.management_endpoints.customer_endpoints import ( + new_end_user, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( global_spend, global_spend_logs, diff --git a/tests/proxy_unit_tests/test_audit_logs_proxy.py b/tests/proxy_unit_tests/test_audit_logs_proxy.py index 275d48670..02303e13d 100644 --- a/tests/proxy_unit_tests/test_audit_logs_proxy.py +++ b/tests/proxy_unit_tests/test_audit_logs_proxy.py @@ -37,7 +37,6 @@ from litellm.proxy.proxy_server import ( image_generation, model_list, moderations, - new_end_user, user_api_key_auth, ) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index fb6e2c7f5..4de451642 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -76,9 +76,11 @@ from litellm.proxy.proxy_server import ( image_generation, model_list, moderations, - new_end_user, user_api_key_auth, ) +from litellm.proxy.management_endpoints.customer_endpoints import ( + new_end_user, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( global_spend, spend_key_fn,