Atualizar app/app.py

app/app.py

@@ -55,6 +55,8 @@ def queue_admin_generate():
 # ================
 #   DECORATORS
 # ================
+
+
 def enforce_plan_limits(f):
     @wraps(f)
     def wrapper(*args, **kwargs):
@@ -1113,6 +1115,27 @@ def queue_generate():
 
     return jsonify({"task_id": task.id})
 
+@app.route("/api/generate", methods=["POST"])
+@login_required
+def generate():
+    data = request.json
+    cid = data["conversation_id"]
+    model = data["model"]
+    prompt = data["prompt"]
+
+    db = SessionLocal()
+    email = session["user"]["email"] if app.config["AUTH_MODE"] == "oauth" else session["db_user"]["email"]
+    user = db.query(User).filter_by(email=email).first()
+    db.close()
+
+    task = enqueue_by_priority(user, cid, model, prompt)
+
+    return jsonify({"task_id": task.id})
+
+
+
+
+
 from celery_app import celery
 
 @app.route("/api/queue/status/<task_id>")
@@ -1127,6 +1150,34 @@ def queue_status(task_id):
 
 
 
+def enqueue_by_priority(user, cid, model, prompt):
+    if user.priority == "system_admin":
+        return system_admin_generate.delay(cid, model, prompt)
+
+    if user.priority == "corp_admin":
+        return corp_admin_generate.delay(cid, model, prompt)
+
+    return user_generate.delay(cid, model, prompt)
+
+
+def enqueue_task(model, prompt, cid, user):
+    db = SessionLocal()
+    plan = db.query(Plan).filter_by(id=user.plan_id).first()
+    db.close()
+
+    queue = plan.priority_queue  # default, premium, admin
+
+    if queue == "admin":
+        return admin_generate.delay(cid, model, prompt)
+    if queue == "premium":
+        return premium_generate.delay(cid, model, prompt)
+    return generate_task.delay(cid, model, prompt)
+
+
+
+
+
+
 # ***
 
 import json