content-find-agent-new/content_agent/integrations/video_fetch.py

"""视频获取链 (V3-M2A).

从 play_url 下载视频(带平台下载头)→ imageio-ffmpeg 压到 ~4MB 低清 →
base64 data URL,供 GeminiVideoClient 投喂(OpenRouter image_url)。
真实下载/压缩只在 M7 live smoke 跑;单测全 mock。
原片不做本地持久化;长期复盘资产走 OSS,这里仅下载到内存并压缩后投喂模型。
"""

from __future__ import annotations

import base64
import os
import subprocess
import tempfile
import time
from typing import Any

import httpx
import imageio_ffmpeg

# platform_profiles 里写的是 "iOS UA"/"PC UA" 占位,这里映射成真实串 + Referer。
_PLATFORM_DOWNLOAD_HEADERS = {
    "douyin": {
        "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 16_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.0 Mobile/15E148",
        "Referer": "https://www.douyin.com/",
    },
    "shipinhao": {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120 Safari/537.36",
        "Referer": "https://channels.weixin.qq.com/",
    },
}
# 已拍板压缩档:360p / 1fps / 低清,实测 ~4MB(memory/video-multimodal-analysis)。
_FFMPEG_ARGS = ["-vf", "scale=360:-2,fps=1", "-crf", "33", "-c:a", "aac", "-b:a", "32k", "-ac", "1"]
MAX_INLINE_BYTES = 500 * 1024 * 1024  # 本地 inline data URL 护栏;实际上游上限由 OpenRouter/Gemini 决定
DOWNLOAD_TIMEOUT_SECONDS = 30 * 60.0
COMPRESS_TIMEOUT_SECONDS = 20 * 60.0


class VideoFetchError(RuntimeError):
    """下载/压缩/超限失败,由 GeminiVideoClient 捕获转 fail。"""

    def __init__(self, message: str, *, metrics: dict[str, Any] | None = None) -> None:
        super().__init__(message)
        self.metrics = metrics or {}


def _download_headers(platform: str, override: dict[str, str] | None) -> dict[str, str]:
    if override is not None:
        return override
    return _PLATFORM_DOWNLOAD_HEADERS.get(platform, {})


def _compress(
    source: bytes | str,
    ffmpeg_exe: str,
    *,
    timeout_seconds: float = COMPRESS_TIMEOUT_SECONDS,
    input_path: bool = False,
) -> bytes:
    # 超时保护:坏视频会让 ffmpeg 卡死,进而挂住一个判定并发线程(实测正常压缩 ~8s)。
    cmd = [ffmpeg_exe, "-i", str(source) if input_path else "pipe:0", *_FFMPEG_ARGS, "-f", "mp4",
           "-movflags", "frag_keyframe+empty_moov", "pipe:1"]
    try:
        proc = subprocess.run(
            cmd,
            input=None if input_path else source,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout_seconds,
        )
    except subprocess.TimeoutExpired as exc:
        raise VideoFetchError("ffmpeg compression timeout") from exc
    if proc.returncode != 0 or not proc.stdout:
        raise VideoFetchError("ffmpeg compression failed")
    return proc.stdout


def fetch_and_compress(
    play_url: str,
    platform: str,
    *,
    headers: dict[str, str] | None = None,
    http_client: Any | None = None,
    ffmpeg_exe: str | None = None,
    timeout_seconds: float = DOWNLOAD_TIMEOUT_SECONDS,
    compress_timeout_seconds: float = COMPRESS_TIMEOUT_SECONDS,
    clock: Any = time.monotonic,
) -> str:
    data_url, _ = fetch_and_compress_with_metrics(
        play_url,
        platform,
        headers=headers,
        http_client=http_client,
        ffmpeg_exe=ffmpeg_exe,
        timeout_seconds=timeout_seconds,
        compress_timeout_seconds=compress_timeout_seconds,
        clock=clock,
    )
    return data_url


def fetch_and_compress_with_metrics(
    play_url: str,
    platform: str,
    *,
    headers: dict[str, str] | None = None,
    http_client: Any | None = None,
    ffmpeg_exe: str | None = None,
    timeout_seconds: float = DOWNLOAD_TIMEOUT_SECONDS,
    compress_timeout_seconds: float = COMPRESS_TIMEOUT_SECONDS,
    clock: Any = time.monotonic,
) -> tuple[str, dict[str, Any]]:
    metrics: dict[str, Any] = {
        "download_timeout_seconds": timeout_seconds,
        "ffmpeg_timeout_seconds": compress_timeout_seconds,
    }
    total_started_at = clock()
    if not play_url:
        raise VideoFetchError("missing play_url", metrics=metrics)
    client = http_client or httpx
    try:
        download_started_at = clock()
        video_path = _download_to_tempfile(
            client,
            play_url,
            platform,
            headers=headers,
            timeout_seconds=timeout_seconds,
            started_at=total_started_at,
            clock=clock,
        )
        metrics["download_duration_ms"] = _elapsed_ms(download_started_at, clock())
    except httpx.HTTPError as exc:
        metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
        raise VideoFetchError(f"download failed: {type(exc).__name__}", metrics=metrics) from exc
    except VideoFetchError as exc:
        metrics.update(exc.metrics)
        metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
        raise VideoFetchError(str(exc), metrics=metrics) from exc
    try:
        downloaded_bytes = os.path.getsize(video_path)
        metrics["downloaded_bytes"] = downloaded_bytes
        if downloaded_bytes <= 0:
            metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
            raise VideoFetchError("empty download", metrics=metrics)
        ffmpeg_started_at = clock()
        compressed = _compress(
            video_path,
            ffmpeg_exe or imageio_ffmpeg.get_ffmpeg_exe(),
            timeout_seconds=compress_timeout_seconds,
            input_path=True,
        )
        metrics["ffmpeg_duration_ms"] = _elapsed_ms(ffmpeg_started_at, clock())
        metrics["compressed_bytes"] = len(compressed)
    except VideoFetchError as exc:
        metrics.update(exc.metrics)
        metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
        raise VideoFetchError(str(exc), metrics=metrics) from exc
    finally:
        try:
            os.unlink(video_path)
        except OSError:
            pass
    if len(compressed) > MAX_INLINE_BYTES:
        metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
        raise VideoFetchError(f"compressed video oversize: {len(compressed)} bytes", metrics=metrics)
    metrics["total_duration_ms"] = _elapsed_ms(total_started_at, clock())
    return f"data:video/mp4;base64,{base64.b64encode(compressed).decode('ascii')}", metrics


def _download_to_tempfile(
    client: Any,
    play_url: str,
    platform: str,
    *,
    headers: dict[str, str] | None,
    timeout_seconds: float,
    started_at: float,
    clock: Any,
) -> str:
    download_headers = _download_headers(platform, headers)
    tmp = tempfile.NamedTemporaryFile(prefix="content-agent-video-", suffix=".mp4", delete=False)
    tmp_path = tmp.name
    tmp.close()
    ok = False
    if hasattr(client, "stream"):
        try:
            with client.stream(
                "GET",
                play_url,
                headers=download_headers,
                follow_redirects=True,
                timeout=timeout_seconds,
            ) as response:
                response.raise_for_status()
                with open(tmp_path, "wb") as file:
                    for chunk in response.iter_bytes():
                        if clock() - started_at > timeout_seconds:
                            raise VideoFetchError("download timeout")
                        file.write(chunk)
            ok = True
            return tmp_path
        finally:
            if not ok:
                _unlink_quietly(tmp_path)
    try:
        response = client.get(
            play_url,
            headers=download_headers,
            follow_redirects=True,
            timeout=timeout_seconds,
        )
        response.raise_for_status()
        if clock() - started_at > timeout_seconds:
            raise VideoFetchError("download timeout")
        with open(tmp_path, "wb") as file:
            file.write(response.content)
        ok = True
        return tmp_path
    finally:
        if not ok:
            _unlink_quietly(tmp_path)


def _unlink_quietly(path: str) -> None:
    try:
        os.unlink(path)
    except OSError:
        pass


def _elapsed_ms(start: float, end: float) -> int:
    return max(0, int((end - start) * 1000))