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

# MediaKit 服务端偶发 OOM 导致的可重试错误关键词
_RETRYABLE_ERROR_KEYWORDS = ("signal: killed", "InternalError", "OOM", "out of memory")


def _is_retryable_error(error_msg: str) -> bool:
    """判断错误是否可重试（MediaKit 服务端偶发 OOM 等）."""
    error_lower = error_msg.lower()
    return any(kw.lower() in error_lower for kw in _RETRYABLE_ERROR_KEYWORDS)


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 = "TimeInterval",
        max_frames: int = 10,
        poll_interval: float = 2.0,
        max_poll_attempts: int = 30,
        max_retries: int = 1,
    ) -> Optional[List[Dict[str, Any]]]:
        """调用 MediaKit 视频抽帧接口.

        Args:
            video_url: 视频 URL（需可公开访问）
            strategy: 抽帧策略（默认 TimeInterval，比 SceneChange 更稳定不易 OOM）
                - TimeInterval: 按固定时间间隔（推荐，稳定性好）
                - SpecifiedTime: 按指定时间点
                - SpecifiedFrames: 首尾帧 + 指定帧数
                - SceneChange: 场景变化检测（封面选取可用，但高分辨率视频易 OOM）
            max_frames: 最大返回帧数
            poll_interval: 轮询间隔（秒）
            max_poll_attempts: 最大轮询次数
            max_retries: 失败后自动重试次数（仅对可重试错误如 OOM 生效）

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

        for attempt in range(1 + max_retries):
            # 提交抽帧任务
            task_id = self._submit_extract_task(video_url, strategy, max_frames)
            if not task_id:
                return None

            # 轮询任务状态
            result, error_msg = self._poll_task_result_with_error(task_id, poll_interval, max_poll_attempts)

            if result is not None:
                return result

            # 任务失败，判断是否可重试
            if error_msg and _is_retryable_error(error_msg) and attempt < max_retries:
                logger.warning(
                    "MediaKit 抽帧遇到可重试错误，%ds 后重试: " "task_id=%s attempt=%d/%d error=%s",
                    2,
                    task_id,
                    attempt + 1,
                    max_retries,
                    error_msg,
                )
                time.sleep(2)
                continue

            # 不可重试或已用尽重试次数
            if error_msg:
                logger.error(
                    "MediaKit 抽帧最终失败: task_id=%s retryable=%s error=%s",
                    task_id,
                    _is_retryable_error(error_msg),
                    error_msg,
                )
            return None

        return None

    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 analyze_videos(
        self,
        video_urls: List[str],
        prompt: str,
        level: str = "Economy",
        poll_interval: float = 2.0,
        max_poll_attempts: int = 15,
        max_retries: int = 1,
    ) -> Optional[List[str]]:
        """调用 MediaKit 视频理解智能策略 API.

        基于火山方舟视觉大模型，对输入的视频 URL 列表进行内容分析，
        返回每个视频的自然语言描述（用于智能选片段等场景）。

        Args:
            video_urls: 视频 URL 列表（最多 10 个，需公网可访问）
            prompt: 指导大模型分析的自然语言指令
            level: 分析档位 Economy / Balanced / Quality
            poll_interval: 轮询间隔（秒）
            max_poll_attempts: 最大轮询次数
            max_retries: 失败后自动重试次数（仅对可重试错误如 OOM 生效）

        Returns:
            分析结果列表，每个元素对应 video_urls 中同索引视频的分析文本。
            失败返回 None。
        """
        if not self.is_available:
            logger.warning("MediaKit 未配置，跳过视频理解")
            return None

        if not video_urls:
            return None

        for attempt in range(1 + max_retries):
            task_id = self._submit_video_understand_task(video_urls, prompt, level)
            if not task_id:
                return None

            result, error_msg = self._poll_video_understand_result_with_error(task_id, poll_interval, max_poll_attempts)

            if result is not None:
                return result

            # 任务失败，判断是否可重试
            if error_msg and _is_retryable_error(error_msg) and attempt < max_retries:
                logger.warning(
                    "MediaKit 视频理解遇到可重试错误，%ds 后重试: " "task_id=%s attempt=%d/%d error=%s",
                    2,
                    task_id,
                    attempt + 1,
                    max_retries,
                    error_msg,
                )
                time.sleep(2)
                continue

            # 不可重试或已用尽重试次数
            if error_msg:
                logger.error(
                    "MediaKit 视频理解最终失败: task_id=%s retryable=%s error=%s",
                    task_id,
                    _is_retryable_error(error_msg),
                    error_msg,
                )
            return None

        return None

    def _submit_video_understand_task(
        self,
        video_urls: List[str],
        prompt: str,
        level: str,
    ) -> Optional[str]:
        """提交视频理解任务，返回 task_id."""
        url = f"{self.base_url}/tools/video-understand-router"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "video_urls": video_urls,
            "prompt": prompt,
            "level": level,
        }

        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 videos=%d level=%s",
                task_id,
                len(video_urls),
                level,
            )
            return task_id

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

    def _poll_video_understand_result_with_error(
        self,
        task_id: str,
        poll_interval: float,
        max_poll_attempts: int,
    ) -> tuple[Optional[List[str]], Optional[str]]:
        """轮询视频理解任务结果，返回 (contents, error_msg).

        成功时返回 (contents, None)，失败时返回 (None, error_message)。
        """
        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 in ("completed", "success"):
                    result = data.get("result", {})
                    contents = result.get("contents", [])
                    logger.info(
                        "MediaKit 视频理解完成: task_id=%s videos=%d",
                        task_id,
                        len(contents),
                    )
                    return contents if contents else None, None

                elif status == "failed":
                    error_detail = data.get("error", "unknown error")
                    error_node = data.get("error_node", "")
                    full_error = f"{error_detail}" + (f" [node={error_node}]" if error_node else "")
                    logger.error(
                        "MediaKit 视频理解任务失败: task_id=%s error=%s",
                        task_id,
                        full_error,
                    )
                    return None, full_error

                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)

        error_msg = f"超时: max_attempts={max_poll_attempts}"
        logger.error(
            "MediaKit 视频理解超时: task_id=%s max_attempts=%d",
            task_id,
            max_poll_attempts,
        )
        return None, error_msg

    def _poll_video_understand_result(
        self,
        task_id: str,
        poll_interval: float,
        max_poll_attempts: int,
    ) -> Optional[List[str]]:
        """轮询视频理解任务结果，返回 contents 列表（兼容旧接口）."""
        result, _ = self._poll_video_understand_result_with_error(task_id, poll_interval, max_poll_attempts)
        return result

    def _poll_task_result_with_error(
        self,
        task_id: str,
        poll_interval: float,
        max_poll_attempts: int,
    ) -> tuple[Optional[List[Dict[str, Any]]], Optional[str]]:
        """轮询抽帧任务状态，返回 (snapshots, error_msg).

        成功时返回 (snapshots, None)，失败时返回 (None, error_message)。
        """
        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 in ("completed", "success"):
                    result = data.get("result", {})
                    snapshots = result.get("snapshots", [])
                    logger.info(
                        "MediaKit 抽帧完成: task_id=%s frames=%d",
                        task_id,
                        len(snapshots),
                    )
                    return snapshots, None

                elif status == "failed":
                    error_detail = data.get("error", "unknown error")
                    error_node = data.get("error_node", "")
                    full_error = f"{error_detail}" + (f" [node={error_node}]" if error_node else "")
                    logger.error(
                        "MediaKit 抽帧任务失败: task_id=%s error=%s",
                        task_id,
                        full_error,
                    )
                    return None, full_error

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

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

    def _poll_task_result(
        self,
        task_id: str,
        poll_interval: float,
        max_poll_attempts: int,
    ) -> Optional[List[Dict[str, Any]]]:
        """轮询抽帧任务状态，返回结果（兼容旧接口）."""
        result, _ = self._poll_task_result_with_error(task_id, poll_interval, max_poll_attempts)
        return result


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


_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