diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2fac0db143..2c970cd4be 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 #### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 101cf9b7f0..5377fe90b9 100644 --- a/litellm/proxy/schema.prisma +++ b/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 } \ No newline at end of file diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3fdb80b918..bd0c0eaf77 100644 --- a/litellm/proxy/utils.py +++ b/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}") diff --git a/schema.prisma b/schema.prisma index 101cf9b7f0..8663db1b06 100644 --- a/schema.prisma +++ b/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 } \ No newline at end of file