Atualizar app/tasks.py

app/tasks.py

@@ -38,3 +38,18 @@ def generate_task(self, conversation_id, model, prompt):
 
     except Exception as e:
         raise self.retry(exc=e, countdown=2)
+
+
+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)
+