211 lines
8.5 KiB
Python
211 lines
8.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.models.chat_generation_task import ChatGenerationTask
|
|
from app.models.user import User
|
|
from app.schemas.generation_ai import GenerationAIReference, GenerationAITaskCreate
|
|
from app.services.generation_ai_service import (
|
|
IMAGE_DEFAULT_PROPORTION,
|
|
IMAGE_DEFAULT_PX,
|
|
IMAGE_DEFAULT_SIZE,
|
|
VIDEO_DEFAULT_RATIO,
|
|
VIDEO_DEFAULT_RESOLUTION,
|
|
_build_image_snapshot,
|
|
_build_video_snapshot,
|
|
_get_image_engine,
|
|
_get_video_engine,
|
|
_image_supported_sizes,
|
|
_parse_list,
|
|
normalize_px,
|
|
)
|
|
from app.services.generation_billing_service import OWNER_CHAT_GENERATION_TASK, charge_generation_media_by_params
|
|
from app.services.resource_capacity_service import assert_user_resource_capacity_available
|
|
from app.services.private_portrait.reference_resolver import resolve_private_portrait_references
|
|
from app.utils.id_gen import generate_id
|
|
|
|
|
|
def _json(data: Any) -> str | None:
|
|
if data is None:
|
|
return None
|
|
return json.dumps(data, ensure_ascii=False, default=str)
|
|
|
|
|
|
def _build_backend_idempotency_key(*, generation_mode: str, gen_type: str, task_id: str) -> str:
|
|
"""模块生成关联 ChatGenerationTask 的幂等键由后端生成。
|
|
|
|
不再接收前端透传，避免 user_id + generation_mode + idempotency_key
|
|
唯一索引被前端固定 key 或重复 key 拦截。
|
|
"""
|
|
return f"{generation_mode}:{gen_type}:{task_id}"[:64]
|
|
|
|
|
|
async def create_chat_generation_task_for_module(
|
|
db: AsyncSession,
|
|
*,
|
|
current_user: User,
|
|
generation_mode: str,
|
|
gen_type: str,
|
|
original_prompt: str,
|
|
optimized_prompt: str | None = None,
|
|
engine_id: str | None = None,
|
|
media_references: list[dict[str, Any]] | None = None,
|
|
idempotency_key: str | None = None,
|
|
image_size: str | None = None,
|
|
image_proportion: str | None = None,
|
|
image_px: str | None = None,
|
|
duration: int | None = None,
|
|
aspect_ratio: str | None = None,
|
|
resolution: str | None = None,
|
|
billing_project_name: str = "模块生成任务",
|
|
billing_description_prefix: str = "模块生成-",
|
|
billing_source_module: str | None = None,
|
|
billing_source_project_id: str | None = None,
|
|
billing_source_step_id: str | None = None,
|
|
billing_source_step_code: str | None = None,
|
|
billing_scene: str | None = None,
|
|
) -> ChatGenerationTask:
|
|
"""创建可复用的 ChatGenerationTask 子任务。
|
|
|
|
和 /generation-ai 普通任务不同，generation_mode 由业务模块传入，
|
|
但仍复用同一套引擎校验、扣费、Celery 创建/轮询/下载逻辑。
|
|
"""
|
|
gen_type = gen_type.lower().strip()
|
|
if gen_type not in ("image", "video"):
|
|
raise HTTPException(status_code=400, detail="gen_type 仅支持 image 或 video")
|
|
|
|
task_id = generate_id()
|
|
now = datetime.now(timezone.utc)
|
|
refs = media_references or []
|
|
refs = await resolve_private_portrait_references(
|
|
db,
|
|
user_id=current_user.id,
|
|
media_references=refs,
|
|
gen_type=gen_type,
|
|
)
|
|
backend_idempotency_key = _build_backend_idempotency_key(
|
|
generation_mode=generation_mode,
|
|
gen_type=gen_type,
|
|
task_id=task_id,
|
|
)
|
|
|
|
await assert_user_resource_capacity_available(db, current_user.id)
|
|
|
|
if gen_type == "image":
|
|
engine = await _get_image_engine(db, engine_id)
|
|
sizes = _image_supported_sizes(engine)
|
|
size = image_size or engine.default_size or IMAGE_DEFAULT_SIZE
|
|
proportion = image_proportion or IMAGE_DEFAULT_PROPORTION
|
|
px = normalize_px(image_px)
|
|
if sizes:
|
|
if size not in sizes:
|
|
raise HTTPException(status_code=400, detail=f"图片分辨率档位不支持: {size}")
|
|
if proportion not in sizes.get(size, {}):
|
|
raise HTTPException(status_code=400, detail=f"图片比例不支持: {proportion}")
|
|
px = px or normalize_px((sizes.get(size) or {}).get(proportion))
|
|
px = px or IMAGE_DEFAULT_PX
|
|
media_billing = await charge_generation_media_by_params(
|
|
db,
|
|
user_id=current_user.id,
|
|
record_id=task_id,
|
|
gen_type="image",
|
|
image_size=size,
|
|
engine_id=engine.id,
|
|
project_name=billing_project_name,
|
|
description_prefix=billing_description_prefix,
|
|
owner_type=OWNER_CHAT_GENERATION_TASK,
|
|
attempt_no=1,
|
|
source_module=billing_source_module,
|
|
source_project_id=billing_source_project_id,
|
|
source_step_id=billing_source_step_id,
|
|
source_step_code=billing_source_step_code,
|
|
billing_scene=billing_scene,
|
|
)
|
|
snapshot = _build_image_snapshot(engine, size, proportion, px)
|
|
task = ChatGenerationTask(
|
|
id=task_id,
|
|
user_id=current_user.id,
|
|
original_prompt=original_prompt,
|
|
optimized_prompt=optimized_prompt,
|
|
gen_type="image",
|
|
image_size=size,
|
|
image_proportion=proportion,
|
|
image_px=px,
|
|
status="generating",
|
|
generation_mode=generation_mode,
|
|
pipeline_stage="queued",
|
|
engine_id=engine.id,
|
|
engine_snapshot_json=_json(snapshot),
|
|
media_references=_json(refs) if refs else None,
|
|
credits_cost=round(media_billing.total_charged, 2),
|
|
idempotency_key=backend_idempotency_key,
|
|
deadline_at=now + timedelta(minutes=settings.CHATAPI_ASYNC_IMAGE_DEADLINE_MINUTES),
|
|
)
|
|
else:
|
|
engine = await _get_video_engine(db, engine_id)
|
|
ratio = aspect_ratio or VIDEO_DEFAULT_RATIO
|
|
selected_resolution = resolution or VIDEO_DEFAULT_RESOLUTION
|
|
selected_duration = duration or 4
|
|
ratios = _parse_list(engine.supported_ratios, [])
|
|
resolutions = _parse_list(engine.supported_resolutions, [])
|
|
durations = _parse_list(engine.supported_durations, [])
|
|
if ratios and ratio not in ratios:
|
|
raise HTTPException(status_code=400, detail=f"视频比例不支持: {ratio}")
|
|
if resolutions and selected_resolution not in resolutions:
|
|
raise HTTPException(status_code=400, detail=f"视频分辨率不支持: {selected_resolution}")
|
|
if durations and selected_duration not in durations:
|
|
raise HTTPException(status_code=400, detail=f"视频时长不支持: {selected_duration}")
|
|
if engine.max_duration and selected_duration > engine.max_duration:
|
|
raise HTTPException(status_code=400, detail=f"视频时长不能超过 {engine.max_duration} 秒")
|
|
media_billing = await charge_generation_media_by_params(
|
|
db,
|
|
user_id=current_user.id,
|
|
record_id=task_id,
|
|
gen_type="video",
|
|
duration=selected_duration,
|
|
resolution=selected_resolution,
|
|
engine_id=engine.id,
|
|
project_name=billing_project_name,
|
|
description_prefix=billing_description_prefix,
|
|
owner_type=OWNER_CHAT_GENERATION_TASK,
|
|
attempt_no=1,
|
|
source_module=billing_source_module,
|
|
source_project_id=billing_source_project_id,
|
|
source_step_id=billing_source_step_id,
|
|
source_step_code=billing_source_step_code,
|
|
billing_scene=billing_scene,
|
|
)
|
|
snapshot = _build_video_snapshot(engine, ratio, selected_resolution, selected_duration)
|
|
task = ChatGenerationTask(
|
|
id=task_id,
|
|
user_id=current_user.id,
|
|
original_prompt=original_prompt,
|
|
optimized_prompt=optimized_prompt,
|
|
gen_type="video",
|
|
duration=selected_duration,
|
|
aspect_ratio=ratio,
|
|
resolution=selected_resolution,
|
|
image_size=image_size or IMAGE_DEFAULT_SIZE,
|
|
image_proportion=image_proportion or IMAGE_DEFAULT_PROPORTION,
|
|
image_px=normalize_px(image_px) or IMAGE_DEFAULT_PX,
|
|
status="generating",
|
|
generation_mode=generation_mode,
|
|
pipeline_stage="queued",
|
|
engine_id=engine.id,
|
|
engine_snapshot_json=_json(snapshot),
|
|
media_references=_json(refs) if refs else None,
|
|
credits_cost=round(media_billing.total_charged, 2),
|
|
idempotency_key=backend_idempotency_key,
|
|
deadline_at=now + timedelta(hours=settings.CHATAPI_ASYNC_VIDEO_FINAL_DEADLINE_HOURS),
|
|
)
|
|
|
|
db.add(task)
|
|
await db.flush()
|
|
return task
|