56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
# All comments are in English.
|
|
|
|
from celery_app import celery
|
|
import requests
|
|
import json
|
|
import uuid
|
|
from models import Message, Conversation
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy import create_engine
|
|
from config import Config
|
|
|
|
engine = create_engine(Config.DATABASE_URL)
|
|
SessionLocal = sessionmaker(bind=engine)
|
|
|
|
@celery.task(bind=True, max_retries=3)
|
|
def generate_task(self, conversation_id, model, prompt):
|
|
"""Runs an Ollama generation asynchronously."""
|
|
try:
|
|
url = f"{Config.OLLAMA_BASE_URL}/api/generate"
|
|
payload = {"model": model, "prompt": prompt}
|
|
|
|
r = requests.post(url, json=payload)
|
|
r.raise_for_status()
|
|
output = r.json().get("response", "")
|
|
|
|
db = SessionLocal()
|
|
msg = Message(
|
|
id=str(uuid.uuid4()),
|
|
conversation_id=conversation_id,
|
|
role="assistant",
|
|
content=output,
|
|
)
|
|
db.add(msg)
|
|
db.commit()
|
|
db.close()
|
|
|
|
return output
|
|
|
|
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)
|
|
|