xiaoxia-saas/packages/shared/mediakit_client.py

"""MediaKit API 客户端（共享层）.

火山引擎 AI MediaKit 提供视频抽帧、视频理解等能力。
API 和 Worker 共用此客户端。

使用方式：
    from packages.shared.mediakit_client import get_mediakit_client

    client = get_mediakit_client()
    if client.is_available:
        frames = client.extract_frames(video_url="https://...")
"""

from __future__ import annotations

import logging
import time
from typing import Any, Dict, List, Optional

import httpx

from packages.shared.config import get_shared_settings

logger = logging.getLogger(__name__)


class MediaKitClient:
    """MediaKit API 客户端.

    封装视频抽帧、视频理解等接口，支持异步任务轮询。
    未配置 API Key 时 is_available 为 False，调用方应降级处理。
    """

    def __init__(self) -> None:
        settings = get_shared_settings()
        self.api_key: str = settings.mediakit_api_key
        self.base_url: str = settings.mediakit_base_url.rstrip("/")
        self.timeout: int = settings.mediakit_timeout

    @property
    def is_available(self) -> bool:
        """是否可用（配置了 API Key）."""
        return bool(self.api_key)

    def extract_frames(
        self,
        video_url: str,
        strategy: str = "SceneChange",
        max_frames: int = 10,
        poll_interval: float = 2.0,
        max_poll_attempts: int = 30,
    ) -> Optional[List[Dict[str, Any]]]:
        """调用 MediaKit 视频抽帧接口.

        Args:
            video_url: 视频 URL（需可公开访问）
            strategy: 抽帧策略
                - TimeInterval: 按固定时间间隔
                - SpecifiedTime: 按指定时间点
                - SpecifiedFrames: 首尾帧 + 指定帧数
                - SceneChange: 场景变化检测（推荐用于封面选取）
            max_frames: 最大返回帧数
            poll_interval: 轮询间隔（秒）
            max_poll_attempts: 最大轮询次数

        Returns:
            帧列表 [{"image_url": "...", "timestamp": 1.5}, ...]
            失败返回 None
        """
        if not self.is_available:
            logger.warning("MediaKit 未配置，跳过抽帧")
            return None

        # 提交抽帧任务
        task_id = self._submit_extract_task(video_url, strategy, max_frames)
        if not task_id:
            return None

        # 轮询任务状态
        return self._poll_task_result(task_id, poll_interval, max_poll_attempts)

    def _submit_extract_task(
        self,
        video_url: str,
        strategy: str,
        max_frames: int,
    ) -> Optional[str]:
        """提交抽帧任务，返回 task_id."""
        url = f"{self.base_url}/tools/extract-frames"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "video_url": video_url,
            "strategy": strategy,
            "max_frames": max_frames,
        }

        try:
            response = httpx.post(url, headers=headers, json=payload, timeout=self.timeout)
            response.raise_for_status()
            data = response.json()

            task_id = data.get("task_id")
            if not task_id:
                logger.error("MediaKit 抽帧任务提交失败: 无 task_id. response=%s", data)
                return None

            logger.info("MediaKit 抽帧任务已提交: task_id=%s strategy=%s", task_id, strategy)
            return task_id

        except Exception as e:
            logger.exception("MediaKit 抽帧任务提交异常: %s", str(e))
            return None

    def _poll_task_result(
        self,
        task_id: str,
        poll_interval: float,
        max_poll_attempts: int,
    ) -> Optional[List[Dict[str, Any]]]:
        """轮询任务状态，返回结果."""
        url = f"{self.base_url}/tasks/{task_id}"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
        }

        for attempt in range(max_poll_attempts):
            try:
                response = httpx.get(url, headers=headers, timeout=self.timeout)
                response.raise_for_status()
                data = response.json()

                status = data.get("status")
                if status == "success":
                    result = data.get("result", {})
                    snapshots = result.get("snapshots", [])
                    logger.info(
                        "MediaKit 抽帧完成: task_id=%s frames=%d",
                        task_id,
                        len(snapshots),
                    )
                    return snapshots

                elif status == "failed":
                    error_msg = data.get("error", "unknown error")
                    logger.error("MediaKit 抽帧任务失败: task_id=%s error=%s", task_id, error_msg)
                    return None

                # status == "processing" or "pending"
                logger.debug(
                    "MediaKit 抽帧进行中: task_id=%s status=%s attempt=%d/%d",
                    task_id,
                    status,
                    attempt + 1,
                    max_poll_attempts,
                )
                time.sleep(poll_interval)

            except Exception as e:
                logger.exception("MediaKit 轮询异常: task_id=%s error=%s", task_id, str(e))
                time.sleep(poll_interval)

        logger.error("MediaKit 抽帧超时: task_id=%s max_attempts=%d", task_id, max_poll_attempts)
        return None


# ── 单例管理 ────────────────────────────────────────────────────────────────


_client_instance: Optional[MediaKitClient] = None


def get_mediakit_client() -> MediaKitClient:
    """获取 MediaKit 客户端单例."""
    global _client_instance
    if _client_instance is None:
        _client_instance = MediaKitClient()
    return _client_instance