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

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 _retry_execute(self, param: TaskParam, remaining: int = 30) -> dict | None:
        logger.info("[%s] Submitting task: POST %s  (remaining_retries=%d)",
                     self.get_type(), self.execute_url, remaining)
        try:
            with logged_client(timeout=60) as client:
                resp = client.post(self.execute_url, headers=self._headers(), json=param.inputs)
                logger.info("[%s] Submit response: status=%d", self.get_type(), resp.status_code)
                bean = resp.json()
        except Exception as e:
            raise RuntimeError(f"Request failed: {e}")

        code = bean.get("code", 0)
        if code in (1302, 1303) and remaining > 0:
            logger.warning("[%s] Rate limited (code=%d), retrying in 5s ... remaining=%d",
                            self.get_type(), code, remaining - 1)
            time.sleep(5)
            return self._retry_execute(param, remaining - 1)
        return bean

    def execute(self, param: TaskParam) -> TaskResult:
        t0 = time.time()
        bean = self._retry_execute(param)
        if not bean:
            return TaskResult.failure("submit task fail !")
        if not bean.get("data"):
            msg = bean.get("message", "unknown error")
            logger.error("[%s] Submit failed: %s", self.get_type(), msg)
            return TaskResult.failure(msg)

        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()