AI创作批量生成任务 main V1 init

video-gen-api/app/services/recent_generation_service.py

@@ -4,7 +4,7 @@ from collections.abc import Iterable
 from datetime import datetime
 from typing import Any, TypedDict
 
-from sqlalchemy import and_, func, or_, select
+from sqlalchemy import and_, case, func, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.enums.generation_task import GenerationType
@@ -12,7 +12,7 @@ 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_MODULE_TO_TASK_MODES,
     RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE,
     RecentGenerationModuleEnum,
     RecentGenerationResourceTypeEnum,
@@ -175,11 +175,12 @@ async def _list_chat_task_recent_rows(
     modules: list[RecentGenerationModuleEnum],
     limit: int,
 ) -> list[dict[str, Any]]:
-    task_mode_values = [
-        RECENT_GENERATION_MODULE_TO_TASK_MODE[module].value
+    task_mode_values = list(dict.fromkeys(
+        task_mode.value
         for module in modules
         if module in RECENT_GENERATION_CHAT_TASK_MODULES
-    ]
+        for task_mode in RECENT_GENERATION_MODULE_TO_TASK_MODES[module]
+    ))
     if not task_mode_values:
         return []
 
@@ -189,10 +190,19 @@ async def _list_chat_task_recent_rows(
         ChatGenerationTask.created_at,
     )
 
+    module_partition_expr = case(
+        (
+            ChatGenerationTask.generation_mode.in_(["chatapi_async", "chatapi_child"]),
+            RecentGenerationModuleEnum.CHAT_AI.value,
+        ),
+        else_=ChatGenerationTask.generation_mode,
+    )
+
     ranked_subquery = (
         select(
             ChatGenerationTask.id.label("generation_id"),
             ChatGenerationTask.generation_mode.label("generation_mode"),
+            module_partition_expr.label("module_key"),
             ChatGenerationTask.gen_type.label("gen_type"),
             ChatGenerationTask.image_url.label("image_url"),
             ChatGenerationTask.video_url.label("video_url"),
@@ -200,7 +210,7 @@ async def _list_chat_task_recent_rows(
             generated_time_expr.label("generated_time"),
             func.row_number()
             .over(
-                partition_by=ChatGenerationTask.generation_mode,
+                partition_by=module_partition_expr,
                 order_by=(generated_time_expr.desc(), ChatGenerationTask.created_at.desc()),
             )
             .label("row_num"),
@@ -218,7 +228,7 @@ async def _list_chat_task_recent_rows(
     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())
+        .order_by(ranked_subquery.c.module_key.asc(), ranked_subquery.c.generated_time.desc())
     )
 
     return [dict(row) for row in (await db.execute(stmt)).mappings().all()]