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

import logging
import time

import jwt

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 KelingImage2VideoTask(BaseTask):
    def get_type(self) -> str:
        return "keling_image2video"

    @property
    def execute_url(self) -> str:
        return settings.keling_image2video_url

    @property
    def search_url(self) -> str:
        return settings.keling_image2video_url

    def _generate_jwt(self) -> str:
        now = time.time()
        payload = {
            "iss": settings.keling_app_key,
            "exp": int(now) + 1800,
            "nbf": int(now) - 5,
        }
        return jwt.encode(payload, settings.keling_app_secret, algorithm="HS256")

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

    def _submit(self, inputs: dict) -> dict:
        with logged_client(timeout=60) as client:
            resp = client.post(self.execute_url, headers=self._headers(), json=inputs)
            logger.info("[%s] Submit response: status=%d", self.get_type(), resp.status_code)
            resp.raise_for_status()
            bean = resp.json()
        code = bean.get("code", 0)
        if code in (1302, 1303):
            raise RuntimeError(f"Rate limited (code={code})")
        if not bean.get("data"):
            msg = bean.get("message", "unknown error")
            raise RuntimeError(f"Submit failed: {msg}")
        return bean

    def execute(self, param: TaskParam) -> TaskResult:
        t0 = time.time()
        logger.info("[%s] Submitting task: POST %s", self.get_type(), self.execute_url)
        try:
            bean = retry_on_failure(
                lambda: self._submit(param.inputs),
                max_retries=5,
                backoff_base=3.0,
                task_name=self.get_type(),
            )
        except Exception as e:
            elapsed = time.time() - t0
            logger.error("[%s] Submit failed in %.2fs: %s", self.get_type(), elapsed, e)
            return TaskResult.failure(f"submit task fail: {e}")

        task_id = bean["data"].get("task_id")
        logger.info("[%s] Task submitted, remote_id=%s", self.get_type(), task_id)

        poll_count = 0
        poll_t0 = time.time()
        while True:
            poll_count += 1
            poll = self._poll(task_id)
            poll_data = poll.get("data", {})
            status = poll_data.get("task_status")
            logger.info("[%s] Poll #%d: remote_id=%s, status=%s, elapsed=%.1fs",
                         self.get_type(), poll_count, task_id, status, time.time() - poll_t0)

            if status in ("submitted", "processing"):
                time.sleep(3)
            elif status == "succeed":
                task_result = poll_data.get("task_result", {})
                break
            else:
                logger.warning("[%s] Unexpected status '%s', treating as failure", self.get_type(), status)
                time.sleep(3)
                return TaskResult.failure(str(poll_data))

        videos_data = task_result.get("videos")
        if not videos_data:
            return TaskResult.failure("No video results found")

        videos = [v["url"] for v in videos_data if v.get("url", "").strip()]
        logger.info("[%s] Completed: %d videos, total_elapsed=%.2fs",
                     self.get_type(), len(videos), time.time() - t0)
        return TaskResult.success({"videos": videos})

    def _poll(self, task_id: str) -> dict:
        url = f"{self.search_url}/{task_id}"
        with logged_client(timeout=60) as client:
            resp = client.get(url, headers=self._headers())
            return resp.json()