114 lines
5.6 KiB
Python
114 lines
5.6 KiB
Python
from app.tasks.async_runner import run_async
|
|
import json
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.models.base import async_session
|
|
from app.models.chat_generation_task import ChatGenerationTask
|
|
from app.services.error_codes import extract_error_message
|
|
from app.services.generation_billing_service import charge_chatapi_prompt_usage
|
|
from app.services.generation_log_service import log_task_event
|
|
from app.services.generation_prompt_service import build_prompt_with_chatapi
|
|
from app.services.generation_provider_service import create_provider_task
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
|
|
async def _run(task_id: str):
|
|
async with async_session() as db:
|
|
result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id))
|
|
task = result.scalar_one_or_none()
|
|
if not task or task.generation_mode != "chatapi_async":
|
|
return
|
|
if task.status != "generating":
|
|
return
|
|
if task.deadline_at and datetime.now(timezone.utc) > task.deadline_at:
|
|
task.status = "failed"
|
|
task.pipeline_stage = "timeout"
|
|
task.error_message = "任务超时"
|
|
await db.commit()
|
|
await log_task_event(task, event_type="TASK_TIMEOUT", to_status="failed", to_stage="timeout")
|
|
return
|
|
if task.pipeline_stage not in ("queued", "preparing", "creating_provider_task"):
|
|
return
|
|
|
|
try:
|
|
if not task.optimized_prompt:
|
|
old_stage = task.pipeline_stage
|
|
task.pipeline_stage = "preparing"
|
|
await db.commit()
|
|
await log_task_event(task, event_type="CHATAPI_START", from_stage=old_stage, to_stage="preparing")
|
|
|
|
optimized, usage = await build_prompt_with_chatapi(db, task)
|
|
await charge_chatapi_prompt_usage(db, record=task, usage=usage, project_name="AI生成任务")
|
|
task.optimized_prompt = optimized
|
|
task.text_tokens_used = usage["total_tokens"]
|
|
await db.commit()
|
|
await log_task_event(task, event_type="CHATAPI_SUCCESS", to_stage="preparing", detail=usage)
|
|
|
|
if task.seedance_task_id or task.provider_task_id:
|
|
task.pipeline_stage = "waiting_remote"
|
|
await db.commit()
|
|
elif task.remote_result_url:
|
|
task.pipeline_stage = "result_ready"
|
|
await db.commit()
|
|
from app.tasks.generation_download_tasks import download_generation_result_task
|
|
download_generation_result_task.delay(task.id)
|
|
return
|
|
else:
|
|
old_stage = task.pipeline_stage
|
|
task.pipeline_stage = "creating_provider_task"
|
|
await db.commit()
|
|
await log_task_event(task, event_type="PROVIDER_CREATE_START", from_stage=old_stage, to_stage="creating_provider_task")
|
|
|
|
created = await create_provider_task(db, task)
|
|
provider_task_id = created.get("task_id")
|
|
if provider_task_id:
|
|
task.provider_task_id = provider_task_id
|
|
task.seedance_task_id = provider_task_id
|
|
task.remote_result_url = created.get("remote_result_url") or task.remote_result_url
|
|
if task.gen_type == "image":
|
|
task.image_tokens_used = created.get("image_tokens", task.image_tokens_used or 0) or 0
|
|
task.provider_response_json = json.dumps(created.get("response_data") or {}, ensure_ascii=False, default=str)
|
|
if task.remote_result_url and not task.seedance_task_id:
|
|
# Sync image path: original SDK already returned final URL.
|
|
task.pipeline_stage = "result_ready"
|
|
else:
|
|
# Video path: provider returns task id, poll later.
|
|
task.pipeline_stage = "waiting_remote"
|
|
task.status = "generating"
|
|
await db.commit()
|
|
await log_task_event(task, event_type="PROVIDER_CREATE_SUCCESS", to_stage=task.pipeline_stage, detail=created)
|
|
|
|
if task.pipeline_stage == "result_ready":
|
|
from app.tasks.generation_download_tasks import download_generation_result_task
|
|
download_generation_result_task.delay(task.id)
|
|
else:
|
|
from app.tasks.generation_poll_tasks import poll_generation_task
|
|
poll_generation_task.delay(task.id)
|
|
except Exception as exc:
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id))
|
|
task = result.scalar_one_or_none()
|
|
if task:
|
|
task.status = "failed"
|
|
task.error_message = extract_error_message(exc, "生成任务") if callable(extract_error_message) else str(exc)
|
|
await db.commit()
|
|
await log_task_event(task, event_type="TASK_FAILED", message=task.error_message)
|
|
|
|
|
|
if celery_app:
|
|
@celery_app.task(name="generation.chatapi_create_generation_task", bind=True, max_retries=3, default_retry_delay=30)
|
|
def chatapi_create_generation_task(self, task_id: str):
|
|
return run_async(_run(task_id))
|
|
else:
|
|
class _DisabledTask:
|
|
def delay(self, *args, **kwargs):
|
|
raise RuntimeError("Celery is disabled")
|
|
def apply_async(self, *args, **kwargs):
|
|
raise RuntimeError("Celery is disabled")
|
|
chatapi_create_generation_task = _DisabledTask()
|