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.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 [] backend_idempotency_key = _build_backend_idempotency_key( generation_mode=generation_mode, gen_type=gen_type, task_id=task_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(minutes=settings.CHATAPI_ASYNC_VIDEO_DEADLINE_MINUTES), ) db.add(task) await db.flush() return task