Files
video-gen/video-gen-api/app/services/shot_replicate_taskset_service.py

618 lines
23 KiB
Python

from __future__ import annotations
import uuid
from datetime import datetime, timezone
from typing import Any
from fastapi import HTTPException
from app.config import settings
from sqlalchemy import String, cast, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession
from app.enums.shot_replicate import (
ModuleCodeEnum,
ShotAnalysisStatusEnum,
ShotSegmentAnalysisStatusEnum,
ShotSegmentReplicateStatusEnum,
ShotSegmentSourceModeEnum,
ShotSplitStatusEnum,
ShotTaskSetStatusEnum,
)
from app.models.module_generation_project import ModuleGenerationProject
from app.models.shot_replicate_segment import ShotReplicateSegment
from app.models.shot_replicate_task_set import ShotReplicateTaskSet
from app.models.user import User
from app.schemas.shot_replicate import (
ShotAISuggestionOut,
ShotSegmentDetailOut,
ShotSegmentListOut,
ShotSegmentOut,
ShotSplitByAIOut,
ShotSplitByAIRequest,
ShotSplitCustomOut,
ShotSplitCustomRequest,
ShotTaskSetCreate,
ShotTaskSetDetailOut,
ShotTaskSetListOut,
ShotTaskSetOut,
)
from app.services.module_generation_log_service import log_module_event_file
from app.services.upload_video_asset_service import (
build_time_node,
validate_split_range,
validate_upload_video_asset,
)
from app.tasks.celery_app import celery_app
from app.utils.id_gen import generate_id
MODULE = ModuleCodeEnum.SHOT_REPLICATE.value
def _now() -> datetime:
return datetime.now(timezone.utc)
def _normalize_suggestions(value: Any) -> list[dict[str, Any]]:
if not isinstance(value, list):
return []
normalized: list[dict[str, Any]] = []
for idx, item in enumerate(value, start=1):
if not isinstance(item, dict):
continue
start = item.get("拆镜开始秒")
end = item.get("拆镜结束秒")
try:
start_f = float(start)
end_f = float(end)
except Exception:
continue
if start_f < 0 or end_f <= start_f:
continue
normalized.append(
{
"index": idx,
"start_second": start_f,
"end_second": end_f,
"duration_seconds": round(end_f - start_f, 3),
"time_node": str(item.get("拆镜时间节点") or build_time_node(start_f, end_f)),
"content": str(item.get("对应时间节点内的内容") or "无"),
"category": str(item.get("分类") or "无"),
"audience": str(item.get("受众人群") or "无"),
"raw": item,
}
)
return normalized
def _task_set_to_out(task_set: ShotReplicateTaskSet, user_name: str | None = None) -> ShotTaskSetOut:
data = ShotTaskSetOut.model_validate(task_set)
data.user_name = user_name
return data
def _task_set_to_detail_out(task_set: ShotReplicateTaskSet, user_name: str | None = None) -> ShotTaskSetDetailOut:
suggestions = [ShotAISuggestionOut(**{k: v for k, v in item.items() if k != "raw"}) for item in _normalize_suggestions(task_set.ai_suggestion_json)]
base = ShotTaskSetDetailOut.model_validate(task_set)
base.user_name = user_name
base.ai_suggestions = suggestions
return base
def _segment_to_out(segment: ShotReplicateSegment, project: ModuleGenerationProject | None = None) -> ShotSegmentOut:
data = ShotSegmentOut.model_validate(segment)
data.segment_name = f"片段{segment.segment_index}"
if project:
data.module_project_title = project.title
data.module_project_status = project.status
data.module_project_current_step_code = project.current_step_code
return data
def _segment_to_detail_out(segment: ShotReplicateSegment, project: ModuleGenerationProject | None = None) -> ShotSegmentDetailOut:
data = ShotSegmentDetailOut.model_validate(segment)
data.segment_name = f"片段{segment.segment_index}"
if project:
data.module_project_title = project.title
data.module_project_status = project.status
data.module_project_current_step_code = project.current_step_code
return data
async def get_task_set_for_user(
db: AsyncSession,
*,
task_set_id: str,
user: User,
for_update: bool = False,
) -> ShotReplicateTaskSet:
query = select(ShotReplicateTaskSet).where(
ShotReplicateTaskSet.id == task_set_id,
ShotReplicateTaskSet.deleted_at.is_(None),
)
if not user.is_admin:
query = query.where(ShotReplicateTaskSet.user_id == user.id)
if for_update:
query = query.with_for_update()
result = await db.execute(query.limit(1))
task_set = result.scalar_one_or_none()
if not task_set:
raise HTTPException(status_code=404, detail="拆镜总任务集不存在")
return task_set
async def get_segment_for_user(
db: AsyncSession,
*,
segment_id: str,
user: User,
for_update: bool = False,
) -> ShotReplicateSegment:
query = select(ShotReplicateSegment).where(
ShotReplicateSegment.id == segment_id,
ShotReplicateSegment.deleted_at.is_(None),
)
if not user.is_admin:
query = query.where(ShotReplicateSegment.user_id == user.id)
if for_update:
query = query.with_for_update()
result = await db.execute(query.limit(1))
segment = result.scalar_one_or_none()
if not segment:
raise HTTPException(status_code=404, detail="拆镜片段不存在")
return segment
async def create_task_set(db: AsyncSession, *, current_user: User, req: ShotTaskSetCreate) -> ShotReplicateTaskSet:
if req.idempotency_key:
existing_result = await db.execute(
select(ShotReplicateTaskSet).where(
ShotReplicateTaskSet.user_id == current_user.id,
ShotReplicateTaskSet.idempotency_key == req.idempotency_key,
ShotReplicateTaskSet.deleted_at.is_(None),
).limit(1)
)
existing = existing_result.scalar_one_or_none()
if existing:
return existing
asset = validate_upload_video_asset(req.video_url, req.video_duration_seconds)
task_set = ShotReplicateTaskSet(
id=generate_id(),
user_id=current_user.id,
title=req.title or "拆镜复刻任务",
video_url=asset.url,
video_path=str(asset.path),
video_duration_seconds=asset.duration_seconds,
status=ShotTaskSetStatusEnum.PENDING_ANALYSIS.value,
analysis_status=ShotAnalysisStatusEnum.PENDING.value,
split_status=ShotSplitStatusEnum.NONE.value,
segment_count=0,
completed_segment_count=0,
failed_segment_count=0,
idempotency_key=req.idempotency_key,
)
db.add(task_set)
await db.flush()
log_module_event_file(
module=MODULE,
event_type="SHOT_TASK_SET_CREATED",
project_id=task_set.id,
user_id=task_set.user_id,
message="创建拆镜总任务集",
detail={
"task_set_id": task_set.id,
"title": task_set.title,
"video_url": task_set.video_url,
"video_path": task_set.video_path,
"video_duration_seconds": task_set.video_duration_seconds,
"idempotency_key": task_set.idempotency_key,
},
)
return task_set
def _user_name_filter_subquery(value: str):
"""后台按用户名筛选时使用的子查询。
只在传入 user_name 时查询 users 表;keyword 不再关联 users,避免后台关键词搜索扩大查询范围。
"""
like = f"%{value.strip()}%"
return select(User.id).where(User.username.ilike(like))
async def _user_name_map_by_ids(db: AsyncSession, user_ids: set[str]) -> dict[str, str | None]:
"""一次性查询当前页涉及的用户,避免列表逐条查询用户表。"""
if not user_ids:
return {}
result = await db.execute(select(User.id, User.username).where(User.id.in_(list(user_ids))))
return {user_id: username for user_id, username in result.all()}
async def list_task_sets(
db: AsyncSession,
*,
current_user: User,
status: str | None = None,
analysis_status: str | None = None,
split_status: str | None = None,
keyword: str | None = None,
user_id: str | None = None,
user_name: str | None = None,
created_start: datetime | None = None,
created_end: datetime | None = None,
page: int = 1,
page_size: int = 20,
) -> ShotTaskSetListOut:
"""
拆镜总任务集列表查询。
性能策略:
1. 主列表不 join User,避免 count/list 复杂化。
2. 管理员按 user_name 搜索时使用 IN (SELECT users.id ...) 子查询;keyword 不查询 users 表。
3. 当前页数据取出后,再按 user_id 去重批量查 user_name,用于后台渲染。
"""
query = select(ShotReplicateTaskSet).where(ShotReplicateTaskSet.deleted_at.is_(None))
if not current_user.is_admin:
query = query.where(ShotReplicateTaskSet.user_id == current_user.id)
else:
if user_id and user_id.strip():
query = query.where(ShotReplicateTaskSet.user_id == user_id.strip())
if user_name and user_name.strip():
query = query.where(ShotReplicateTaskSet.user_id.in_(_user_name_filter_subquery(user_name)))
if created_start:
query = query.where(ShotReplicateTaskSet.created_at >= created_start)
if created_end:
query = query.where(ShotReplicateTaskSet.created_at <= created_end)
if status:
query = query.where(ShotReplicateTaskSet.status == status)
if analysis_status:
query = query.where(ShotReplicateTaskSet.analysis_status == analysis_status)
if split_status:
query = query.where(ShotReplicateTaskSet.split_status == split_status)
if keyword and keyword.strip():
like = f"%{keyword.strip()}%"
conditions = [
ShotReplicateTaskSet.id.ilike(like),
ShotReplicateTaskSet.title.ilike(like),
ShotReplicateTaskSet.original_video_content.ilike(like),
ShotReplicateTaskSet.original_video_category.ilike(like),
ShotReplicateTaskSet.original_video_audience.ilike(like),
cast(ShotReplicateTaskSet.ai_suggestion_json, String).ilike(like),
]
query = query.where(or_(*conditions))
total_result = await db.execute(select(func.count()).select_from(query.subquery()))
total = int(total_result.scalar() or 0)
result = await db.execute(
query.order_by(ShotReplicateTaskSet.created_at.desc())
.offset((page - 1) * page_size)
.limit(page_size)
)
task_sets = list(result.scalars().unique().all())
user_name_map: dict[str, str | None] = {}
if current_user.is_admin:
user_ids = {task_set.user_id for task_set in task_sets if task_set.user_id}
user_name_map = await _user_name_map_by_ids(db, user_ids)
return ShotTaskSetListOut(
total=total,
page=page,
page_size=page_size,
items=[
_task_set_to_out(
task_set,
user_name_map.get(task_set.user_id) if current_user.is_admin and task_set.user_id else None,
)
for task_set in task_sets
],
)
async def task_set_detail(db: AsyncSession, *, current_user: User, task_set_id: str) -> ShotTaskSetDetailOut:
task_set = await get_task_set_for_user(db, task_set_id=task_set_id, user=current_user)
user_result = await db.execute(select(User.username).where(User.id == task_set.user_id).limit(1))
return _task_set_to_detail_out(task_set, user_result.scalar_one_or_none())
async def _next_segment_index(db: AsyncSession, task_set_id: str) -> int:
result = await db.execute(
select(func.max(ShotReplicateSegment.segment_index)).where(
ShotReplicateSegment.task_set_id == task_set_id,
ShotReplicateSegment.deleted_at.is_(None),
)
)
return int(result.scalar() or 0) + 1
async def refresh_task_set_split_summary(db: AsyncSession, task_set_id: str) -> None:
task_set_result = await db.execute(select(ShotReplicateTaskSet).where(ShotReplicateTaskSet.id == task_set_id).with_for_update().limit(1))
task_set = task_set_result.scalar_one_or_none()
if not task_set:
return
result = await db.execute(
select(ShotReplicateSegment).where(
ShotReplicateSegment.task_set_id == task_set_id,
ShotReplicateSegment.deleted_at.is_(None),
)
)
segments = list(result.scalars().all())
total = len(segments)
completed = len([s for s in segments if s.split_status == ShotSplitStatusEnum.COMPLETED.value])
failed = len([s for s in segments if s.split_status == ShotSplitStatusEnum.FAILED.value])
task_set.segment_count = total
task_set.completed_segment_count = completed
task_set.failed_segment_count = failed
if total <= 0:
task_set.split_status = ShotSplitStatusEnum.NONE.value
if task_set.analysis_status == ShotAnalysisStatusEnum.COMPLETED.value:
task_set.status = ShotTaskSetStatusEnum.ANALYSIS_COMPLETED.value
return
old_status = task_set.status
old_split_status = task_set.split_status
if completed == total:
task_set.split_status = ShotSplitStatusEnum.COMPLETED.value
task_set.status = ShotTaskSetStatusEnum.SPLIT_COMPLETED.value
elif failed == total:
task_set.split_status = ShotSplitStatusEnum.FAILED.value
task_set.status = ShotTaskSetStatusEnum.FAILED.value
elif failed > 0:
task_set.split_status = ShotSplitStatusEnum.FAILED.value
task_set.status = ShotTaskSetStatusEnum.PARTIAL_FAILED.value
else:
task_set.split_status = ShotSplitStatusEnum.PROCESSING.value
task_set.status = ShotTaskSetStatusEnum.SPLITTING.value
if old_status != task_set.status or old_split_status != task_set.split_status:
log_module_event_file(
module=MODULE,
event_type="SHOT_SPLIT_STATUS_CHANGED",
project_id=task_set.id,
user_id=task_set.user_id,
message="拆镜总任务集拆分状态变更",
detail={
"task_set_id": task_set.id,
"from_status": old_status,
"to_status": task_set.status,
"from_split_status": old_split_status,
"to_split_status": task_set.split_status,
"segment_count": total,
"completed_segment_count": completed,
"failed_segment_count": failed,
},
)
async def create_segments_by_ai(
db: AsyncSession,
*,
current_user: User,
task_set_id: str,
req: ShotSplitByAIRequest,
) -> ShotSplitByAIOut:
task_set = await get_task_set_for_user(db, task_set_id=task_set_id, user=current_user, for_update=True)
if task_set.analysis_status != ShotAnalysisStatusEnum.COMPLETED.value:
raise HTTPException(status_code=400, detail="原视频分析未完成,不能按 AI 建议拆镜")
suggestions = _normalize_suggestions(task_set.ai_suggestion_json)
if not suggestions:
raise HTTPException(status_code=400, detail="当前没有可用 AI 建议拆镜方案,请使用自定义拆镜")
if req.selected_indices:
selected_set = {int(x) for x in req.selected_indices}
suggestions = [item for item in suggestions if int(item["index"]) in selected_set]
if not suggestions:
raise HTTPException(status_code=400, detail="selected_indices 没有匹配到可用 AI 建议")
old_result = await db.execute(
select(ShotReplicateSegment).where(
ShotReplicateSegment.task_set_id == task_set.id,
ShotReplicateSegment.source_mode == ShotSegmentSourceModeEnum.AI_SUGGESTION.value,
ShotReplicateSegment.deleted_at.is_(None),
)
)
old_segments = list(old_result.scalars().all())
if old_segments and not req.replace_existing:
raise HTTPException(status_code=409, detail="已存在 AI 建议拆镜片段,如需重拆请传 replace_existing=true")
if old_segments and req.replace_existing:
now = _now()
for segment in old_segments:
segment.deleted_at = now
created: list[ShotReplicateSegment] = []
next_index = await _next_segment_index(db, task_set.id)
for item in suggestions:
start, end, duration = validate_split_range(
start_second=item["start_second"],
end_second=item["end_second"],
video_duration_seconds=task_set.video_duration_seconds,
)
segment = ShotReplicateSegment(
id=generate_id(),
task_set_id=task_set.id,
user_id=task_set.user_id,
segment_index=next_index,
source_mode=ShotSegmentSourceModeEnum.AI_SUGGESTION.value,
start_second=start,
end_second=end,
duration_seconds=duration,
time_node=build_time_node(start, end),
split_status=ShotSplitStatusEnum.PENDING.value,
analysis_status=ShotSegmentAnalysisStatusEnum.NOT_REQUIRED.value,
replicate_status=ShotSegmentReplicateStatusEnum.NOT_STARTED.value,
original_video_content=task_set.original_video_content,
original_video_category=task_set.original_video_category,
original_video_audience=task_set.original_video_audience,
segment_content=item.get("content"),
segment_category=item.get("category"),
segment_audience=item.get("audience"),
ai_suggestion_json=item.get("raw") or item,
split_enqueued_at=_now(),
split_celery_task_id=f"shot-split:{uuid.uuid4().hex}",
)
db.add(segment)
created.append(segment)
next_index += 1
task_set.status = ShotTaskSetStatusEnum.SPLITTING.value
task_set.split_status = ShotSplitStatusEnum.PROCESSING.value
await db.flush()
await refresh_task_set_split_summary(db, task_set.id)
await db.flush()
log_module_event_file(
module=MODULE,
event_type="SHOT_SPLIT_BY_AI_SUBMITTED",
project_id=task_set.id,
user_id=task_set.user_id,
message="按 AI 建议创建拆镜片段",
detail={
"task_set_id": task_set.id,
"selected_indices": req.selected_indices,
"replace_existing": req.replace_existing,
"created_segment_count": len(created),
"segment_ids": [segment.id for segment in created],
},
)
return ShotSplitByAIOut(
task_set_id=task_set.id,
status=task_set.status,
split_status=task_set.split_status,
created_segment_count=len(created),
segments=[_segment_to_out(segment) for segment in created],
)
async def create_custom_segment(
db: AsyncSession,
*,
current_user: User,
task_set_id: str,
req: ShotSplitCustomRequest,
) -> ShotSplitCustomOut:
task_set = await get_task_set_for_user(db, task_set_id=task_set_id, user=current_user, for_update=True)
start, end, duration = validate_split_range(
start_second=req.start_second,
end_second=req.end_second,
video_duration_seconds=task_set.video_duration_seconds,
)
next_index = await _next_segment_index(db, task_set.id)
segment = ShotReplicateSegment(
id=generate_id(),
task_set_id=task_set.id,
user_id=task_set.user_id,
segment_index=next_index,
source_mode=ShotSegmentSourceModeEnum.CUSTOM.value,
start_second=start,
end_second=end,
duration_seconds=duration,
time_node=build_time_node(start, end),
split_status=ShotSplitStatusEnum.PENDING.value,
analysis_status=ShotSegmentAnalysisStatusEnum.PENDING.value,
replicate_status=ShotSegmentReplicateStatusEnum.NOT_STARTED.value,
split_enqueued_at=_now(),
split_celery_task_id=f"shot-split:{uuid.uuid4().hex}",
)
db.add(segment)
task_set.status = ShotTaskSetStatusEnum.SPLITTING.value
task_set.split_status = ShotSplitStatusEnum.PROCESSING.value
await db.flush()
await refresh_task_set_split_summary(db, task_set.id)
await db.flush()
log_module_event_file(
module=MODULE,
event_type="SHOT_SPLIT_CUSTOM_SUBMITTED",
project_id=task_set.id,
step_id=segment.id,
user_id=task_set.user_id,
message="按用户自定义时间创建拆镜片段",
detail={
"task_set_id": task_set.id,
"segment_id": segment.id,
"start_second": start,
"end_second": end,
"duration_seconds": duration,
"time_node": segment.time_node,
},
)
return ShotSplitCustomOut(task_set_id=task_set.id, segment=_segment_to_out(segment))
async def enqueue_segment_split(segment_id: str, *, countdown: int | None = None, recover: bool = False) -> None:
if not celery_app:
return
from app.tasks.shot_replicate_tasks import split_one_segment
split_one_segment.apply_async(
args=[segment_id],
queue="gen_result_download",
countdown=countdown,
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER if recover else settings.DOWNLOAD_TASK_PRIORITY_NORMAL,
)
async def list_segments(
db: AsyncSession,
*,
current_user: User,
task_set_id: str,
source_mode: str | None = None,
split_status: str | None = None,
analysis_status: str | None = None,
replicate_status: str | None = None,
page: int = 1,
page_size: int = 20,
) -> ShotSegmentListOut:
await get_task_set_for_user(db, task_set_id=task_set_id, user=current_user)
query = (
select(ShotReplicateSegment, ModuleGenerationProject)
.outerjoin(ModuleGenerationProject, ModuleGenerationProject.id == ShotReplicateSegment.module_project_id)
.where(
ShotReplicateSegment.task_set_id == task_set_id,
ShotReplicateSegment.deleted_at.is_(None),
)
)
if not current_user.is_admin:
query = query.where(ShotReplicateSegment.user_id == current_user.id)
if source_mode:
query = query.where(ShotReplicateSegment.source_mode == source_mode)
if split_status:
query = query.where(ShotReplicateSegment.split_status == split_status)
if analysis_status:
query = query.where(ShotReplicateSegment.analysis_status == analysis_status)
if replicate_status:
query = query.where(ShotReplicateSegment.replicate_status == replicate_status)
total_result = await db.execute(select(func.count()).select_from(query.subquery()))
total = int(total_result.scalar() or 0)
rows = await db.execute(
query.order_by(ShotReplicateSegment.segment_index.asc())
.offset((page - 1) * page_size)
.limit(page_size)
)
return ShotSegmentListOut(
total=total,
page=page,
page_size=page_size,
items=[_segment_to_out(segment, project) for segment, project in rows.all()],
)
async def segment_detail(db: AsyncSession, *, current_user: User, segment_id: str) -> ShotSegmentDetailOut:
segment = await get_segment_for_user(db, segment_id=segment_id, user=current_user)
project: ModuleGenerationProject | None = None
if segment.module_project_id:
project_result = await db.execute(select(ModuleGenerationProject).where(ModuleGenerationProject.id == segment.module_project_id).limit(1))
project = project_result.scalar_one_or_none()
return _segment_to_detail_out(segment, project)