ai-flow-python/app/service/task/impl/seedance_video_task.py

import logging
import time

from app.config import settings
from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util.http_util import logged_client
from app.util.retry_util import retry_on_failure

logger = logging.getLogger(__name__)


class SeedanceVideoTask(BaseTask):
    def get_type(self) -> str:
        return "seedance_video"

    def _headers(self) -> dict[str, str]:
        return {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {settings.seed_api_key}",
        }

    def _submit(self, url: str, inputs: dict) -> dict:
        with logged_client(timeout=60) as client:
            resp = client.post(url, headers=self._headers(), json=inputs)
            logger.info("[seedance_video] Submit response: status=%d", resp.status_code)
            if not resp.is_success:
                raise RuntimeError(f"HTTP {resp.status_code}: {resp.text}")
            data = resp.json()
        task_id = data.get("id")
        if not task_id:
            raise RuntimeError(f"No task ID in response: {data}")
        return data

    def execute(self, param: TaskParam) -> TaskResult:
        execute_url = settings.seedance_video_url
        t0 = time.time()

        logger.info("[seedance_video] Submitting task: POST %s", execute_url)
        try:
            data = retry_on_failure(
                lambda: self._submit(execute_url, param.inputs),
                task_name="seedance_video",
            )
            task_id = data.get("id")
            logger.info("[seedance_video] Task submitted, remote_id=%s", task_id)
        except Exception as e:
            logger.error("[seedance_video] Submit failed: %s", e)
            return TaskResult.failure(f"Failed to submit task: {e}")

        result = self._poll(task_id)
        if result is None:
            return TaskResult.failure("Failed to poll task status")

        if result.get("status") != "succeeded":
            return TaskResult.failure(f"Task failed with status: {result.get('status')}")

        videos: list[str] = []
        content = result.get("content", {})
        video_url = content.get("video_url")
        if video_url:
            videos.append(video_url)

        logger.info("[seedance_video] Completed: %d videos, total_elapsed=%.2fs",
                     len(videos), time.time() - t0)
        return TaskResult.success({"videos": videos})

    def _poll(self, task_id: str) -> dict | None:
        search_url = f"{settings.seedance_video_url}/{task_id}"
        max_attempts = 30
        poll_t0 = time.time()

        for attempt in range(1, max_attempts + 1):
            try:
                with logged_client(timeout=60) as client:
                    resp = client.get(search_url, headers=self._headers())
                    data = resp.json()

                status = data.get("status")
                logger.info("[seedance_video] Poll #%d/%d: remote_id=%s, status=%s, elapsed=%.1fs",
                             attempt, max_attempts, task_id, status, time.time() - poll_t0)

                if status == "succeeded":
                    return data
                if status in ("failed", "cancelled"):
                    logger.error("[seedance_video] Task %s terminated with status=%s", task_id, status)
                    return None

                if attempt < max_attempts:
                    time.sleep(2)
            except Exception as e:
                logger.error("[seedance_video] Poll #%d error for %s: %s", attempt, task_id, e)
                if attempt < max_attempts:
                    time.sleep(2)

        logger.warning("[seedance_video] Max polling attempts (%d) reached for %s", max_attempts, task_id)
        return None