ai-flow-python/app/service/task/impl/kontext_pro_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 KontextProTask(BaseTask):
    def get_type(self) -> str:
        return "kontext_pro"

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

    def _headers(self) -> dict[str, str]:
        return {
            "Content-Type": "application/json",
            "x-key": settings.kontext_api_key,
        }

    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()
        if not bean:
            raise RuntimeError("Empty response from API")
        return bean

    def execute(self, param: TaskParam) -> TaskResult:
        logger.info("[%s] Submitting task: POST %s", self.get_type(), self.execute_url)
        t0 = time.time()
        try:
            bean = retry_on_failure(
                lambda: self._submit(param.inputs),
                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.get("id")
        polling_url = bean.get("polling_url")
        logger.info("[%s] Task submitted, remote_id=%s, polling_url=%s",
                     self.get_type(), task_id, polling_url)

        poll_count = 0
        poll_t0 = time.time()
        while True:
            poll_count += 1
            poll = self._poll(polling_url)
            status = poll.get("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 == "Pending":
                time.sleep(3)
            elif status == "Ready":
                result = poll.get("result", {})
                break
            else:
                logger.warning("[%s] Unexpected status '%s', treating as failure", self.get_type(), status)
                time.sleep(3)
                return TaskResult.failure(str(poll))

        if not result:
            return TaskResult.failure("No result found")

        images: list[str] = []
        sample = result.get("sample")
        if sample and sample.strip():
            images.append(sample)

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

    def _poll(self, polling_url: str) -> dict:
        with logged_client(timeout=60) as client:
            resp = client.get(polling_url, headers=self._headers())
            return resp.json()