407 lines
14 KiB
Python
407 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from datetime import datetime
|
|
from typing import Any, TypedDict
|
|
|
|
from sqlalchemy import and_, func, or_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.generation_task import GenerationType
|
|
from app.enums.recent_generation import (
|
|
RECENT_GENERATION_ALL_MODULES,
|
|
RECENT_GENERATION_CHAT_TASK_MODULES,
|
|
RECENT_GENERATION_COMPLETED_STATUS,
|
|
RECENT_GENERATION_MODULE_TO_TASK_MODE,
|
|
RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE,
|
|
RecentGenerationModuleEnum,
|
|
RecentGenerationResourceTypeEnum,
|
|
)
|
|
from app.models.chat_generation_task import ChatGenerationTask
|
|
from app.models.generation_record import GenerationRecord
|
|
from app.models.module_generation_step import ModuleGenerationStep
|
|
from app.models.shot_replicate_segment import ShotReplicateSegment
|
|
from app.schemas.recent_generation import RecentGenerationGroupOut, RecentGenerationItemOut
|
|
from app.services.resource_signed_url_service import build_resource_signed_url
|
|
|
|
DEFAULT_RECENT_GENERATION_LIMIT = 5
|
|
MAX_RECENT_GENERATION_LIMIT = 100
|
|
|
|
|
|
class _StepLinkInfo(TypedDict):
|
|
module_project_id: str | None
|
|
module_step_id: str | None
|
|
module: str | None
|
|
|
|
|
|
class _ShotLinkInfo(TypedDict):
|
|
shot_task_set_id: str | None
|
|
shot_segment_id: str | None
|
|
|
|
|
|
def _normalize_limit(limit: int | None) -> int:
|
|
if limit is None:
|
|
return DEFAULT_RECENT_GENERATION_LIMIT
|
|
return min(max(int(limit), 1), MAX_RECENT_GENERATION_LIMIT)
|
|
|
|
|
|
def _normalize_modules(
|
|
modules: Iterable[RecentGenerationModuleEnum] | None,
|
|
) -> list[RecentGenerationModuleEnum]:
|
|
if not modules:
|
|
return list(RECENT_GENERATION_ALL_MODULES)
|
|
|
|
normalized: list[RecentGenerationModuleEnum] = []
|
|
seen: set[RecentGenerationModuleEnum] = set()
|
|
for module in modules:
|
|
module_enum = RecentGenerationModuleEnum(module)
|
|
if module_enum not in seen:
|
|
normalized.append(module_enum)
|
|
seen.add(module_enum)
|
|
return normalized
|
|
|
|
|
|
def _has_url(column) -> Any:
|
|
return and_(column.is_not(None), column != "")
|
|
|
|
|
|
def _detect_resource_type(
|
|
gen_type: str | None,
|
|
image_url: str | None,
|
|
video_url: str | None,
|
|
) -> RecentGenerationResourceTypeEnum:
|
|
gen_type_value = (gen_type or "").strip().lower()
|
|
|
|
if gen_type_value == GenerationType.IMAGE.value and image_url:
|
|
return RecentGenerationResourceTypeEnum.IMAGE
|
|
if gen_type_value == GenerationType.VIDEO.value and video_url:
|
|
return RecentGenerationResourceTypeEnum.VIDEO
|
|
if video_url:
|
|
return RecentGenerationResourceTypeEnum.VIDEO
|
|
return RecentGenerationResourceTypeEnum.IMAGE
|
|
|
|
|
|
def _sign_url(url: str | None) -> str | None:
|
|
if not url:
|
|
return None
|
|
return build_resource_signed_url(url)
|
|
|
|
|
|
def _build_item(
|
|
*,
|
|
generation_id: str,
|
|
module: RecentGenerationModuleEnum,
|
|
gen_type: str | None,
|
|
image_url: str | None,
|
|
video_url: str | None,
|
|
video_cover_url: str | None,
|
|
generated_time: datetime | None,
|
|
step_info: _StepLinkInfo | None = None,
|
|
shot_info: _ShotLinkInfo | None = None,
|
|
) -> RecentGenerationItemOut:
|
|
resource_type = _detect_resource_type(
|
|
gen_type=gen_type,
|
|
image_url=image_url,
|
|
video_url=video_url,
|
|
)
|
|
|
|
raw_result_url = video_url if resource_type == RecentGenerationResourceTypeEnum.VIDEO else image_url
|
|
raw_cover_url = video_cover_url if resource_type == RecentGenerationResourceTypeEnum.VIDEO else None
|
|
|
|
return RecentGenerationItemOut(
|
|
generated_time=generated_time,
|
|
result_url=_sign_url(raw_result_url),
|
|
cover_url=_sign_url(raw_cover_url),
|
|
module=module,
|
|
shot_task_set_id=shot_info["shot_task_set_id"] if shot_info else None,
|
|
shot_segment_id=shot_info["shot_segment_id"] if shot_info else None,
|
|
module_project_id=step_info["module_project_id"] if step_info else None,
|
|
module_step_id=step_info["module_step_id"] if step_info else None,
|
|
generation_id=generation_id,
|
|
resource_type=resource_type,
|
|
)
|
|
|
|
|
|
async def _list_project_recent_items(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
limit: int,
|
|
) -> list[RecentGenerationItemOut]:
|
|
generated_time_expr = func.coalesce(
|
|
GenerationRecord.generated_at,
|
|
GenerationRecord.updated_at,
|
|
GenerationRecord.created_at,
|
|
).label("generated_time")
|
|
|
|
stmt = (
|
|
select(
|
|
GenerationRecord.id.label("generation_id"),
|
|
GenerationRecord.gen_type.label("gen_type"),
|
|
GenerationRecord.image_url.label("image_url"),
|
|
GenerationRecord.video_url.label("video_url"),
|
|
GenerationRecord.video_cover_url.label("video_cover_url"),
|
|
generated_time_expr,
|
|
)
|
|
.where(
|
|
GenerationRecord.user_id == user_id,
|
|
GenerationRecord.deleted_at.is_(None),
|
|
GenerationRecord.status == RECENT_GENERATION_COMPLETED_STATUS,
|
|
or_(_has_url(GenerationRecord.image_url), _has_url(GenerationRecord.video_url)),
|
|
)
|
|
.order_by(generated_time_expr.desc(), GenerationRecord.created_at.desc())
|
|
.limit(limit)
|
|
)
|
|
|
|
rows = (await db.execute(stmt)).mappings().all()
|
|
return [
|
|
_build_item(
|
|
generation_id=row["generation_id"],
|
|
module=RecentGenerationModuleEnum.PROJECT,
|
|
gen_type=row["gen_type"],
|
|
image_url=row["image_url"],
|
|
video_url=row["video_url"],
|
|
video_cover_url=row["video_cover_url"],
|
|
generated_time=row["generated_time"],
|
|
)
|
|
for row in rows
|
|
]
|
|
|
|
|
|
async def _list_chat_task_recent_rows(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
modules: list[RecentGenerationModuleEnum],
|
|
limit: int,
|
|
) -> list[dict[str, Any]]:
|
|
task_mode_values = [
|
|
RECENT_GENERATION_MODULE_TO_TASK_MODE[module].value
|
|
for module in modules
|
|
if module in RECENT_GENERATION_CHAT_TASK_MODULES
|
|
]
|
|
if not task_mode_values:
|
|
return []
|
|
|
|
generated_time_expr = func.coalesce(
|
|
ChatGenerationTask.generated_at,
|
|
ChatGenerationTask.updated_at,
|
|
ChatGenerationTask.created_at,
|
|
)
|
|
|
|
ranked_subquery = (
|
|
select(
|
|
ChatGenerationTask.id.label("generation_id"),
|
|
ChatGenerationTask.generation_mode.label("generation_mode"),
|
|
ChatGenerationTask.gen_type.label("gen_type"),
|
|
ChatGenerationTask.image_url.label("image_url"),
|
|
ChatGenerationTask.video_url.label("video_url"),
|
|
ChatGenerationTask.video_cover_url.label("video_cover_url"),
|
|
generated_time_expr.label("generated_time"),
|
|
func.row_number()
|
|
.over(
|
|
partition_by=ChatGenerationTask.generation_mode,
|
|
order_by=(generated_time_expr.desc(), ChatGenerationTask.created_at.desc()),
|
|
)
|
|
.label("row_num"),
|
|
)
|
|
.where(
|
|
ChatGenerationTask.user_id == user_id,
|
|
ChatGenerationTask.deleted_at.is_(None),
|
|
ChatGenerationTask.status == RECENT_GENERATION_COMPLETED_STATUS,
|
|
ChatGenerationTask.generation_mode.in_(task_mode_values),
|
|
or_(_has_url(ChatGenerationTask.image_url), _has_url(ChatGenerationTask.video_url)),
|
|
)
|
|
.subquery()
|
|
)
|
|
|
|
stmt = (
|
|
select(ranked_subquery)
|
|
.where(ranked_subquery.c.row_num <= limit)
|
|
.order_by(ranked_subquery.c.generation_mode.asc(), ranked_subquery.c.generated_time.desc())
|
|
)
|
|
|
|
return [dict(row) for row in (await db.execute(stmt)).mappings().all()]
|
|
|
|
|
|
async def _load_step_link_map(
|
|
db: AsyncSession,
|
|
*,
|
|
chat_task_ids: list[str],
|
|
) -> dict[str, _StepLinkInfo]:
|
|
if not chat_task_ids:
|
|
return {}
|
|
|
|
stmt = (
|
|
select(
|
|
ModuleGenerationStep.chat_task_id.label("chat_task_id"),
|
|
ModuleGenerationStep.id.label("module_step_id"),
|
|
ModuleGenerationStep.project_id.label("module_project_id"),
|
|
ModuleGenerationStep.module.label("module"),
|
|
ModuleGenerationStep.is_current.label("is_current"),
|
|
ModuleGenerationStep.updated_at.label("updated_at"),
|
|
)
|
|
.where(
|
|
ModuleGenerationStep.deleted_at.is_(None),
|
|
ModuleGenerationStep.chat_task_id.in_(chat_task_ids),
|
|
ModuleGenerationStep.module.in_(
|
|
[
|
|
RecentGenerationModuleEnum.HOT_OPENING_REPLICATE.value,
|
|
RecentGenerationModuleEnum.SHOT_REPLICATE.value,
|
|
]
|
|
),
|
|
)
|
|
.order_by(
|
|
ModuleGenerationStep.chat_task_id.asc(),
|
|
ModuleGenerationStep.is_current.desc(),
|
|
ModuleGenerationStep.updated_at.desc(),
|
|
)
|
|
)
|
|
|
|
link_map: dict[str, _StepLinkInfo] = {}
|
|
rows = (await db.execute(stmt)).mappings().all()
|
|
for row in rows:
|
|
chat_task_id = row["chat_task_id"]
|
|
if not chat_task_id or chat_task_id in link_map:
|
|
continue
|
|
link_map[chat_task_id] = {
|
|
"module_project_id": row["module_project_id"],
|
|
"module_step_id": row["module_step_id"],
|
|
"module": row["module"],
|
|
}
|
|
return link_map
|
|
|
|
|
|
async def _load_shot_link_map(
|
|
db: AsyncSession,
|
|
*,
|
|
module_project_ids: list[str],
|
|
) -> dict[str, _ShotLinkInfo]:
|
|
if not module_project_ids:
|
|
return {}
|
|
|
|
stmt = (
|
|
select(
|
|
ShotReplicateSegment.module_project_id.label("module_project_id"),
|
|
ShotReplicateSegment.id.label("shot_segment_id"),
|
|
ShotReplicateSegment.task_set_id.label("shot_task_set_id"),
|
|
)
|
|
.where(
|
|
ShotReplicateSegment.deleted_at.is_(None),
|
|
ShotReplicateSegment.module_project_id.in_(module_project_ids),
|
|
)
|
|
.order_by(ShotReplicateSegment.updated_at.desc())
|
|
)
|
|
|
|
link_map: dict[str, _ShotLinkInfo] = {}
|
|
rows = (await db.execute(stmt)).mappings().all()
|
|
for row in rows:
|
|
module_project_id = row["module_project_id"]
|
|
if not module_project_id or module_project_id in link_map:
|
|
continue
|
|
link_map[module_project_id] = {
|
|
"shot_task_set_id": row["shot_task_set_id"],
|
|
"shot_segment_id": row["shot_segment_id"],
|
|
}
|
|
return link_map
|
|
|
|
|
|
async def _build_chat_task_group_items(
|
|
db: AsyncSession,
|
|
*,
|
|
rows: list[dict[str, Any]],
|
|
) -> dict[RecentGenerationModuleEnum, list[RecentGenerationItemOut]]:
|
|
grouped: dict[RecentGenerationModuleEnum, list[RecentGenerationItemOut]] = {
|
|
RecentGenerationModuleEnum.CHAT_AI: [],
|
|
RecentGenerationModuleEnum.HOT_OPENING_REPLICATE: [],
|
|
RecentGenerationModuleEnum.SHOT_REPLICATE: [],
|
|
}
|
|
if not rows:
|
|
return grouped
|
|
|
|
module_task_rows: list[dict[str, Any]] = []
|
|
module_chat_task_ids: list[str] = []
|
|
for row in rows:
|
|
module = RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE.get(row["generation_mode"])
|
|
if module in (
|
|
RecentGenerationModuleEnum.HOT_OPENING_REPLICATE,
|
|
RecentGenerationModuleEnum.SHOT_REPLICATE,
|
|
):
|
|
module_task_rows.append(row)
|
|
module_chat_task_ids.append(row["generation_id"])
|
|
|
|
step_link_map = await _load_step_link_map(db, chat_task_ids=module_chat_task_ids)
|
|
|
|
shot_project_ids = [
|
|
step_info["module_project_id"]
|
|
for row in module_task_rows
|
|
if (step_info := step_link_map.get(row["generation_id"]))
|
|
and step_info.get("module") == RecentGenerationModuleEnum.SHOT_REPLICATE.value
|
|
and step_info.get("module_project_id")
|
|
]
|
|
shot_link_map = await _load_shot_link_map(db, module_project_ids=shot_project_ids)
|
|
|
|
for row in rows:
|
|
module = RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE.get(row["generation_mode"])
|
|
if not module:
|
|
continue
|
|
|
|
step_info = step_link_map.get(row["generation_id"])
|
|
shot_info = None
|
|
if module == RecentGenerationModuleEnum.SHOT_REPLICATE and step_info:
|
|
module_project_id = step_info.get("module_project_id")
|
|
if module_project_id:
|
|
shot_info = shot_link_map.get(module_project_id)
|
|
|
|
grouped[module].append(
|
|
_build_item(
|
|
generation_id=row["generation_id"],
|
|
module=module,
|
|
gen_type=row["gen_type"],
|
|
image_url=row["image_url"],
|
|
video_url=row["video_url"],
|
|
video_cover_url=row["video_cover_url"],
|
|
generated_time=row["generated_time"],
|
|
step_info=step_info,
|
|
shot_info=shot_info,
|
|
)
|
|
)
|
|
|
|
return grouped
|
|
|
|
|
|
async def list_recent_generations(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
modules: Iterable[RecentGenerationModuleEnum] | None = None,
|
|
limit: int | None = None,
|
|
) -> RecentGenerationGroupOut:
|
|
"""获取当前用户各模块最近生成成功的图片/视频记录。"""
|
|
|
|
normalized_limit = _normalize_limit(limit)
|
|
normalized_modules = _normalize_modules(modules)
|
|
|
|
response = RecentGenerationGroupOut()
|
|
|
|
if RecentGenerationModuleEnum.PROJECT in normalized_modules:
|
|
response.project = await _list_project_recent_items(
|
|
db,
|
|
user_id=user_id,
|
|
limit=normalized_limit,
|
|
)
|
|
|
|
chat_modules = [module for module in normalized_modules if module in RECENT_GENERATION_CHAT_TASK_MODULES]
|
|
if chat_modules:
|
|
chat_rows = await _list_chat_task_recent_rows(
|
|
db,
|
|
user_id=user_id,
|
|
modules=chat_modules,
|
|
limit=normalized_limit,
|
|
)
|
|
chat_grouped = await _build_chat_task_group_items(db, rows=chat_rows)
|
|
response.chat_ai = chat_grouped[RecentGenerationModuleEnum.CHAT_AI]
|
|
response.hot_opening_replicate = chat_grouped[RecentGenerationModuleEnum.HOT_OPENING_REPLICATE]
|
|
response.shot_replicate = chat_grouped[RecentGenerationModuleEnum.SHOT_REPLICATE]
|
|
|
|
return response |