ForcePilot/backend/package/yuxi/knowledge/parser/video.py

"""视频解析器 — 提取音频 ASR + 关键帧 OCR，输出 Markdown 文本。"""

import tempfile
from pathlib import Path
from typing import Any

from yuxi.knowledge.parser.base import BaseDocumentProcessor, DocumentProcessorException
from yuxi.knowledge.parser.ffmpeg import (
    MAX_VIDEO_BYTES,
    FFmpegError,
    extract_audio_from_video,
    extract_keyframes,
)
from yuxi.utils import logger


class VideoProcessor(BaseDocumentProcessor):
    """视频文件解析器：提取音频 → ASR + 关键帧 OCR。"""

    def get_service_name(self) -> str:
        return "video"

    def get_supported_extensions(self) -> list[str]:
        return [".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"]

    def process_file(self, file_path: str, params: dict[str, Any] | None = None) -> str:
        """
        解析视频文件，返回 Markdown 文本。

        流程:
        1. 提取音频轨道 → ASR 转写
        2. 提取关键帧 → OCR（复用现有 DocumentProcessorFactory）
        3. 合并为 Markdown
        """
        params = params or {}

        # 文件大小检查
        file_size = Path(file_path).stat().st_size
        if file_size > MAX_VIDEO_BYTES:
            raise DocumentProcessorException(
                f"视频文件大小 {file_size / 1024 / 1024:.1f}MB 超过限制 {MAX_VIDEO_BYTES / 1024 / 1024:.0f}MB",
                service_name=self.get_service_name(),
            )

        sections: list[str] = []

        # 1. 提取音频 → ASR
        asr_markdown = self._process_audio_track(file_path, params)
        if asr_markdown:
            sections.append(asr_markdown)

        # 2. 提取关键帧 → OCR
        ocr_markdown = self._process_keyframes(file_path, params)
        if ocr_markdown:
            sections.append(ocr_markdown)

        if not sections:
            raise DocumentProcessorException(
                "视频解析未产生任何内容",
                service_name=self.get_service_name(),
            )

        return "\n\n---\n\n".join(sections)

    def _process_audio_track(self, video_path: str, params: dict) -> str | None:
        """提取音频轨道并 ASR 转写。"""
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            wav_path = tmp.name
        try:
            extract_audio_from_video(video_path, wav_path)
            from yuxi.knowledge.parser.factory import DocumentProcessorFactory
            return DocumentProcessorFactory.process_file("asr", wav_path, params)
        except FFmpegError:
            logger.info("视频无音轨或音频提取失败: %s", video_path)
            return None
        except DocumentProcessorException:
            logger.warning("视频音频 ASR 转写失败: %s", video_path)
            return None
        finally:
            Path(wav_path).unlink(missing_ok=True)

    def _process_keyframes(self, video_path: str, params: dict) -> str | None:
        """提取关键帧并 OCR。"""
        from yuxi.knowledge.parser.factory import DocumentProcessorFactory

        ocr_engine = params.get("ocr_engine", "rapid_ocr")
        if ocr_engine == "disable":
            return None

        with tempfile.TemporaryDirectory() as tmpdir:
            try:
                frames = extract_keyframes(video_path, tmpdir)
            except FFmpegError as e:
                logger.warning("关键帧提取失败: %s — %s", video_path, e)
                return None

            if not frames:
                return None

            ocr_texts: list[str] = []
            processor = DocumentProcessorFactory.get_processor(ocr_engine)
            for frame_path, timestamp in frames:
                try:
                    text = processor.process_file(frame_path, params)
                    if text.strip():
                        mm_ss = f"{int(timestamp) // 60:02d}:{int(timestamp) % 60:02d}"
                        ocr_texts.append(f"**[{mm_ss}]** {text.strip()}")
                except Exception as e:
                    logger.debug("关键帧 OCR 失败: %s — %s", frame_path, e)
                    continue

        if not ocr_texts:
            return None

        return "## 视频画面文字\n\n" + "\n\n".join(ocr_texts)

    def check_health(self) -> dict[str, Any]:
        """检查视频处理依赖（ffmpeg + ASR + OCR）。"""
        # 检查 ffmpeg
        try:
            from yuxi.knowledge.parser.ffmpeg import run_ffmpeg
            run_ffmpeg(["-version"], timeout_ms=3000)
        except Exception:
            return {"status": "unavailable", "message": "ffmpeg 不可用"}

        # 检查 ASR
        from yuxi.knowledge.parser.factory import DocumentProcessorFactory
        asr_health = DocumentProcessorFactory.check_health("asr")

        # 检查 OCR
        ocr_health = DocumentProcessorFactory.check_health("rapid_ocr")

        if asr_health["status"] == "healthy" and ocr_health["status"] == "healthy":
            return {"status": "healthy", "message": "视频解析服务正常"}
        return {"status": "unhealthy", "message": "ASR 或 OCR 服务不可用"}