diff --git a/app/tasks.py b/app/tasks.py index 32528ad..6dfb5e8 100644 --- a/app/tasks.py +++ b/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) +