新增素材列表前测提交

video-gen-api/app/api/v1/resources_material.py

@@ -2,10 +2,15 @@ from typing import Any, Optional
 
 from fastapi import APIRouter, Depends, Query
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select
+from app.models.resources_material import ResourcesMaterial
+from app.models.pre_test_template import PreTestTemplate
 
-from app.dependencies import get_db
+from app.dependencies import get_db, get_current_user
 from app.schemas.resources_material import ResourcesMaterialListResponse
 from app.services.resources_material_service import get_resources_material_list
+from app.services.pre_test_queue import pre_test_queue
+from app.utils.id_gen import generate_id
 
 router = APIRouter(prefix="/resources-material", tags=["resources-material"])
 
@@ -42,4 +47,98 @@ async def get_resources_material_list_api(
         "message": "查询成功",
         "data": items,
         "total": total,
-    }
+    }
+
+#如果素材上传的时候没有指定前测，现在可以对已经上传好的素材进行前测
+@router.post(
+    "/pre-commit",
+    summary="素材列表提交前测",
+    description="通过素材列表提交未前测的素材，异步处理，立即返回任务ID，结果稍后通过列表查询",
+)
+async def pre_test_material(
+    resources_material_ids: list[str] = Query(..., description="资源素材表id"),
+    pre_test_template_id: str = Query(..., description="前测模板id"),
+    db: AsyncSession = Depends(get_db),
+    current_user: dict = Depends(get_current_user),
+) -> Any | dict:
+
+    try:
+        invalid_ids = []
+        grouped_videos = {}
+        
+        #检查前测模板是否有效
+        pre_test_template = await db.execute(
+            select(PreTestTemplate)
+            .where(
+                PreTestTemplate.id == pre_test_template_id,
+                PreTestTemplate.user_id == current_user.id,
+                PreTestTemplate.deleted_at.is_(None),
+            )
+        )
+        pre_test_template = pre_test_template.scalar_one_or_none()
+        if not pre_test_template:
+            return {"code": 1, "message": f"前测模板{pre_test_template_id}不存在或不属于当前用户"}
+        
+        for resource_material_id in resources_material_ids:
+            resource_material = await db.execute(
+                select(ResourcesMaterial)
+                .where(
+                    ResourcesMaterial.id == resource_material_id,
+                    ResourcesMaterial.user_id == current_user.id,
+                    ResourcesMaterial.deleted_at.is_(None),
+                )
+            )
+            resource_material = resource_material.scalar_one_or_none()
+            
+            if not resource_material:
+                invalid_ids.append(f"{resource_material_id}: 资源不存在或不属于当前用户")
+                continue
+            
+            if resource_material.status is not None:
+                invalid_ids.append(f"{resource_material_id}: 该资源已进行过前测")
+                continue
+
+            if not resource_material.upload_id:
+                invalid_ids.append(f"{resource_material_id}: 该资源未上传到平台")
+                continue
+            
+            if resource_material.resource_type != "video":
+                invalid_ids.append(f"{resource_material_id}: 前测仅支持视频类型")
+                continue
+            
+            key = f"{resource_material.oauth_id}|{resource_material.advertiser_id}"
+            if key not in grouped_videos:
+                grouped_videos[key] = []
+            grouped_videos[key].append({
+                "id": resource_material.id,
+                "upload_id": resource_material.upload_id,
+            })
+        
+        if invalid_ids:
+            logger.error(f"前测素材提交失败，无效素材ID: {'; '.join(invalid_ids)}")
+            return {"code": 1, "message": "; ".join(invalid_ids)}
+        
+        if not grouped_videos:
+            return {"code": 1, "message": "没有有效的视频资源"}
+        
+        task_id = generate_id()
+        
+        await pre_test_queue.enqueue({
+            "task_id": task_id,
+            "grouped_videos": grouped_videos,
+            "pre_test_template_id": pre_test_template_id,
+        })
+        
+        return {
+            "code": 0,
+            "message": "任务已提交，正在处理中",
+            "data": {
+                "task_id": task_id,
+                "total_groups": len(grouped_videos),
+            },
+        }
+        
+    except ValueError as e:
+        return {"code": 1, "message": str(e)}
+    except Exception as e:
+        return {"code": 1, "message": str(e)}

video-gen-api/app/main.py

@@ -80,6 +80,10 @@ async def lifespan(app: FastAPI):
     from app.tasks.pre_test_result_task import poll_pre_test_results
     pre_test_poll_task = asyncio.create_task(poll_pre_test_results())
     
+    # 启动前测任务队列
+    from app.services.pre_test_queue import pre_test_queue
+    pre_test_queue_task = asyncio.create_task(pre_test_queue.run())
+    
     # 启动时立即同步一次未支付订单
     asyncio.create_task(asyncio.sleep(5))  # 等待5秒后再同步，让系统完全启动
     async def startup_sync():
@@ -104,6 +108,8 @@ async def lifespan(app: FastAPI):
     await queue_task
     upload_queue.stop()
     await upload_queue_task
+    pre_test_queue.stop()
+    await pre_test_queue_task
     material_consumption_queue.stop()
     await consumption_queue_task
     consumption_schedule_task.cancel()

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

@@ -0,0 +1,246 @@
+import asyncio
+from sqlalchemy import select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.models.base import async_session
+from app.models.resources_material import ResourcesMaterial
+from app.models.user_oauth import UserOAuth
+from app.models.pre_test_template import PreTestTemplate
+from app.utils.logger import get_logger
+from app.utils.douyinApi import DouyinApi
+import json
+
+logger = get_logger("pre_test_queue", "pre_test_queue")
+douyin_api = DouyinApi()
+
+
+class PreTestQueue:
+    def __init__(self):
+        self.queue: asyncio.Queue[dict] = asyncio.Queue()
+        self.running = False
+
+    async def enqueue(self, task_data: dict):
+        """Add a pre-test task to the queue."""
+        await self.queue.put(task_data)
+        logger.info(f"Enqueued pre-test task: {task_data.get('task_id')}")
+
+    async def run(self):
+        """Main processing loop."""
+        self.running = True
+        logger.info("Pre-test queue started")
+
+        while self.running:
+            try:
+                task_data = await asyncio.wait_for(self.queue.get(), timeout=5.0)
+            except asyncio.TimeoutError:
+                continue
+
+            try:
+                await self._process(task_data)
+            except Exception as e:
+                logger.error(f"Error processing pre-test task: {e}", exc_info=True)
+            finally:
+                self.queue.task_done()
+
+        logger.info("Pre-test queue stopped")
+
+    async def _process(self, task_data: dict):
+        """Process a single pre-test task."""
+        task_id = task_data.get("task_id")
+        grouped_videos = task_data.get("grouped_videos", {})
+        pre_test_template_id = task_data.get("pre_test_template_id")
+
+        logger.info(f"Processing pre-test task: {task_id}")
+
+        for key, videos in grouped_videos.items():
+            oauth_id, advertiser_id = key.split("|")
+            video_ids = [v["upload_id"] for v in videos]
+            
+            logger.info(f"Processing {len(video_ids)} videos for oauth_id={oauth_id}, advertiser_id={advertiser_id}")
+            
+            try:
+                async with async_session() as db:
+                    await _pre_test_material_batch(
+                        oauth_id=oauth_id,
+                        advertiser_id=advertiser_id,
+                        video_ids=video_ids,
+                        pre_test_template_id=pre_test_template_id,
+                        db=db,
+                        resource_material_ids=[v["id"] for v in videos],
+                    )
+            except Exception as e:
+                logger.error(f"Failed to process pre-test for oauth_id={oauth_id}, advertiser_id={advertiser_id}: {e}", exc_info=True)
+
+    def stop(self):
+        """Stop the queue."""
+        self.running = False
+
+
+async def _pre_test_material_batch(
+    oauth_id: str,
+    advertiser_id: str,
+    video_ids: list[str],
+    pre_test_template_id: str,
+    db: AsyncSession,
+    resource_material_ids: list[str],
+) -> any:
+    oauth = await db.execute(
+        select(UserOAuth).where(
+            UserOAuth.id == oauth_id,
+            UserOAuth.deleted_at.is_(None),
+        )
+    )
+    oauth = oauth.scalar_one_or_none()
+    if not oauth:
+        note = "授权记录不存在"
+        for video_id, resource_material_id in zip(video_ids, resource_material_ids):
+            await _update_material_pre_test_status(
+                db, resource_material_id,
+                task_id=None,
+                status="FAILED",
+                note=note,
+                pre_result=None,
+                pre_test_template_id=pre_test_template_id,
+            )
+        await db.commit()
+        return {"code": -1, "message": note, "data": {}}
+
+    pre_test_template = await db.execute(
+        select(PreTestTemplate).where(
+            PreTestTemplate.id == pre_test_template_id,
+            PreTestTemplate.deleted_at.is_(None),
+            PreTestTemplate.user_id == oauth.user_id,
+        )
+    )
+    pre_test_template = pre_test_template.scalar_one_or_none()
+    if not pre_test_template:
+        note = f"前测模板 {pre_test_template_id} 不存在"
+        for video_id, resource_material_id in zip(video_ids, resource_material_ids):
+            await _update_material_pre_test_status(
+                db, resource_material_id,
+                task_id=None,
+                status="FAILED",
+                note=note,
+                pre_result=None,
+                pre_test_template_id=pre_test_template_id,
+            )
+        await db.commit()
+        return {"code": -1, "message": note, "data": {}}
+    
+    diagnose_config = {}
+    if pre_test_template.platform:
+        diagnose_config["platform"] = pre_test_template.platform
+    if pre_test_template.external_action:
+        diagnose_config["external_action"] = pre_test_template.external_action
+    if pre_test_template.cpa_bid:
+        diagnose_config["cpa_bid"] = pre_test_template.cpa_bid
+    if pre_test_template.audience_gender:
+        diagnose_config["audience_gender"] = pre_test_template.audience_gender
+    if pre_test_template.audience_age:
+        diagnose_config["audience_age"] = json.loads(pre_test_template.audience_age)
+    if pre_test_template.audience_region:
+        diagnose_config["audience_region"] = json.loads(pre_test_template.audience_region)
+    if pre_test_template.audience_network:
+        diagnose_config["audience_network"] = json.loads(pre_test_template.audience_network)
+    if pre_test_template.cus_name:
+        diagnose_config["cus_name"] = pre_test_template.cus_name
+    if pre_test_template.pricing_type:
+        diagnose_config["pricing_type"] = pre_test_template.pricing_type
+    if pre_test_template.cost_cap:
+        diagnose_config["cost_cap"] = pre_test_template.cost_cap
+    if pre_test_template.target_cost:
+        diagnose_config["target_cost"] = pre_test_template.target_cost
+    if pre_test_template.nobid:
+        diagnose_config["nobid"] = pre_test_template.nobid
+    if pre_test_template.cpc_bid:
+        diagnose_config["cpc_bid"] = pre_test_template.cpc_bid
+    if pre_test_template.budget:
+        diagnose_config["budget"] = pre_test_template.budget
+
+    params = {
+        "advertiser_id": int(advertiser_id),
+        "video_ids": video_ids,
+        "diagnose_config": diagnose_config,
+    }
+    response = await douyin_api.pre_test_material(oauth_id, params)
+    
+    code = response.get("code", -1)
+    
+    if code != 0:
+        #记录日志
+        logger.error(f"前测提交失败, oauth_id={oauth_id}, advertiser_id={advertiser_id}: {json.dumps(response)}")
+        note = response.get("message", "未知错误")
+        for video_id, resource_material_id in zip(video_ids, resource_material_ids):
+            await _update_material_pre_test_status(
+                db, resource_material_id,
+                task_id=None,
+                status="FAILED",
+                note=note,
+                pre_result=None,
+                pre_test_template_id=pre_test_template_id,
+            )
+        await db.commit()
+        return response
+    
+    data = response.get("data", {})
+    task_ids = data.get("task_ids", [])
+    fail_video_ids = data.get("fail_video_ids", {})
+    
+    success_count = 0
+    for i, (video_id, resource_material_id) in enumerate(zip(video_ids, resource_material_ids)):
+        if video_id in fail_video_ids:
+            fail_info = fail_video_ids[video_id]
+            err_code = fail_info.get("err_code", "")
+            err_message = fail_info.get("err_message", "未知错误")
+            note = f"失败[{err_code}]: {err_message}"
+            
+            await _update_material_pre_test_status(
+                db, resource_material_id,
+                task_id=None,
+                status="FAILED",
+                note=note,
+                pre_result=None,
+                pre_test_template_id=pre_test_template_id,
+            )
+        else:
+            task_id = str(task_ids[success_count]) if success_count < len(task_ids) else None
+            
+            await _update_material_pre_test_status(
+                db, resource_material_id,
+                task_id=task_id,
+                status="PENDING",
+                note="",
+                pre_result=None,
+                pre_test_template_id=pre_test_template_id,
+            )
+            success_count += 1
+    
+    await db.commit()
+    
+    return response
+
+
+async def _update_material_pre_test_status(
+    db: AsyncSession,
+    resource_material_id: str,
+    task_id: str | None,
+    status: str,
+    note: str,
+    pre_result: str | None,
+    pre_test_template_id: str | None,
+):
+    await db.execute(
+        update(ResourcesMaterial).where(
+            ResourcesMaterial.id == resource_material_id,
+            ResourcesMaterial.deleted_at.is_(None),
+        ).values(
+            task_id=task_id,
+            status=status,
+            note=note,
+            pre_result=pre_result,
+            pre_test_template_id=pre_test_template_id,
+        )
+    )
+
+
+pre_test_queue = PreTestQueue()