Atualizar app/tasks.py
This commit is contained in:
+12
-10
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user