xiaoxia-saas/apps/api/app/api/routes/generation_tasks.py

from typing import Any

from app.api.routes.permissions import require_workspace_member
from app.auth import AuthenticatedUser, get_current_user
from app.core.celery_app import celery_app
from app.dependencies import (
    get_asset_library_repository,
    get_asset_repository,
    get_generated_video_repository,
    get_generation_task_repository,
    get_project_repository,
    get_project_title_repository,
    get_workspace_member_repository,
)
from app.schemas.generated_video import (
    GeneratedVideoResponse,
    ListGeneratedVideosResponse,
)
from app.schemas.generation_task import (
    CreateGenerationTaskRequest,
    GenerationTaskResponse,
)
from fastapi import APIRouter, Depends, HTTPException

from packages.application import (
    CreateGenerationTaskCommand,
    CreateGenerationTaskUseCase,
    GetGenerationTaskUseCase,
    ListGeneratedVideosByTaskUseCase,
)
from packages.ports.workspace_member_repository import WorkspaceMemberRepository

router = APIRouter()


def _to_generation_task_response(task) -> GenerationTaskResponse:
    return GenerationTaskResponse(
        id=task.id,
        workspace_id=task.workspace_id,
        project_id=task.project_id,
        asset_library_id=task.asset_library_id,
        strategy_id=task.strategy_id,
        voice_library_id=task.voice_library_id,
        edit_plan_id=task.edit_plan_id,
        status=task.status.value,
        progress=task.progress,
        result_count=task.result_count,
        error_message=task.error_message,
    )


def _to_generated_video_response(item) -> GeneratedVideoResponse:
    return GeneratedVideoResponse(
        id=item.id,
        workspace_id=item.workspace_id,
        project_id=item.project_id,
        generation_task_id=item.generation_task_id,
        name=item.name,
        file_url=item.file_url,
        file_size=item.file_size,
        duration=item.duration,
        thumbnail_url=item.thumbnail_url,
        width=item.width,
        height=item.height,
        fps=item.fps,
    )


def _ensure_library_has_ready_video_assets(assets) -> None:
    ready_video_assets = [
        asset for asset in assets if asset.status.value == "ready" and asset.mime_type.startswith("video")
    ]
    if not ready_video_assets:
        raise HTTPException(
            status_code=422,
            detail="当前素材库没有 ready 状态的视频素材，请先上传并等待导入完成后再生成。",
        )


def _select_title_id(project_title_repository: Any, project_id: str) -> str:
    active_titles = project_title_repository.list_by_project(project_id, active_only=True)
    if not active_titles:
        return ""
    selected = sorted(
        active_titles,
        key=lambda title: (
            0 if getattr(title, "favorite", False) else 1,
            int(title.usage_count or 0),
            title.created_at,
        ),
    )[0]
    return selected.id


@router.post("/tasks", response_model=GenerationTaskResponse)
def create_generation_task(
    request: CreateGenerationTaskRequest,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    project_repository: Any = Depends(get_project_repository),
    asset_library_repository: Any = Depends(get_asset_library_repository),
    asset_repository: Any = Depends(get_asset_repository),
    project_title_repository: Any = Depends(get_project_title_repository),
    workspace_member_repository: WorkspaceMemberRepository = Depends(get_workspace_member_repository),
) -> GenerationTaskResponse:
    require_workspace_member(request.workspace_id, authenticated_user, workspace_member_repository)
    project = project_repository.find_by_id(request.project_id)
    if project is None or project.workspace_id != request.workspace_id:
        raise HTTPException(status_code=404, detail=f"Project {request.project_id} not found")
    library = asset_library_repository.get(request.asset_library_id)
    if library is None or library.workspace_id != request.workspace_id or library.project_id != request.project_id:
        raise HTTPException(status_code=404, detail=f"AssetLibrary {request.asset_library_id} not found")
    assets = asset_repository.list_by_library(request.asset_library_id)
    _ensure_library_has_ready_video_assets(assets)

    use_case = CreateGenerationTaskUseCase(generation_task_repository)
    task = use_case.execute(
        CreateGenerationTaskCommand(
            workspace_id=request.workspace_id,
            project_id=request.project_id,
            asset_library_id=request.asset_library_id,
            strategy_id=request.strategy_id or _select_title_id(project_title_repository, request.project_id),
            voice_library_id=request.voice_library_id,
            edit_plan_id=request.edit_plan_id,
            created_by_user_id=authenticated_user.user.id,
        )
    )
    celery_app.send_task("worker.generate_video", args=[task.id])
    return _to_generation_task_response(task)


@router.get("/tasks/{task_id}", response_model=GenerationTaskResponse)
def get_generation_task(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    workspace_member_repository: WorkspaceMemberRepository = Depends(get_workspace_member_repository),
) -> GenerationTaskResponse:
    use_case = GetGenerationTaskUseCase(generation_task_repository)
    task = use_case.execute(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail=f"GenerationTask {task_id} not found")
    require_workspace_member(task.workspace_id, authenticated_user, workspace_member_repository)
    return _to_generation_task_response(task)


@router.get("/tasks/{task_id}/results", response_model=ListGeneratedVideosResponse)
def list_generation_results(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    generated_video_repository: Any = Depends(get_generated_video_repository),
    workspace_member_repository: WorkspaceMemberRepository = Depends(get_workspace_member_repository),
) -> ListGeneratedVideosResponse:
    task = generation_task_repository.get(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail=f"GenerationTask {task_id} not found")
    require_workspace_member(task.workspace_id, authenticated_user, workspace_member_repository)
    use_case = ListGeneratedVideosByTaskUseCase(generated_video_repository)
    items = use_case.execute(task_id)
    return ListGeneratedVideosResponse(items=[_to_generated_video_response(item) for item in items])