449 lines
16 KiB
Python
449 lines
16 KiB
Python
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.generation.pipeline.db_lock_service import (
|
|
apply_short_lock_timeout,
|
|
execute_with_lock_timeout,
|
|
raise_if_database_lock_busy,
|
|
)
|
|
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:
|
|
await apply_short_lock_timeout(db)
|
|
query = query.with_for_update()
|
|
try:
|
|
result = await db.execute(query.limit(1))
|
|
except Exception as exc:
|
|
raise_if_database_lock_busy(exc)
|
|
raise
|
|
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:
|
|
await apply_short_lock_timeout(db)
|
|
query = query.with_for_update()
|
|
try:
|
|
result = await db.execute(query.limit(1))
|
|
except Exception as exc:
|
|
raise_if_database_lock_busy(exc)
|
|
raise
|
|
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 execute_with_lock_timeout(db, stmt)
|
|
else:
|
|
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 execute_with_lock_timeout(
|
|
db,
|
|
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 execute_with_lock_timeout(
|
|
db,
|
|
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)
|