from __future__ import annotations from collections.abc import Awaitable, Callable from datetime import datetime from typing import Any from fastapi import HTTPException from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.enums.common import ModuleEventTypeEnum, ModuleStepStatusEnum from app.enums.generation_task import ChatGenerationPipelineStage, ChatGenerationTaskStatus from app.models.chat_generation_task import ChatGenerationTask from app.models.module_generation_project import ModuleGenerationProject from app.models.module_generation_step import ModuleGenerationStep from app.models.user import User from app.enums.module_generation_flow import ModuleGenerationFlowConfig from app.services.module_generation_step_common_service import build_step_input, build_step_output, utc_now from app.services.resource_accounting_service import ( SOURCE_MODEL_CHAT_TASK, soft_delete_resources_by_source, ) from app.utils.id_gen import generate_id LogModuleEventCallable = Callable[..., Awaitable[None]] async def get_project_for_user( db: AsyncSession, *, project_id: str, user: User, config: ModuleGenerationFlowConfig, for_update: bool = False, populate_existing: bool = False, ) -> ModuleGenerationProject: query = select(ModuleGenerationProject).where( ModuleGenerationProject.id == project_id, ModuleGenerationProject.module == config.module, ModuleGenerationProject.deleted_at.is_(None), ) if not user.is_admin: query = query.where(ModuleGenerationProject.user_id == user.id) if populate_existing: query = query.execution_options(populate_existing=True) if for_update: query = query.with_for_update() result = await db.execute(query.limit(1)) project = result.scalar_one_or_none() if not project: raise HTTPException(status_code=404, detail=config.project_not_found_message) return project async def get_step_for_user( db: AsyncSession, *, project_id: str, step_id: str, user: User, config: ModuleGenerationFlowConfig, for_update: bool = False, ) -> ModuleGenerationStep: await get_project_for_user(db, project_id=project_id, user=user, config=config, for_update=for_update) query = select(ModuleGenerationStep).where( ModuleGenerationStep.id == step_id, ModuleGenerationStep.project_id == project_id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.deleted_at.is_(None), ModuleGenerationStep.is_current == True, ) if not user.is_admin: query = query.where(ModuleGenerationStep.user_id == user.id) if for_update: query = query.with_for_update() result = await db.execute(query.limit(1)) step = result.scalar_one_or_none() if not step: raise HTTPException(status_code=404, detail=config.step_not_found_message) return step async def get_current_steps(db: AsyncSession, *, project_id: str, config: ModuleGenerationFlowConfig) -> list[ModuleGenerationStep]: result = await db.execute( select(ModuleGenerationStep) .where( ModuleGenerationStep.project_id == project_id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.deleted_at.is_(None), ModuleGenerationStep.is_current == True, ) .order_by(ModuleGenerationStep.step_index.asc(), ModuleGenerationStep.created_at.asc()) ) return list(result.scalars().all()) async def get_current_step_by_code( db: AsyncSession, *, project_id: str, step_code: str, config: ModuleGenerationFlowConfig, ) -> ModuleGenerationStep | None: result = await db.execute( select(ModuleGenerationStep) .where( ModuleGenerationStep.project_id == project_id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.step_code == step_code, ModuleGenerationStep.is_current == True, ModuleGenerationStep.deleted_at.is_(None), ) .order_by(ModuleGenerationStep.version.desc(), ModuleGenerationStep.created_at.desc()) .limit(1) ) return result.scalar_one_or_none() async def next_version( db: AsyncSession, *, project_id: str, step_code: str, config: ModuleGenerationFlowConfig, ) -> int: result = await db.execute( select(func.max(ModuleGenerationStep.version)).where( ModuleGenerationStep.project_id == project_id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.step_code == step_code, ) ) return int(result.scalar_one_or_none() or 0) + 1 async def create_module_step( db: AsyncSession, *, project: ModuleGenerationProject, step_code: str, config: ModuleGenerationFlowConfig, log_module_event: LogModuleEventCallable, status: str = ModuleStepStatusEnum.PENDING.value, parent_step_id: str | None = None, source_step_id: str | None = None, chat_task_id: str | None = None, input_data: dict[str, Any] | None = None, output_data: dict[str, Any] | None = None, ) -> ModuleGenerationStep: version = await next_version(db, project_id=project.id, step_code=step_code, config=config) step = ModuleGenerationStep( id=generate_id(), project_id=project.id, user_id=project.user_id, module=project.module, step_index=config.step_index_map[step_code], step_code=step_code, status=status, version=version, is_current=True, parent_step_id=parent_step_id, source_step_id=source_step_id, chat_task_id=chat_task_id, input_json=( build_step_input( step_code=step_code, payload=input_data, source_step_id=source_step_id, parent_step_id=parent_step_id, schema_version=config.step_io_schema_version, ) if input_data is not None else None ), output_json=( build_step_output( step_code=step_code, status=status, result=output_data, schema_version=config.step_io_schema_version, ) if output_data is not None else None ), started_at=utc_now() if status == ModuleStepStatusEnum.PROCESSING.value else None, completed_at=utc_now() if status == ModuleStepStatusEnum.COMPLETED.value else None, ) db.add(step) project.current_step_code = step_code await db.flush() await log_module_event( db, project=project, step=step, event_type=ModuleEventTypeEnum.STEP_CREATED.value, detail={"step_code": step_code, "version": version}, ) return step def _clear_project_final_resources_by_deleted_steps( project: ModuleGenerationProject, *, deleted_step_codes: set[str], config: ModuleGenerationFlowConfig, ) -> list[str]: """根据被软删的步骤,清空项目表中对应的最终资源字段。 第 3 步图片生成被软删时,下游第 5 步视频也基于旧图片失效,所以图片、视频、封面都要清空。 第 5 步视频生成被软删时,只清空视频和封面,保留第 3 步图片结果。 """ cleared_fields: list[str] = [] def clear_field(field_name: str) -> None: if getattr(project, field_name, None): cleared_fields.append(field_name) setattr(project, field_name, None) if config.image_generate_step_code in deleted_step_codes: clear_field("final_image_url") clear_field("final_video_url") clear_field("final_video_cover_url") elif config.video_generate_step_code in deleted_step_codes: clear_field("final_video_url") clear_field("final_video_cover_url") return cleared_fields ACTIVE_CHAT_TASK_BLOCK_STATUSES = { ChatGenerationTaskStatus.PENDING.value, ChatGenerationTaskStatus.GENERATING.value, } ACTIVE_CHAT_TASK_BLOCK_STAGES = { ChatGenerationPipelineStage.QUEUED.value, ChatGenerationPipelineStage.PREPARING.value, ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value, ChatGenerationPipelineStage.WAITING_REMOTE.value, ChatGenerationPipelineStage.POLLING.value, ChatGenerationPipelineStage.RESULT_READY.value, ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value, ChatGenerationPipelineStage.DOWNLOADING.value, ChatGenerationPipelineStage.RETRY_WAITING.value, ChatGenerationPipelineStage.UPSCALE_QUEUED.value, ChatGenerationPipelineStage.UPSCALE_PROCESSING.value, ChatGenerationPipelineStage.UPSCALE_POLLING.value, ChatGenerationPipelineStage.UPSCALE_DOWNLOADING.value, ChatGenerationPipelineStage.UPSCALE_FINALIZING.value, ChatGenerationPipelineStage.UPSCALE_RETRY_WAITING.value, } def is_active_chat_generation_task(task: ChatGenerationTask) -> bool: """判断 ChatGenerationTask 是否仍处于不可主动删除/废弃的处理中状态。""" return bool( task.status in ACTIVE_CHAT_TASK_BLOCK_STATUSES or (task.pipeline_stage in ACTIVE_CHAT_TASK_BLOCK_STAGES) ) def _unique_ids(values: list[str | None] | tuple[str | None, ...]) -> list[str]: return list(dict.fromkeys(str(value) for value in values if value)) async def load_chat_tasks_for_steps( db: AsyncSession, steps: list[ModuleGenerationStep], *, for_update: bool = False, ) -> dict[str, ChatGenerationTask]: ids = _unique_ids([step.chat_task_id for step in steps]) if not ids: return {} stmt = select(ChatGenerationTask).where( ChatGenerationTask.id.in_(ids), ChatGenerationTask.deleted_at.is_(None), ) if for_update: stmt = stmt.with_for_update() result = await db.execute(stmt) return {task.id: task for task in result.scalars().all()} async def assert_no_active_chat_tasks_for_steps( db: AsyncSession, steps: list[ModuleGenerationStep], *, detail_message: str = "当前存在生成中任务,请等待生成完成或失败后再操作", for_update: bool = True, ) -> dict[str, ChatGenerationTask]: """批量校验步骤关联任务是否有进行中任务。 用户主动修改上游步骤、删除项目、删除历史时,统一采用“生成中拦截、不退款”。 返回已批量加载的未删除 ChatGenerationTask,调用方可继续复用,避免重复查询。 """ task_map = await load_chat_tasks_for_steps(db, steps, for_update=for_update) active_task_ids = [task.id for task in task_map.values() if is_active_chat_generation_task(task)] if active_task_ids: raise HTTPException( status_code=409, detail={ "message": detail_message, "active_count": len(active_task_ids), }, ) return task_map async def assert_project_has_no_active_chat_tasks( db: AsyncSession, *, project: ModuleGenerationProject, config: ModuleGenerationFlowConfig, detail_message: str = "当前存在生成中任务,请等待生成完成或失败后再操作", ) -> dict[str, ChatGenerationTask]: result = await db.execute( select(ModuleGenerationStep) .where( ModuleGenerationStep.project_id == project.id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.deleted_at.is_(None), ModuleGenerationStep.chat_task_id.is_not(None), ) .with_for_update() ) steps = list(result.scalars().all()) return await assert_no_active_chat_tasks_for_steps( db, steps, detail_message=detail_message, for_update=True, ) async def soft_delete_steps_from_index( db: AsyncSession, *, project: ModuleGenerationProject, start_index: int, config: ModuleGenerationFlowConfig, log_module_event: LogModuleEventCallable, deleted_at: datetime | None = None, refund_unfinished: bool = False, block_active_tasks: bool = True, release_stats: dict[str, int] | None = None, ) -> list[ModuleGenerationStep]: """软删指定步骤及其后续当前版本步骤。 说明: - 兼容旧调用方保留 refund_unfinished 参数,但用户主动修改/删除链路不再退款。 - 存在 pending/generating/下载中/轮询中等任务时直接 409 拦截。 - 已完成任务只做软删任务与 generated_resources,释放容量统计;失败任务只软删任务。 """ deleted_at = deleted_at or utc_now() result = await db.execute( select(ModuleGenerationStep) .where( ModuleGenerationStep.project_id == project.id, ModuleGenerationStep.module == config.module, ModuleGenerationStep.is_current == True, ModuleGenerationStep.deleted_at.is_(None), ModuleGenerationStep.step_index >= start_index, ) .with_for_update() ) steps = list(result.scalars().all()) if not steps: return [] task_map = await load_chat_tasks_for_steps(db, steps, for_update=True) if block_active_tasks: active_task_ids = [task.id for task in task_map.values() if is_active_chat_generation_task(task)] if active_task_ids: raise HTTPException( status_code=409, detail={ "message": "当前存在生成中任务,请等待生成完成或失败后再操作", "active_count": len(active_task_ids), }, ) deleted_step_codes = {step.step_code for step in steps} cleared_project_fields = _clear_project_final_resources_by_deleted_steps( project, deleted_step_codes=deleted_step_codes, config=config, ) completed_task_ids = [ task.id for task in task_map.values() if task.status == ChatGenerationTaskStatus.COMPLETED.value ] if completed_task_ids: released_size = await soft_delete_resources_by_source( db, source_model=SOURCE_MODEL_CHAT_TASK, source_ids=completed_task_ids, deleted_at=deleted_at, ) if release_stats is not None: release_stats["released_size_bytes"] = int(release_stats.get("released_size_bytes", 0)) + int(released_size or 0) for step in steps: step.is_current = False step.deleted_at = deleted_at if step.chat_task_id and step.chat_task_id in task_map: task_map[step.chat_task_id].deleted_at = deleted_at if steps: await log_module_event( db, project=project, event_type=ModuleEventTypeEnum.SOFT_DELETE_STEPS.value, message=f"软删除第 {start_index} 步及之后的旧子任务", detail={ "step_ids": [step.id for step in steps], "step_codes": [step.step_code for step in steps], "cleared_project_fields": cleared_project_fields, "block_active_tasks": block_active_tasks, "refund_unfinished": False, }, ) return steps async def chat_tasks_by_id(db: AsyncSession, steps: list[ModuleGenerationStep]) -> dict[str, ChatGenerationTask]: return await load_chat_tasks_for_steps(db, steps, for_update=False)