259 lines
9.2 KiB
Python
259 lines
9.2 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.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.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
|
|
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 soft_delete_chat_task_resources
|
|
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
|
|
|
|
|
|
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,
|
|
) -> list[ModuleGenerationStep]:
|
|
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())
|
|
for step in steps:
|
|
step.is_current = False
|
|
step.deleted_at = deleted_at
|
|
if step.chat_task_id:
|
|
chat_result = await db.execute(
|
|
select(ChatGenerationTask)
|
|
.where(ChatGenerationTask.id == step.chat_task_id, ChatGenerationTask.deleted_at.is_(None))
|
|
.with_for_update()
|
|
.limit(1)
|
|
)
|
|
chat_task = chat_result.scalar_one_or_none()
|
|
if chat_task:
|
|
if chat_task.status == "completed":
|
|
await soft_delete_chat_task_resources(db, chat_task.id, deleted_at=deleted_at)
|
|
elif chat_task.status != "failed":
|
|
await mark_chat_generation_task_failed_and_refund_once(
|
|
db,
|
|
task=chat_task,
|
|
error_message=config.cancel_chat_task_error_message,
|
|
pipeline_stage="failed",
|
|
)
|
|
chat_task.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]},
|
|
)
|
|
return steps
|
|
|
|
|
|
async def chat_tasks_by_id(db: AsyncSession, steps: list[ModuleGenerationStep]) -> dict[str, ChatGenerationTask]:
|
|
ids = [step.chat_task_id for step in steps if step.chat_task_id]
|
|
if not ids:
|
|
return {}
|
|
result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id.in_(ids)))
|
|
return {task.id: task for task in result.scalars().all()}
|