AI创作批量生成任务 main V1 init

2026-07-15 13:03:13 +08:00
parent 5b937f652b
commit 51f9deecde
59 changed files with 8091 additions and 413 deletions
@@ -0,0 +1,212 @@
+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.engine_service import (
+    IMAGE_DEFAULT_PROPORTION,
+    IMAGE_DEFAULT_PX,
+    IMAGE_DEFAULT_SIZE,
+    VIDEO_DEFAULT_RATIO,
+    VIDEO_DEFAULT_RESOLUTION,
+    build_image_snapshot as _build_image_snapshot,
+    build_video_snapshot as _build_video_snapshot,
+    get_image_engine as _get_image_engine,
+    get_video_engine as _get_video_engine,
+    image_supported_sizes as _image_supported_sizes,
+    normalize_px,
+    parse_json_list as _parse_list,
+)
+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)
+        snapshot["generation_count"] = 1
+        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)
+        snapshot["generation_count"] = 1
+        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