feat(proxy_server.py): add CRUD endpoints for 'end_user' management

allow admin to specify region + default models for end users

.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

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

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(
+@router.post(
-    "/user/get_users",
+    "/end_user/new",
-    tags=["user management"],
+    tags=["End User Management"],
     dependencies=[Depends(user_api_key_auth)],
 )
-async def get_users(
+async def new_end_user(
-    role: str = fastapi.Query(
+    data: NewEndUserRequest,
-        default=None,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
-        description="Either 'proxy_admin', 'proxy_viewer', 'app_owner', 'app_user'",
-    )
 ):
     """
-    [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}
         )
 
-    return all_users
+    ## 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
+
+
+@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 ####

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)

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
 

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)