80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
import os
|
|
from datetime import datetime
|
|
|
|
from dashscope import ImageSynthesis
|
|
from http import HTTPStatus
|
|
|
|
from sqlalchemy import desc
|
|
|
|
from llm.factory import get_adapter
|
|
from models.ai import Drawing
|
|
from sqlalchemy.orm import Session
|
|
|
|
|
|
def create_drawing_task(db: Session, user, platform: str, model: str, prompt: str, size: str, rsp,
|
|
options: str = None):
|
|
# 写入数据库
|
|
|
|
drawing = Drawing(
|
|
user_id=user['user_id'],
|
|
creator=user['username'],
|
|
modifier=user['username'],
|
|
platform=platform,
|
|
model=model,
|
|
prompt=prompt,
|
|
width=int(size.split('*')[0]),
|
|
height=int(size.split('*')[1]),
|
|
options=options,
|
|
status=rsp['output']['task_status'],
|
|
task_id=rsp['output']['task_id'],
|
|
error_message=rsp['message']
|
|
)
|
|
db.add(drawing)
|
|
db.commit()
|
|
db.refresh(drawing)
|
|
return drawing
|
|
|
|
def fetch_drawing_task_status(db: Session, drawing_id: int):
|
|
drawing = db.query(Drawing).filter(Drawing.id == drawing_id).first()
|
|
if not drawing or not drawing.task_id:
|
|
return None, "任务不存在"
|
|
if drawing.status in ("PENDING", 'RUNNING'):
|
|
api_key = os.getenv("DASHSCOPE_API_KEY")
|
|
adapter = get_adapter('tongyi', api_key=api_key, model='')
|
|
rsp = adapter.fetch_drawing_task_status(drawing.task_id)
|
|
if rsp['status_code'] == HTTPStatus.OK:
|
|
# 可根据 status.output.task_status 更新数据库
|
|
if rsp['output']['task_status'] == 'SUCCEEDED':
|
|
drawing.update_time = datetime.now()
|
|
drawing.status = rsp['output']['task_status']
|
|
drawing.pic_url = rsp['output']['results'][0]['url']
|
|
db.commit()
|
|
db.refresh(drawing)
|
|
elif rsp['output']['task_status'] == 'FAILED':
|
|
drawing.update_time = datetime.now()
|
|
drawing.status = rsp['output']['task_status']
|
|
drawing.error_message = rsp['output']['message']
|
|
db.commit()
|
|
db.refresh(drawing)
|
|
elif rsp['output']['task_status'] == 'RUNNING':
|
|
drawing.update_time = datetime.now()
|
|
drawing.status = rsp['output']['task_status']
|
|
db.commit()
|
|
db.refresh(drawing)
|
|
return drawing, None
|
|
else:
|
|
return drawing, None
|
|
|
|
|
|
def get_drawing_page(db: Session, user_id: int = None, page: int = 1, page_size: int = 12):
|
|
query = db.query(Drawing)
|
|
if user_id:
|
|
query = query.filter(Drawing.user_id == user_id)
|
|
total = query.count()
|
|
items = query.order_by(desc(Drawing.id)).offset((page - 1) * page_size).limit(page_size).all()
|
|
return {
|
|
'total': total,
|
|
'page': page,
|
|
'page_size': page_size,
|
|
'items': items
|
|
} |