diff --git a/app/tasks.py b/app/tasks.py index 6dfb5e8..7e0930b 100644 --- a/app/tasks.py +++ b/app/tasks.py @@ -40,16 +40,18 @@ def generate_task(self, conversation_id, model, prompt): 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 +@celery.task(name="tasks.system_admin_generate", bind=True) +def system_admin_generate(self, cid, model, prompt): + logger.info("system_admin_task_started") + return run_ollama(cid, model, prompt) - 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) +@celery.task(name="tasks.corp_admin_generate", bind=True) +def corp_admin_generate(self, cid, model, prompt): + logger.info("corp_admin_task_started") + return run_ollama(cid, model, prompt) +@celery.task(name="tasks.user_generate", bind=True) +def user_generate(self, cid, model, prompt): + logger.info("user_task_started") + return run_ollama(cid, model, prompt)