增加前测任务获取前测结果

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

@@ -158,6 +158,11 @@ async def async_batch_upload_material(
         
         for advertiser_id in task.advertiser_ids:
             for resource_id in resource_ids_to_upload:
+                other_info = {}
+                if task.is_pre_test == "1":
+                    other_info["is_pre_test"] = task.is_pre_test
+                    other_info["pre_test_template"] = task.pre_test_template
+
                 task_id = generate_id()
                 upload_task = UploadTask(
                     id=task_id,
@@ -167,6 +172,7 @@ async def async_batch_upload_material(
                     resource_id=resource_id,
                     status=1,
                     note=None,
+                    other_info=json.dumps(other_info) if other_info else None,
                 )
                 
                 db.add(upload_task)

video-gen-api/app/main.py

@@ -76,6 +76,10 @@ async def lifespan(app: FastAPI):
     from app.tasks.material_consumption_task import schedule_daily_sync
     consumption_schedule_task = asyncio.create_task(schedule_daily_sync())
     
+    # 启动前测结果轮询任务（每分钟检查一次）
+    from app.tasks.pre_test_result_task import poll_pre_test_results
+    pre_test_poll_task = asyncio.create_task(poll_pre_test_results())
+    
     # 启动时立即同步一次未支付订单
     asyncio.create_task(asyncio.sleep(5))  # 等待5秒后再同步，让系统完全启动
     async def startup_sync():
@@ -103,6 +107,7 @@ async def lifespan(app: FastAPI):
     material_consumption_queue.stop()
     await consumption_queue_task
     consumption_schedule_task.cancel()
+    pre_test_poll_task.cancel()
     expiry_task.cancel()
     token_refresh_task.cancel()
     await close_database()

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

@@ -4,6 +4,7 @@ from datetime import datetime, timezone
 
 from sqlalchemy import select, update
 from sqlalchemy.orm import selectinload
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.models.base import async_session
 from app.models.upload_task import UploadTask
@@ -11,6 +12,7 @@ from app.models.generated_resource import GeneratedResource
 from app.models.user_oauth import UserOAuth
 from app.models.resources_material import ResourcesMaterial
 from app.models.user_oauth_account import UserOAuthAccount
+from app.models.pre_test_template import PreTestTemplate
 from app.utils.id_gen import generate_id
 from app.utils.douyinApi import DouyinApi
 
@@ -51,7 +53,7 @@ if not logger.handlers:
 
 douyin_api = DouyinApi()
 
-
+#上传素材队列，处理上传素材的任务
 class UploadQueue:
     def __init__(self):
         self.queue: asyncio.Queue[str] = asyncio.Queue()
@@ -124,7 +126,8 @@ class UploadQueue:
                 task.oauth_id,
                 task.advertiser_id,
                 task.resource_id,
-                db=None
+                db=None,
+                other_info=json.loads(task.other_info) if task.other_info else None,
             )
 
             async with async_session() as db:
@@ -215,7 +218,8 @@ async def _upload_single_material(
     oauth_id: str,
     advertiser_id: str,
     resource_id: str,
-    db=None
+    db=None,
+    other_info=None,
 ) -> dict:
     own_db = False
     if db is None:
@@ -269,9 +273,20 @@ async def _upload_single_material(
                 "error": "资源本地存储路径为空",
             }
 
-        return await _upload_to_juliang(
+        result = await _upload_to_juliang(
             oauth_id, storage_path, resource_type, advertiser_id, resource, db, user_id, file_name
         )
+
+        if result.get("success") and resource_type == "video" and other_info and "is_pre_test" in other_info and other_info["is_pre_test"] == "1" and "pre_test_template" in other_info:
+            await _pre_test_material(
+                oauth_id,
+                advertiser_id,
+                [result.get("upload_id")],
+                other_info["pre_test_template"],
+                db,
+            )
+            
+        return result
     finally:
         if own_db and db is not None:
             await db.close()
@@ -470,4 +485,172 @@ async def _upload_to_juliang(
         }
 
 
+async def _pre_test_material(
+    oauth_id: str,
+    advertiser_id: str,
+    video_ids: list[str],
+    pre_test_template_id: str,
+    db: AsyncSession,
+) -> 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 in video_ids:
+            await _update_material_pre_test_status(
+                db, oauth_id, advertiser_id, video_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 in video_ids:
+            await _update_material_pre_test_status(
+                db, oauth_id, advertiser_id, video_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:
+        note = response.get("message", "未知错误")
+        for video_id in video_ids:
+            await _update_material_pre_test_status(
+                db, oauth_id, advertiser_id, video_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 in enumerate(video_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, oauth_id, advertiser_id, video_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, oauth_id, advertiser_id, video_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,
+    oauth_id: str,
+    advertiser_id: str,
+    upload_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.oauth_id == oauth_id,
+            ResourcesMaterial.advertiser_id == advertiser_id,
+            ResourcesMaterial.upload_id == upload_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,
+        )
+    )
+
+
 upload_queue = UploadQueue()

video-gen-api/app/tasks/pre_test_result_task.py

@@ -0,0 +1,154 @@
+from datetime import datetime, timezone
+import asyncio
+import os
+import logging
+import json
+
+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.utils.douyinApi import DouyinApi
+
+LOG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "logs")
+os.makedirs(LOG_DIR, exist_ok=True)
+
+#前测结果和上传素材属于一种任务，放到一起日志里边
+logger = logging.getLogger("upload_queue")
+logger.setLevel(logging.INFO)
+
+
+class DailyRotatingFileHandler(logging.FileHandler):
+    def __init__(self, directory, encoding=None):
+        self.directory = directory
+        filename = self._get_log_filename()
+        super().__init__(filename, encoding=encoding)
+    
+    def _get_log_filename(self):
+        return os.path.join(self.directory, f"pre_test_result_task-{datetime.now(timezone.utc).strftime('%Y-%m-%d')}.log")
+    
+    def emit(self, record):
+        current_filename = self._get_log_filename()
+        if self.baseFilename != current_filename:
+            self.close()
+            self.baseFilename = current_filename
+            self.stream = self._open()
+        super().emit(record)
+
+
+if not logger.handlers:
+    handler = DailyRotatingFileHandler(LOG_DIR, encoding="utf-8")
+    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"))
+    logger.addHandler(handler)
+
+douyin_api = DouyinApi()
+
+#获取前测结果并更新数据库，计划任务，每2分钟执行一次
+async def poll_pre_test_results():
+    """每2分钟轮询前测结果并更新数据库"""
+    logger.info("Pre-test result polling task started")
+    
+    while True:
+        try:
+            await process_pending_pre_tests()
+            await asyncio.sleep(120)
+        
+        except asyncio.CancelledError:
+            logger.info("Pre-test result polling task cancelled")
+            break
+        except Exception as e:
+            logger.error(f"Error in poll_pre_test_results: {e}")
+            await asyncio.sleep(120)
+
+
+async def process_pending_pre_tests():
+    """处理所有待查询的前测任务"""
+    async with async_session() as db:
+        result = await db.execute(
+            select(ResourcesMaterial).where(
+                ResourcesMaterial.status == "PENDING",
+                ResourcesMaterial.task_id.is_not(None),
+                ResourcesMaterial.deleted_at.is_(None),
+            )
+        )
+        pending_materials = result.scalars().all()
+        
+        if not pending_materials:
+            return
+        
+        logger.info(f"Found {len(pending_materials)} pending pre-test tasks to process")
+        
+        for material in pending_materials:
+            try:
+                await update_single_pre_test_result(db, material)
+            except Exception as e:
+                logger.error(f"Error processing pre-test for material {material.id}: {e}")
+
+
+async def update_single_pre_test_result(db: AsyncSession, material: ResourcesMaterial):
+    """更新单个素材的前测结果"""
+    oauth = await db.execute(
+        select(UserOAuth).where(
+            UserOAuth.id == material.oauth_id,
+            UserOAuth.deleted_at.is_(None),
+        )
+    )
+    oauth = oauth.scalar_one_or_none()
+    
+    if not oauth:
+        logger.error(f"OAuth record not found for material {material.id}")
+        await db.execute(
+            update(ResourcesMaterial).where(ResourcesMaterial.id == material.id).values(
+                status="FAILED",
+                note="授权记录不存在",
+            )
+        )
+        await db.commit()
+        return
+    
+    params = {
+        "advertiser_id": int(material.advertiser_id),
+        "task_ids": json.dumps([int(material.task_id)]),
+    }
+    
+    try:
+        response = await douyin_api.get_material_pre_test_result(oauth.id, params)
+    except Exception as e:
+        logger.error(f"Failed to get pre-test result for task {material.task_id}: {e}")
+        return
+    
+    code = response.get("code", -1)
+    if code != 0:
+        logger.error(f"API error for task {material.task_id}: {response.get('message', 'Unknown error')}")
+        return
+    
+    data = response.get("data", {})
+    task_details = data.get("task_list", [])
+    
+    if not task_details:
+        return
+    
+    task_detail = task_details[0]
+    status = task_detail.get("status") 
+    pre_result = {
+        "video_id": task_detail.get("video_id") or None,
+        "advertiser_id": task_detail.get("advertiser_id") or None,
+        "material_id": task_detail.get("material_id") or None,
+        "is_ad_high_quality_material": task_detail.get("is_ad_high_quality_material") or None,
+        "is_ecp_high_quality_material": task_detail.get("is_ecp_high_quality_material") or None,
+        "is_inefficient_material": task_detail.get("is_inefficient_material") or None,
+        "is_first_publish_material": task_detail.get("is_first_publish_material") or None,
+        "not_ad_high_quality_reason": task_detail.get("not_ad_high_quality_reason") or None,
+        "not_ecp_high_quality_reason": task_detail.get("not_ecp_high_quality_reason") or None,
+        "is_local_high_quality_material": task_detail.get("is_local_high_quality_material") or None,
+    }
+    
+    await db.execute(
+        update(ResourcesMaterial).where(ResourcesMaterial.id == material.id).values(
+            status=status,
+            pre_result=json.dumps(pre_result, ensure_ascii=False),
+        )
+    )
+    await db.commit()

video-gen-api/app/utils/douyinApi.py

@@ -103,3 +103,15 @@ class DouyinApi:
             'GET',
             {'params': params or {}}
         )
+
+    #获取素材前测结果
+    async def get_material_pre_test_result(self, oauth_id: str, params: any) -> Dict[str, Any]:
+        if not oauth_id:
+            raise RuntimeError('OAuth ID is not set.')
+        url = "https://api.oceanengine.com/open_api/2/diagnosis_task/adv/get/"
+        return await self.request.request_with_token_with_context(
+            oauth_id,
+            url,
+            'GET',
+            {'params': params or {}}
+        )