Merge pull request #2080 from BerriAI/litellm_backend_req_model_access

[FEAT] proxy backend - Save User Model Requests in DB

litellm/proxy/proxy_server.py

@@ -4122,6 +4122,101 @@ async def user_update(data: UpdateUserRequest):
         )
 
 
+@router.post(
+    "/user/request_model",
+    tags=["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
+    """
+    global 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)
+
+        response = 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:
+        traceback.print_exc()
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Authentication Error({str(e)})"),
+                type="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="auth_error",
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
+@router.get(
+    "/user/get_requests",
+    tags=["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
+    """
+    global 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:
+        traceback.print_exc()
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Authentication Error({str(e)})"),
+                type="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="auth_error",
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
 #### TEAM MANAGEMENT ####
 
 

litellm/proxy/schema.prisma

@@ -95,4 +95,13 @@ model LiteLLM_SpendLogs {
   cache_hit           String   @default("")
   cache_key           String   @default("")
   request_tags        Json     @default("[]")
+}
+
+// Beta - allow team members to request access to a model
+model LiteLLM_UserNotifications {
+  request_id          String @unique
+  user_id             String 
+  models              String[]
+  justification       String
+    status              String // approved, disapproved, pending
 }

litellm/proxy/utils.py

@@ -512,7 +512,9 @@ class PrismaClient:
         user_id_list: Optional[list] = None,
         team_id: Optional[str] = None,
         key_val: Optional[dict] = None,
-        table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
+        table_name: Optional[
+            Literal["user", "key", "config", "spend", "team", "user_notification"]
+        ] = None,
         query_type: Literal["find_unique", "find_all"] = "find_unique",
         expires: Optional[datetime] = None,
         reset_at: Optional[datetime] = None,
@@ -677,6 +679,14 @@ class PrismaClient:
                         where={"members": {"has": user_id}}
                     )
                 return response
+            elif table_name == "user_notification":
+                if query_type == "find_unique":
+                    response = await self.db.litellm_usernotifications.find_unique(
+                        where={"user_id": user_id}  # type: ignore
+                    )
+                elif query_type == "find_all":
+                    response = await self.db.litellm_usernotifications.find_many()
+                return response
         except Exception as e:
             print_verbose(f"LiteLLM Prisma Client Exception: {e}")
             import traceback
@@ -696,7 +706,11 @@ class PrismaClient:
         on_backoff=on_backoff,  # specifying the function to call on backoff
     )
     async def insert_data(
-        self, data: dict, table_name: Literal["user", "key", "config", "spend", "team"]
+        self,
+        data: dict,
+        table_name: Literal[
+            "user", "key", "config", "spend", "team", "user_notification"
+        ],
     ):
         """
         Add a key to the database. If it already exists, do nothing.
@@ -778,6 +792,19 @@ class PrismaClient:
                 )
                 verbose_proxy_logger.info(f"Data Inserted into Spend Table")
                 return new_spend_row
+            elif table_name == "user_notification":
+                db_data = self.jsonify_object(data=data)
+                new_user_notification_row = (
+                    await self.db.litellm_usernotifications.upsert(
+                        where={"request_id": data["request_id"]},
+                        data={
+                            "create": {**db_data},  # type: ignore
+                            "update": {},  # don't do anything if it already exists
+                        },
+                    )
+                )
+                verbose_proxy_logger.info(f"Data Inserted into Model Request Table")
+                return new_user_notification_row
 
         except Exception as e:
             print_verbose(f"LiteLLM Prisma Client Exception: {e}")

schema.prisma

@@ -95,4 +95,13 @@ model LiteLLM_SpendLogs {
   cache_hit           String   @default("")
   cache_key           String   @default("")
   request_tags        Json     @default("[]")
+}
+
+// Beta - allow team members to request access to a model
+model LiteLLM_UserNotifications {
+  request_id          String @unique
+  user_id             String 
+  models              String[]
+  justification       String
+  status              String // approved, disapproved, pending
 }