diff --git a/app/tree/video_builder.py b/app/tree/video_builder.py new file mode 100644 index 0000000..2a46d13 --- /dev/null +++ b/app/tree/video_builder.py @@ -0,0 +1,1313 @@ +"""视频树构建模块。 + +将长视频通过 L2 轴心策略 + VLM 帧描述转化为三层 TreeIndex。 + +构建策略:: + + Step 1: _segment_video — 固定步长切分,确定 L1 时间边界 + Step 2: L2 先行 — 从 L3 帧中采样代表帧,VLM 生成 L2Card + Step 3: L3 向下 — 注入 L2 上下文,VLM 批量帧描述,生成 L3Card + Step 4: L1 向上 — 聚合 L2 描述,LLM 生成 L1Card + Step 5: 组装 TreeIndex + Step 6: 字幕 Voronoi 分配(可选) + +并发模型(异步版):: + + build() → asyncio.run(_build_async()) + _build_async(): + asyncio.Semaphore(concurrency) 控制最大 VLM/LLM 并发数 + 各 L1 段并发构建,段内 L2 clip 各启动 _chain 协程: + 提取全部 L3 帧 → 采样 L2 代表帧 → L2 VLM → L3 VLM + 所有 L2+L3 完成后 → L1 LLM + +L2 轴心策略解决了循环依赖: + - L2 描述从 L3 帧中采样代表帧直接生成 + - L3 注入 L2 上下文后批量/逐帧描述 + - L1 聚合 L2 描述,保证完整覆盖 + +帧持久化: + - 帧图像保存到 {cache_dir}/frames/{video_stem}/,长期有效 + - 已提取的帧自动跳过(缓存复用) + +核心算法保真(CLAUDE.md §4.7): + #1 L2 轴心建树策略 + #2 VLM 批量帧描述 + JSON fallback + #3 断点续跑机制 +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import re +import subprocess +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import cv2 +from loguru import logger + +from app.tree.index import ( + IndexMeta, + L1Card, + L1Node, + L2Card, + L2Node, + L3Card, + L3Node, + TreeIndex, + load_l1_json, + save_l1_json, +) +from app.tree.subtitle import ( + SRTEntry, + assign_subtitles_voronoi, + extract_subtitle_for_range, +) + +if TYPE_CHECKING: + from app.tree.config import TreeConfig + from core.protocols import LLMProvider, VLMProvider + +# --------------------------------------------------------------------------- +# Prompt 常量(结构化 JSON 输出版本,保真原始 prompt 风格) +# --------------------------------------------------------------------------- + +_L2_VIDEO_PROMPT = ( + "用1-2句话描述以下视频片段的核心内容,与同级片段形成区分。\n" + "{subtitle_block}" + "返回 JSON 对象,包含以下字段:\n" + "- event_description: 1-2句片段描述\n" + "- entities: 可见实体列表\n" + "- actions: 动作列表\n" + "- action_subjects: 动作主体列表\n" + "- visible_text: 画面中可见文字列表\n" + "- spatial_relations: 空间关系描述\n" + "- state_changes: 状态变化描述(无则 null)\n" + "只返回 JSON 对象,不要其他内容。" +) + +_L3_VIDEO_PROMPT = ( + '该片段的整体内容: "{l2_description}"\n' + "以下是该片段中连续的 {n} 帧画面。\n" + "对每帧用一到两句话描述其具体画面内容。\n" + "重点关注: 动作、物体变化、文字信息、人物表情。\n" + "不要重复片段整体描述,聚焦每帧的区分性信息。\n" + "{subtitle_block}" + "对每帧返回一个 JSON 对象,包含以下字段:\n" + "- frame_summary: 1-2句画面描述\n" + "- visible_entities: 可见实体列表\n" + "- ongoing_actions: 正在进行的动作列表\n" + "- visible_text: 画面中可见文字列表\n" + "- spatial_layout: 画面空间布局\n" + '- visual_attributes: {{"lighting": "...", "dominant_colors": [...], "camera_angle": "..."}}\n' + "只返回 JSON 数组,格式: [{{...}}, {{...}}, ...],不要其他内容。" +) + +_L3_SINGLE_PROMPT = ( + '该片段的整体内容: "{l2_description}"\n' + "用一到两句话描述这帧画面的具体内容。" + "重点关注: 动作、物体变化、文字信息、人物表情。\n" + "{subtitle_block}" + "返回 JSON 对象,包含以下字段:\n" + "- frame_summary: 画面描述\n" + "- visible_entities: 可见实体列表\n" + "- ongoing_actions: 动作列表\n" + "- visible_text: 可见文字列表\n" + "- spatial_layout: 空间布局\n" + '- visual_attributes: {{"lighting": "...", "dominant_colors": [...], "camera_angle": "..."}}\n' + "只返回 JSON 对象,不要其他内容。" +) + +_L1_VIDEO_PROMPT = ( + "以下是一个视频段落中各片段的描述:\n{l2_texts}\n" + "用2-3句话总结该段落的整体内容,涵盖所有片段的主题。\n" + "返回 JSON 对象,包含以下字段:\n" + "- scene_summary: 2-3句段落摘要\n" + "- main_setting: 主要场景\n" + "- key_entities: 关键实体列表\n" + "- main_actions: 主要动作列表\n" + "- topic_keywords: 主题关键词列表\n" + "- visible_text: 出现的文字列表\n" + "- temporal_flow: 时间流向描述\n" + "只返回 JSON 对象,不要其他内容。" +) + +# 每次 VLM 调用携带的最大帧数:5 帧 payload 小、JSON 解析成功率高 +_L3_BATCH_SIZE = 5 + +# ffmpeg 并发提帧的线程池大小(CPU 密集型,避免过度并发) +_FFMPEG_MAX_WORKERS = 8 + + +# --------------------------------------------------------------------------- +# 主类 +# --------------------------------------------------------------------------- + + +class VideoTreeBuilder: + """视频模态树构建器(asyncio 真并发版)。 + + 将长视频通过 L2 轴心策略(先构建 L2,再向下扩展 L3,向上聚合 L1) + 转化为三层 TreeIndex。 + + 并发架构: + build() 为同步壳,内部调用 asyncio.run(_build_async())。 + _build_async() 使用 asyncio.Semaphore(concurrency) 控制并发 VLM/LLM 数量。 + 所有 VLM 调用通过 VLMProvider 的异步接口发起,零线程阻塞。 + 所有 LLM 调用通过 LLMProvider 的异步接口发起(L1 摘要)。 + ffmpeg 提帧在独立 ThreadPoolExecutor 中并行,不阻塞事件循环。 + + 属性: + _vlm: VLM 图文调用端口。 + _llm: LLM 文本调用端口(L1 摘要)。 + _config: 树构建配置。 + _ffmpeg_pool: ffmpeg 专用线程池(max_workers=_FFMPEG_MAX_WORKERS)。 + """ + + def __init__( + self, + vlm: VLMProvider, + llm: LLMProvider, + config: TreeConfig, + ) -> None: + """初始化视频树构建器。 + + 参数: + vlm: VLM 图文调用端口(VLMProvider Protocol)。 + llm: LLM 文本调用端口(LLMProvider Protocol),用于 L1 摘要。 + config: 树构建配置(TreeConfig),关键字段: + l1_segment_duration, l2_clip_duration, l3_fps, + l2_representative_frames, cache_dir, concurrency。 + """ + self._vlm = vlm + self._llm = llm + self._config = config + self._ffmpeg_pool = ThreadPoolExecutor(max_workers=_FFMPEG_MAX_WORKERS) + self._cache_root = Path(self._config.cache_dir) + self._session_id: str = "" + + # ------------------------------------------------------------------ + # URL 流式辅助方法 + # ------------------------------------------------------------------ + + @staticmethod + def _is_url(path_or_url: str) -> bool: + """判断输入是否为网络 URL(而非本地路径)。 + + 参数: + path_or_url: 文件路径或 URL 字符串。 + + 返回: + True 表示 URL,False 表示本地路径。 + """ + return path_or_url.startswith(("http://", "https://")) + + @staticmethod + def _source_stem(video_path: str) -> str: + """从视频路径或 YouTube URL 中提取短标识符,用于帧缓存目录命名。 + + 参数: + video_path: 本地文件路径或 YouTube 视频页面 URL。 + + 返回: + 短字符串标识符(本地文件取 stem,YouTube URL 取 v= 后的视频 ID)。 + """ + if "youtube.com/watch" in video_path or "youtu.be/" in video_path: + match = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{8,15})", video_path) + if match: + return match.group(1) + stem = Path(video_path).stem + return stem[:64] if len(stem) > 64 else stem + + @staticmethod + def _resolve_stream(url: str) -> str: + """通过 yt-dlp 获取 YouTube 视频的 CDN 直链。 + + 参数: + url: YouTube 视频页面 URL。 + + 返回: + CDN HTTPS 直链。 + """ + logger.info("获取 YouTube CDN 直链", url=url) + result = subprocess.run( + [ + "yt-dlp", + "-g", + "--format", + "best[ext=mp4][height<=720]/best[ext=mp4]/best", + url, + ], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, f"yt-dlp 获取直链失败: {result.stderr.strip()}" + stream_url = result.stdout.strip().splitlines()[0] + assert stream_url.startswith("http"), f"yt-dlp 返回非 URL: {stream_url[:100]}" + logger.info("CDN 直链获取成功", stream_url=stream_url[:80]) + return stream_url + + @staticmethod + def _get_video_duration(url: str) -> float: + """通过 yt-dlp --dump-json 获取视频时长(秒)。 + + 参数: + url: YouTube 视频页面 URL。 + + 返回: + 视频总时长(秒,浮点数)。 + """ + logger.info("获取视频时长元数据", url=url) + result = subprocess.run( + ["yt-dlp", "--dump-json", "--no-playlist", url], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, f"yt-dlp 元数据获取失败: {result.stderr.strip()}" + meta = json.loads(result.stdout) + duration = float(meta.get("duration", 0)) + assert duration > 0, f"视频时长读取异常: {duration}" + logger.info("视频时长确认", duration_sec=round(duration, 1)) + return duration + + # ------------------------------------------------------------------ + # 公共接口 + # ------------------------------------------------------------------ + + def build( + self, + video_path: str, + srt_entries: list[SRTEntry] | None = None, + ) -> TreeIndex: + """将长视频构建为三层 TreeIndex(同步壳,内部 asyncio.run 驱动)。 + + 参数: + video_path: 视频文件路径(.mp4/.avi/.mkv 等)或 YouTube URL。 + srt_entries: 可选的 SRT 字幕条目列表, + 若提供则注入 VLM prompt 并执行 Voronoi 字幕分配。 + + 返回: + 三层 TreeIndex 对象。 + """ + return asyncio.run(self._build_async(video_path, srt_entries)) + + # ------------------------------------------------------------------ + # 核心异步构建逻辑(保真算法 #1:L2→L3 链式触发) + # ------------------------------------------------------------------ + + async def _build_async( + self, + video_path: str, + srt_entries: list[SRTEntry] | None = None, + ) -> TreeIndex: + """异步构建三层 TreeIndex(真并发核心,L2→L3 链式触发)。 + + 参数: + video_path: 视频文件路径或 YouTube URL。 + srt_entries: 可选的 SRT 字幕条目列表。 + + 返回: + 三层 TreeIndex 对象。 + + 实现细节: + 并发架构:每个 L1 段内启动一组"L2→L3 链式协程", + L2 完成后立即触发 L3(不等待其他 L2),L3 完成后触发 L1 摘要。 + 各 L1 段独立并发,彼此不阻塞。 + Semaphore(concurrency) 全局限制同时在途 VLM/LLM 调用数量。 + + 关键调用链(每个 L2 clip 独立,保真算法 #1):: + _build_segment(i) → asyncio.gather( + _chain(i,0): extract_frames → sample_l2 → L2_VLM → L3_VLM + _chain(i,1): extract_frames → sample_l2 → L2_VLM → L3_VLM + ... + ) → _build_l1_video_async(i) + """ + # Phase 0: URL vs 本地文件处理 + if self._is_url(video_path): + stream_url = self._resolve_stream(video_path) + duration_hint: float | None = self._get_video_duration(video_path) + logger.info("开始构建视频树索引(URL 流式模式)", source_url=video_path) + else: + assert os.path.isfile(video_path), f"视频文件不存在: {video_path}" + stream_url = video_path + duration_hint = None + logger.info("开始构建视频树索引", video_path=video_path) + + source_id = self._source_stem(video_path) + self._session_id = f"build_{source_id}" + + # Phase 1: 时间切分(同步,仅一次) + l1_ranges = self._segment_video(stream_url, duration_hint=duration_hint) + assert len(l1_ranges) > 0, "视频时间切分结果为空" + logger.info("视频切分完成", l1_count=len(l1_ranges)) + + total_l1 = len(l1_ranges) + + # Phase 1.1: 读取已有进度(保真算法 #3:断点续跑) + finished_l1_ids = self._load_resume_state(source_id, total_l1) + + # 创建 VLM/LLM 并发控制信号量 + vlm_sem = asyncio.Semaphore(self._config.concurrency) + + # Phase 2-5: 按 L1 段并发,段内 L2→L3 链式触发(保真算法 #1) + async def _build_segment( + i: int, + l1_range: tuple[float, float], + ) -> L1Node: + """单个 L1 段的完整构建:L2+L3 并发链式 → L1 摘要。 + + 参数: + i: L1 段索引。 + l1_range: L1 时间区间 (start, end)。 + + 返回: + 完整的 L1Node(含所有 L2 和 L3 子节点)。 + """ + clips = self._get_l2_clips(l1_range) + + async def _chain( + j: int, + clip_range: tuple[float, float], + ) -> tuple[int, L2Node]: + """L2→L3 链:提取帧→采样→L2 VLM→L3 VLM。""" + l2_id = f"l1_{i}_l2_{j}" + + # Phase A: 提取该 clip 的全部 L3 帧 + all_frames = await self._extract_frames_async( + stream_url, + clip_range, + self._config.l3_fps, + source_id=source_id, + ) + assert len(all_frames) > 0, f"L2 clip {l2_id} 帧提取结果为空" + + # Phase B: 从 L3 帧中采样 L2 代表帧 + l2_rep_paths = self._sample_representative_frames( + all_frames, + self._config.l2_representative_frames, + ) + + # Phase C: L2 VLM 描述 + l2_node = await self._build_l2_video_async( + l2_rep_paths, + clip_range, + l2_id, + vlm_sem, + srt_entries, + ) + logger.info("L2 VLM 完成,已触发 L3 任务", l2_id=l2_id) + + # Phase D: L3 VLM 描述(注入 L2 上下文) + l3_nodes = await self._build_l3_video_async( + all_frames, + l2_node.description, + i, + j, + vlm_sem, + srt_entries, + ) + l2_node.children = l3_nodes + logger.info( + "L3 完成", + l2_id=l2_id, + l3_count=len(l3_nodes), + ) + return (j, l2_node) + + # 所有 clip 同时启动(保真算法 #1:asyncio.gather 链式并发) + pairs = await asyncio.gather(*[_chain(j, clip) for j, clip in enumerate(clips)]) + ordered_l2 = [p[1] for p in sorted(pairs, key=lambda x: x[0])] + + logger.info("L1 触发", l1_id=f"l1_{i}") + l1_node = await self._build_l1_video_async( + ordered_l2, + f"l1_{i}", + l1_range, + vlm_sem, + ) + logger.info( + "L1 节点构建完成", + l1_id=f"l1_{i}", + l2_count=len(ordered_l2), + ) + return l1_node + + total_clips = sum(len(self._get_l2_clips(r)) for r in l1_ranges) + logger.info( + "开始并发构建(L2→L3链式,L1段间并发,支持断点续跑)", + total_l2=total_clips, + concurrency=self._config.concurrency, + ) + + # Phase 2: 并发构建尚未完成的 L1 段(保真算法 #3:断点续跑) + tasks: list[asyncio.Task[L1Node]] = [] + task_indices: list[int] = [] + for i, r in enumerate(l1_ranges): + if i in finished_l1_ids and self._has_l1_intermediate(source_id, i): + continue + tasks.append(asyncio.create_task(_build_segment(i, r))) + task_indices.append(i) + + new_l1_nodes: dict[int, L1Node] = {} + if tasks: + results = await asyncio.gather(*tasks) + for idx, node in zip(task_indices, results, strict=False): + self._save_l1_intermediate(source_id, node, idx) + finished_l1_ids.add(idx) + new_l1_nodes[idx] = node + self._save_progress(source_id, total_l1, finished_l1_ids) + + # Phase 3: 汇总所有 L1 段(中间 + 新生成,保真算法 #3) + l1_nodes = self._assemble_roots( + new_l1_nodes, + finished_l1_ids, + total_l1, + source_id, + ) + + # Phase 6: 组装 TreeIndex + metadata = IndexMeta( + source_path=video_path, + modality="video", + created_at=datetime.now().isoformat(), + ) + index = TreeIndex(metadata=metadata, roots=l1_nodes) + + # Phase 7: 字幕 Voronoi 分配(可选) + if srt_entries: + assign_subtitles_voronoi(index, srt_entries) + logger.info("字幕 Voronoi 分配完成", n_entries=len(srt_entries)) + + total_l2_count = sum(len(r.children) for r in l1_nodes) + total_l3_count = sum(len(l2.children) for r in l1_nodes for l2 in r.children) + logger.info( + "视频树索引构建完成", + source_path=video_path, + l1=len(l1_nodes), + l2=total_l2_count, + l3=total_l3_count, + ) + + # Phase 8: 清理中间文件(保真算法 #3:构建成功后清理) + self._cleanup_intermediate_and_progress(source_id) + return index + + # ------------------------------------------------------------------ + # 内部方法:时间切分(同步,仅执行一次) + # ------------------------------------------------------------------ + + def _segment_video( + self, + video_path: str, + duration_hint: float | None = None, + ) -> list[tuple[float, float]]: + """读取视频总时长,按固定步长切分为 L1 时间区间列表。 + + 参数: + video_path: 视频文件路径或 CDN 流式 URL。 + duration_hint: 已知视频时长(秒),传入时跳过 cv2 读取。 + + 返回: + L1 时间区间列表,每项为 (start_sec, end_sec)。 + """ + if duration_hint is not None: + total_duration = duration_hint + else: + cap = cv2.VideoCapture(video_path) + assert cap.isOpened(), f"无法打开视频文件: {video_path}" + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) + cap.release() + assert fps > 0, f"视频 FPS 读取异常: {fps}" + assert total_frames > 0, f"视频总帧数读取异常: {total_frames}" + total_duration = total_frames / fps + + step = self._config.l1_segment_duration + ranges: list[tuple[float, float]] = [] + start = 0.0 + while start < total_duration: + end = min(start + step, total_duration) + ranges.append((start, end)) + start = end + + logger.info( + "L1 时间切分", + total_duration=round(total_duration, 2), + l1_count=len(ranges), + ) + return ranges + + def _get_l2_clips( + self, + l1_range: tuple[float, float], + ) -> list[tuple[float, float]]: + """将 L1 时间区间等分为 L2 clips。 + + 参数: + l1_range: L1 时间区间 (start, end),单位秒。 + + 返回: + L2 clip 时间区间列表。 + """ + start, end = l1_range + step = self._config.l2_clip_duration + clips: list[tuple[float, float]] = [] + t = start + while t < end: + clip_end = min(t + step, end) + clips.append((t, clip_end)) + t = clip_end + return clips + + # ------------------------------------------------------------------ + # 内部方法:帧提取(ffmpeg subprocess,在线程池执行) + # ------------------------------------------------------------------ + + def _ffmpeg_extract_frame( + self, + video_path: str, + ts: float, + out_path: str, + ) -> bool: + """用 ffmpeg subprocess 提取单帧图像。 + + 参数: + video_path: 视频文件路径(本地 MP4 或 CDN URL)。 + ts: 目标时间戳(秒)。 + out_path: 输出 JPEG 文件路径。 + + 返回: + True 表示提取成功,False 表示失败。 + """ + cmd = [ + "ffmpeg", + "-hide_banner", + "-loglevel", + "error", + "-ss", + f"{ts:.3f}", + "-i", + video_path, + "-frames:v", + "1", + "-q:v", + "2", + "-y", + out_path, + ] + result = subprocess.run(cmd, capture_output=True) + return result.returncode == 0 and os.path.isfile(out_path) + + async def _extract_frames_async( + self, + video_path: str, + time_range: tuple[float, float], + fps: float, + source_id: str | None = None, + ) -> list[tuple[str, float]]: + """异步并发提取时间范围内的帧,保存到 cache 目录。 + + 参数: + video_path: 视频文件路径或 CDN 流式 URL。 + time_range: 提取时间区间 (start_sec, end_sec)。 + fps: 提取帧率(帧/秒)。 + source_id: 帧缓存目录名。 + + 返回: + [(frame_path, timestamp_sec), ...],按时间顺序排列。 + """ + video_stem = source_id if source_id is not None else self._source_stem(video_path) + frame_dir = Path(self._config.cache_dir) / "frames" / video_stem + frame_dir.mkdir(parents=True, exist_ok=True) + + start_sec, end_sec = time_range + step = 1.0 / fps + + timestamps: list[float] = [] + t = start_sec + while t < end_sec: + timestamps.append(t) + t += step + + if not timestamps: + logger.warning( + "帧提取时间区间内无有效时间戳", + time_range=time_range, + fps=fps, + ) + return [] + + loop = asyncio.get_running_loop() + + async def _extract_one(ts: float) -> tuple[str, float] | None: + """提取单帧:缓存命中直接返回,否则在线程池中调用 ffmpeg。""" + frame_name = f"{start_sec:.1f}_{ts:.3f}.jpg" + frame_path = str(frame_dir / frame_name) + + if os.path.isfile(frame_path): + return (frame_path, ts) + + success = await loop.run_in_executor( + self._ffmpeg_pool, + self._ffmpeg_extract_frame, + video_path, + ts, + frame_path, + ) + if not success: + logger.warning( + "帧读取失败,跳过", + timestamp=ts, + video_path=video_path, + ) + return None + return (frame_path, ts) + + results = await asyncio.gather(*[_extract_one(ts) for ts in timestamps]) + return [r for r in results if r is not None] + + # ------------------------------------------------------------------ + # 内部方法:帧采样(L2 代表帧复用 L3 帧) + # ------------------------------------------------------------------ + + @staticmethod + def _sample_representative_frames( + frames: list[tuple[str, float]], + n: int, + ) -> list[str]: + """从 L3 帧列表中均匀采样 n 帧路径,用于 L2 VLM 描述。 + + 参数: + frames: L3 帧列表 [(frame_path, timestamp), ...]。 + n: 目标采样数。 + + 返回: + 采样的帧路径列表,长度为 min(n, len(frames))。 + """ + if n >= len(frames): + return [fp for fp, _ in frames] + step = len(frames) / n + return [frames[int(i * step)][0] for i in range(n)] + + # ------------------------------------------------------------------ + # 内部方法:字幕辅助 + # ------------------------------------------------------------------ + + def _build_subtitle_block( + self, + srt_entries: list[SRTEntry] | None, + time_range: tuple[float, float], + ) -> str: + """构建字幕注入文本块。无字幕或无匹配时返回空字符串。 + + 参数: + srt_entries: SRT 字幕条目列表。 + time_range: 时间范围 (start, end)。 + 若 start >= end(如单帧),自动扩展为窗口。 + + 返回: + 字幕文本块字符串,含前后换行。 + """ + if not srt_entries: + return "" + start, end = time_range + if end <= start: + start = max(0.0, start - self._config.srt_window_sec) + end = end + self._config.srt_window_sec + text = extract_subtitle_for_range(srt_entries, (start, end)) + if not text: + return "" + return f"字幕信息:\n{text}\n" + + # ------------------------------------------------------------------ + # 内部方法:L1 中间结果与进度管理(保真算法 #3:断点续跑) + # ------------------------------------------------------------------ + + def _intermediate_dir(self, stem: str) -> Path: + """获取某视频的中间结果目录路径。""" + return self._cache_root / "intermediate" / stem + + def _progress_path(self, stem: str) -> Path: + """获取某视频的进度文件路径。""" + return self._cache_root / "progress" / f"{stem}.json" + + def _has_l1_intermediate(self, stem: str, l1_idx: int) -> bool: + """检查某 L1 段的中间 JSON 是否存在。""" + path = self._intermediate_dir(stem) / f"l1_{l1_idx}.json" + return path.is_file() + + def _save_l1_intermediate( + self, + stem: str, + l1_node: L1Node, + l1_idx: int, + ) -> None: + """将单个 L1 段的中间结果保存到 JSON 文件。""" + dir_path = self._intermediate_dir(stem) + dir_path.mkdir(parents=True, exist_ok=True) + out_path = dir_path / f"l1_{l1_idx}.json" + save_l1_json(str(out_path), l1_node) + + def _load_l1_intermediate( + self, + stem: str, + l1_idx: int, + ) -> L1Node | None: + """从中间 JSON 加载单个 L1 段,若不存在则返回 None。""" + path = self._intermediate_dir(stem) / f"l1_{l1_idx}.json" + if not path.is_file(): + return None + return load_l1_json(str(path)) + + def _load_progress(self, stem: str) -> dict[str, Any] | None: + """加载某视频的进度文件(若不存在则返回 None)。""" + path = self._progress_path(stem) + if not path.is_file(): + return None + with open(path, encoding="utf-8") as f: + try: + data: dict[str, Any] = json.load(f) + except json.JSONDecodeError: + logger.warning("进度文件 JSON 解析失败,忽略", path=str(path)) + return None + return data + + def _save_progress( + self, + stem: str, + total_l1: int, + finished_l1_ids: set[int], + ) -> None: + """将最新进度写回磁盘(保真算法 #3)。""" + path = self._progress_path(stem) + path.parent.mkdir(parents=True, exist_ok=True) + payload: dict[str, Any] = { + "video_id": stem, + "total_l1": total_l1, + "finished_l1_ids": sorted(finished_l1_ids), + "updated_at": datetime.now().isoformat(), + } + if not path.is_file(): + payload["created_at"] = payload["updated_at"] + else: + old = self._load_progress(stem) + if old and isinstance(old.get("created_at"), str): + payload["created_at"] = old["created_at"] + with open(path, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + logger.info( + "进度文件已更新", + path=str(path), + total_l1=total_l1, + finished_l1=sorted(finished_l1_ids), + ) + + def _load_resume_state( + self, + source_id: str, + total_l1: int, + ) -> set[int]: + """加载断点续跑状态,返回已完成的 L1 段索引集合。 + + 参数: + source_id: 视频源标识符(用于查找进度文件)。 + total_l1: 当前切分产生的 L1 段总数。 + + 返回: + 已完成的 L1 段索引集合;无有效进度时返回空集。 + """ + progress = self._load_progress(source_id) + if progress is None: + return set() + + if progress.get("total_l1") != total_l1: + logger.warning( + "进度文件与当前 L1 段数不一致,忽略旧进度", + stem=source_id, + recorded_total_l1=progress.get("total_l1"), + current_total_l1=total_l1, + ) + return set() + + finished: set[int] = set(progress.get("finished_l1_ids", [])) + if finished: + logger.info( + "检测到中间进度,启用断点续跑", + stem=source_id, + finished_l1=sorted(finished), + ) + return finished + + def _assemble_roots( + self, + new_l1_nodes: dict[int, L1Node], + finished_l1_ids: set[int], + total_l1: int, + source_id: str, + ) -> list[L1Node]: + """汇总所有 L1 段(新构建 + 中间缓存),按索引顺序返回。 + + 参数: + new_l1_nodes: 本次新构建的 L1 节点 {索引: 节点}。 + finished_l1_ids: 所有已完成的 L1 段索引(含历史 + 本次)。 + total_l1: L1 段总数。 + source_id: 视频源标识符(用于查找中间 JSON)。 + + 返回: + 按索引排序的 L1Node 列表。 + """ + l1_nodes: list[L1Node] = [] + for i in range(total_l1): + if i in new_l1_nodes: + l1_nodes.append(new_l1_nodes[i]) + continue + node = self._load_l1_intermediate(source_id, i) + assert node is not None, f"L1 段 {i} 缺失中间结果,无法恢复" + l1_nodes.append(node) + return l1_nodes + + def _cleanup_intermediate_and_progress(self, stem: str) -> None: + """在最终构建成功后清理中间结果与进度文件(保真算法 #3)。""" + progress_path = self._progress_path(stem) + if progress_path.is_file(): + try: + progress_path.unlink() + except OSError: + logger.warning("删除进度文件失败", path=str(progress_path)) + + inter_dir = self._intermediate_dir(stem) + if inter_dir.is_dir(): + for child in inter_dir.glob("l1_*.json"): + try: + child.unlink() + except OSError: + logger.warning( + "删除 L1 中间 JSON 失败", + path=str(child), + ) + with contextlib.suppress(OSError): + inter_dir.rmdir() + + # ------------------------------------------------------------------ + # 内部方法:异步节点构建 + # ------------------------------------------------------------------ + + async def _build_l2_video_async( + self, + rep_frame_paths: list[str], + clip_range: tuple[float, float], + l2_id: str, + vlm_sem: asyncio.Semaphore, + srt_entries: list[SRTEntry] | None, + ) -> L2Node: + """异步构建 L2 视频节点(VLM 代表帧描述,输出 L2Card)。 + + 参数: + rep_frame_paths: 已从 L3 帧中采样的代表帧路径列表。 + clip_range: L2 clip 时间区间 (start, end),单位秒。 + l2_id: 节点 ID。 + vlm_sem: VLM 并发控制信号量。 + srt_entries: 可选的 SRT 字幕条目列表。 + + 返回: + L2Node(children 为空,由后续 L3 阶段填充)。 + """ + assert len(rep_frame_paths) > 0, f"L2 节点 {l2_id} 代表帧列表为空" + + subtitle_block = self._build_subtitle_block(srt_entries, clip_range) + prompt = _L2_VIDEO_PROMPT.format(subtitle_block=subtitle_block) + messages = [{"role": "user", "content": prompt}] + + async with vlm_sem: + response = await self._vlm.chat_with_images( + messages, + images=rep_frame_paths, + session_id=self._session_id, + ) + + card = self._parse_l2_card(response.content) + return L2Node(id=l2_id, card=card, time_range=clip_range) + + async def _build_l3_video_async( + self, + frames: list[tuple[str, float]], + l2_description: str, + l1_i: int, + l2_j: int, + vlm_sem: asyncio.Semaphore, + srt_entries: list[SRTEntry] | None, + ) -> list[L3Node]: + """异步批次级并发构建 L3 节点(核心加速点,保真算法 #2)。 + + 参数: + frames: [(frame_path, timestamp), ...]。 + l2_description: L2 节点描述,注入 prompt 上下文。 + l1_i: 父 L1 索引(用于节点 ID 生成)。 + l2_j: 父 L2 索引(用于节点 ID 生成)。 + vlm_sem: VLM 并发控制信号量。 + srt_entries: 可选的 SRT 字幕条目列表。 + + 返回: + L3Node 列表,每项对应一帧。 + """ + assert len(frames) > 0, f"L3 帧列表为空 (l1={l1_i}, l2={l2_j})" + + # Phase 1: 分批并发 VLM 调用(保真算法 #2:_L3_BATCH_SIZE=5) + batches: list[list[tuple[str, float]]] = [] + for batch_start in range(0, len(frames), _L3_BATCH_SIZE): + batches.append(frames[batch_start : batch_start + _L3_BATCH_SIZE]) + + batch_results: list[list[L3Card]] = list( + await asyncio.gather( + *[ + self._call_vlm_batch_async( + batch, + l2_description, + l1_i, + l2_j, + vlm_sem, + srt_entries, + ) + for batch in batches + ] + ) + ) + + # Phase 2: 展平所有批次卡片,构建 L3 节点 + all_cards: list[L3Card] = [card for batch in batch_results for card in batch] + + nodes: list[L3Node] = [] + for k, (card, (frame_path, ts)) in enumerate(zip(all_cards, frames, strict=False)): + nodes.append( + L3Node( + id=f"l1_{l1_i}_l2_{l2_j}_l3_{k}", + card=card, + timestamp=ts, + frame_path=frame_path, + ) + ) + return nodes + + async def _call_vlm_batch_async( + self, + batch: list[tuple[str, float]], + l2_description: str, + l1_i: int, + l2_j: int, + vlm_sem: asyncio.Semaphore, + srt_entries: list[SRTEntry] | None, + ) -> list[L3Card]: + """异步单批次 VLM 调用(保真算法 #2:批量→逐帧 fallback)。 + + 参数: + batch: 本批帧列表 [(frame_path, ts), ...],长度 <= _L3_BATCH_SIZE。 + l2_description: L2 描述,用于 prompt 和 fallback prompt。 + l1_i: 父 L1 索引(日志用)。 + l2_j: 父 L2 索引(日志用)。 + vlm_sem: VLM 并发控制信号量。 + srt_entries: 可选的 SRT 字幕条目列表。 + + 返回: + 与 batch 等长的 L3Card 列表。 + """ + batch_paths = [fp for fp, _ in batch] + n = len(batch_paths) + + batch_time_range = (batch[0][1], batch[-1][1]) + subtitle_block = self._build_subtitle_block(srt_entries, batch_time_range) + + prompt = _L3_VIDEO_PROMPT.format( + l2_description=l2_description, + n=n, + subtitle_block=subtitle_block, + ) + messages = [{"role": "user", "content": prompt}] + + # Phase 1: 尝试批量调用(保真算法 #2) + try: + async with vlm_sem: + response = await self._vlm.chat_with_images( + messages, + images=batch_paths, + session_id=self._session_id, + ) + cards = self._parse_l3_cards(response.content, n) + if cards is not None: + return cards + logger.warning( + "L3 小批量 VLM JSON 解析失败,对本批逐帧 fallback", + l1=l1_i, + l2=l2_j, + batch_n=n, + raw_preview=response.content[:100], + ) + except Exception as exc: + logger.warning( + "L3 小批量 VLM 调用异常,对本批逐帧 fallback: {}", + exc, + l1=l1_i, + l2=l2_j, + batch_n=n, + ) + + # Phase 2: 逐帧 fallback(并发,受信号量保护,保真算法 #2) + async def _single_frame(fp: str, ts: float) -> L3Card: + single_time_range = (ts, ts) + sub_block = self._build_subtitle_block( + srt_entries, + single_time_range, + ) + single_prompt = _L3_SINGLE_PROMPT.format( + l2_description=l2_description, + subtitle_block=sub_block, + ) + single_messages = [{"role": "user", "content": single_prompt}] + async with vlm_sem: + resp = await self._vlm.chat_with_images( + single_messages, + images=[fp], + session_id=self._session_id, + ) + return self._parse_l3_card_single(resp.content) + + return list(await asyncio.gather(*[_single_frame(fp, ts) for fp, ts in batch])) + + async def _build_l1_video_async( + self, + l2_children: list[L2Node], + l1_id: str, + l1_range: tuple[float, float], + vlm_sem: asyncio.Semaphore, + ) -> L1Node: + """异步构建 L1 节点(LLM 文本摘要,输出 L1Card)。 + + 参数: + l2_children: 该 L1 节点下的所有 L2 节点。 + l1_id: 节点 ID。 + l1_range: L1 时间区间 (start, end),单位秒。 + vlm_sem: VLM/LLM 并发控制信号量。 + + 返回: + L1Node(children 已赋值)。 + """ + assert len(l2_children) > 0, f"L1 节点 {l1_id} 没有 L2 子节点" + l2_texts = "\n".join(f"- {node.description}" for node in l2_children) + prompt = _L1_VIDEO_PROMPT.format(l2_texts=l2_texts) + messages = [{"role": "user", "content": prompt}] + + async with vlm_sem: + response = await self._llm.chat( + messages, + session_id=self._session_id, + ) + + card = self._parse_l1_card(response.content) + return L1Node( + id=l1_id, + card=card, + time_range=l1_range, + children=l2_children, + ) + + # ------------------------------------------------------------------ + # 内部方法:JSON 解析(同步,纯 CPU) + # ------------------------------------------------------------------ + + @staticmethod + def _extract_json(raw: str) -> Any: + """从 VLM/LLM 原始输出中提取 JSON(处理 markdown 代码块包裹)。 + + 参数: + raw: 原始返回字符串。 + + 返回: + 解析后的 Python 对象(dict/list),解析失败返回 None。 + """ + raw = raw.strip() + # Phase 1: 尝试提取 markdown 代码块中的 JSON + code_match = re.search( + r"```(?:json)?\s*([\[{].*?[\]}])\s*```", + raw, + re.DOTALL, + ) + if code_match: + raw = code_match.group(1) + + # Phase 2: 直接解析 + try: + return json.loads(raw) + except json.JSONDecodeError: + pass + + # Phase 3: 尝试提取裸 JSON 对象/数组 + json_match = re.search(r"[\[{].*[\]}]", raw, re.DOTALL) + if json_match: + try: + return json.loads(json_match.group()) + except json.JSONDecodeError: + pass + + return None + + def _parse_l2_card(self, raw: str) -> L2Card: + """解析 VLM 输出为 L2Card。解析失败时创建退化卡片。 + + 参数: + raw: VLM 原始返回字符串。 + + 返回: + L2Card 实例。 + """ + data = self._extract_json(raw) + if isinstance(data, dict): + try: + state_changes = data.get("state_changes") + if state_changes is not None: + state_changes = str(state_changes) + return L2Card( + event_description=str(data["event_description"]), + entities=list(data["entities"]), + actions=list(data["actions"]), + action_subjects=list(data["action_subjects"]), + visible_text=list(data["visible_text"]), + spatial_relations=str(data["spatial_relations"]), + state_changes=state_changes, + ) + except (KeyError, TypeError, ValueError): + pass + + logger.warning( + "L2 VLM 输出 JSON 解析失败,使用退化卡片", + raw_preview=raw[:200], + ) + return L2Card( + event_description=raw.strip(), + entities=[], + actions=[], + action_subjects=[], + visible_text=[], + spatial_relations="", + state_changes=None, + ) + + def _parse_l3_cards( + self, + raw: str, + expected_n: int, + ) -> list[L3Card] | None: + """解析 VLM 输出为 L3Card 列表(保真算法 #2)。 + + 解析失败或数量不匹配时返回 None(触发逐帧 fallback)。 + 字段级校验:任何必填字段缺失或类型错误,整批次返回 None。 + + 参数: + raw: VLM 原始返回字符串。 + expected_n: 期望的卡片数量。 + + 返回: + 成功时返回 L3Card 列表,失败时返回 None。 + """ + data = self._extract_json(raw) + if not isinstance(data, list) or len(data) != expected_n: + return None + + cards: list[L3Card] = [] + for item in data: + if not isinstance(item, dict): + return None + try: + cards.append( + L3Card( + frame_summary=str(item["frame_summary"]), + visible_entities=list(item["visible_entities"]), + ongoing_actions=list(item["ongoing_actions"]), + visible_text=list(item["visible_text"]), + spatial_layout=str(item["spatial_layout"]), + visual_attributes=dict(item["visual_attributes"]), + ) + ) + except (KeyError, TypeError, ValueError): + return None + + return cards + + def _parse_l3_card_single(self, raw: str) -> L3Card: + """解析单帧 VLM 输出为 L3Card。解析失败时创建退化卡片。 + + 参数: + raw: VLM 原始返回字符串。 + + 返回: + L3Card 实例。 + """ + data = self._extract_json(raw) + if isinstance(data, dict): + try: + return L3Card( + frame_summary=str(data["frame_summary"]), + visible_entities=list(data["visible_entities"]), + ongoing_actions=list(data["ongoing_actions"]), + visible_text=list(data["visible_text"]), + spatial_layout=str(data["spatial_layout"]), + visual_attributes=dict(data["visual_attributes"]), + ) + except (KeyError, TypeError, ValueError): + pass + + logger.warning( + "L3 单帧 VLM 输出 JSON 解析失败,使用退化卡片", + raw_preview=raw[:200], + ) + return L3Card( + frame_summary=raw.strip(), + visible_entities=[], + ongoing_actions=[], + visible_text=[], + spatial_layout="", + visual_attributes={}, + ) + + def _parse_l1_card(self, raw: str) -> L1Card: + """解析 LLM 输出为 L1Card。解析失败时创建退化卡片。 + + 参数: + raw: LLM 原始返回字符串。 + + 返回: + L1Card 实例。 + """ + data = self._extract_json(raw) + if isinstance(data, dict): + try: + return L1Card( + scene_summary=str(data["scene_summary"]), + main_setting=str(data["main_setting"]), + key_entities=list(data["key_entities"]), + main_actions=list(data["main_actions"]), + topic_keywords=list(data["topic_keywords"]), + visible_text=list(data["visible_text"]), + temporal_flow=str(data["temporal_flow"]), + ) + except (KeyError, TypeError, ValueError): + pass + + logger.warning( + "L1 LLM 输出 JSON 解析失败,使用退化卡片", + raw_preview=raw[:200], + ) + return L1Card( + scene_summary=raw.strip(), + main_setting="", + key_entities=[], + main_actions=[], + topic_keywords=[], + visible_text=[], + temporal_flow="", + ) diff --git a/tests/unit/test_video_builder.py b/tests/unit/test_video_builder.py new file mode 100644 index 0000000..5a9addf --- /dev/null +++ b/tests/unit/test_video_builder.py @@ -0,0 +1,999 @@ +"""VideoTreeBuilder 单元测试。 + +测试覆盖: +- _segment_video: 时间切分 +- _get_l2_clips: 片段切分 +- _sample_representative_frames: 帧采样 +- _parse_l3_cards: L3 JSON 解析 + fallback 条件 +- _parse_l2_card: L2 JSON 解析 +- _parse_l1_card: L1 JSON 解析 +- _parse_l3_card_single: 单帧 L3 JSON 解析 +- build: 完整构建流程(mock VLM/LLM) +- checkpoint/resume: 断点续跑机制 +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest + +from app.tree.config import TreeConfig +from app.tree.index import ( + L1Card, + L1Node, + L2Card, + L2Node, + L3Card, + L3Node, +) +from app.tree.video_builder import VideoTreeBuilder +from core.types import LLMResponse + +# --------------------------------------------------------------------------- +# Mock 提供者 +# --------------------------------------------------------------------------- + + +def _make_llm_response(content: str) -> LLMResponse: + """构造 LLMResponse 测试工具。""" + return LLMResponse( + content=content, + thinking="", + model="mock", + provider="mock", + prompt_tokens=10, + completion_tokens=10, + latency_ms=10, + ttft_ms=1.0, + max_inter_token_ms=1.0, + cache_hit=False, + call_id="mock-call", + ) + + +def _l3_card_dict(idx: int = 0) -> dict[str, Any]: + """构造单个 L3Card 的字典表示。""" + return { + "frame_summary": f"帧 {idx} 的描述", + "visible_entities": ["实体A"], + "ongoing_actions": ["动作A"], + "visible_text": [], + "spatial_layout": "居中", + "visual_attributes": { + "lighting": "自然光", + "dominant_colors": ["白"], + "camera_angle": "正面", + }, + } + + +def _l2_card_dict() -> dict[str, Any]: + """构造 L2Card 的字典表示。""" + return { + "event_description": "视频片段描述", + "entities": ["实体A"], + "actions": ["动作A"], + "action_subjects": ["主体A"], + "visible_text": [], + "spatial_relations": "居中", + "state_changes": None, + } + + +def _l1_card_dict() -> dict[str, Any]: + """构造 L1Card 的字典表示。""" + return { + "scene_summary": "场景摘要描述", + "main_setting": "室内", + "key_entities": ["实体A"], + "main_actions": ["动作A"], + "topic_keywords": ["关键词A"], + "visible_text": [], + "temporal_flow": "从开始到结束", + } + + +class MockVLMProvider: + """模拟 VLM 提供者,根据 prompt 类型返回固定 JSON 响应。""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def chat_with_images( + self, + messages: list[dict[str, Any]], + images: list[str | Path], + *, + session_id: str | None = None, + parent_call_id: str | None = None, + ) -> LLMResponse: + """根据 prompt 内容判断调用类型,返回对应 JSON。""" + self.calls.append({"messages": messages, "images": images}) + content = messages[0]["content"] + n_images = len(images) + + if "JSON 数组" in content: + # L3 batch prompt + cards = [_l3_card_dict(i) for i in range(n_images)] + return _make_llm_response(json.dumps(cards, ensure_ascii=False)) + if "用一到两句话描述这帧" in content: + # L3 single prompt + return _make_llm_response( + json.dumps(_l3_card_dict(0), ensure_ascii=False), + ) + # L2 prompt + return _make_llm_response( + json.dumps(_l2_card_dict(), ensure_ascii=False), + ) + + +class MockLLMProvider: + """模拟 LLM 提供者,返回固定 L1Card JSON 响应。""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def chat( + self, + messages: list[dict[str, Any]], + *, + session_id: str | None = None, + parent_call_id: str | None = None, + ) -> LLMResponse: + """返回 L1Card JSON。""" + self.calls.append({"messages": messages}) + return _make_llm_response( + json.dumps(_l1_card_dict(), ensure_ascii=False), + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def tree_config(tmp_path: Path) -> TreeConfig: + """简化配置:10 秒视频→1 个 L1 段→2 个 L2 clip→每 clip 5 帧。""" + return TreeConfig( + l1_segment_duration=10.0, + l2_clip_duration=5.0, + l3_fps=1.0, + l2_representative_frames=2, + cache_dir=str(tmp_path / "cache"), + concurrency=4, + ) + + +@pytest.fixture() +def mock_vlm() -> MockVLMProvider: + """返回 MockVLMProvider 实例。""" + return MockVLMProvider() + + +@pytest.fixture() +def mock_llm() -> MockLLMProvider: + """返回 MockLLMProvider 实例。""" + return MockLLMProvider() + + +@pytest.fixture() +def builder( + mock_vlm: MockVLMProvider, + mock_llm: MockLLMProvider, + tree_config: TreeConfig, +) -> VideoTreeBuilder: + """构造带 mock 依赖的 VideoTreeBuilder。""" + return VideoTreeBuilder(vlm=mock_vlm, llm=mock_llm, config=tree_config) + + +# --------------------------------------------------------------------------- +# 测试:_segment_video +# --------------------------------------------------------------------------- + + +class TestSegmentVideo: + """测试时间切分逻辑。""" + + def test_exact_division(self, builder: VideoTreeBuilder) -> None: + """总时长能被 L1 段时长整除时的切分。""" + ranges = builder._segment_video("dummy", duration_hint=20.0) + assert ranges == [(0.0, 10.0), (10.0, 20.0)] + + def test_non_divisible(self, builder: VideoTreeBuilder) -> None: + """总时长不能被整除时,末段截断。""" + ranges = builder._segment_video("dummy", duration_hint=15.0) + assert len(ranges) == 2 + assert ranges[0] == (0.0, 10.0) + assert ranges[1] == (10.0, 15.0) + + def test_short_video(self, builder: VideoTreeBuilder) -> None: + """短视频(时长 < L1 段时长)产生单段。""" + ranges = builder._segment_video("dummy", duration_hint=5.0) + assert ranges == [(0.0, 5.0)] + + +# --------------------------------------------------------------------------- +# 测试:_get_l2_clips +# --------------------------------------------------------------------------- + + +class TestGetL2Clips: + """测试 L2 clip 切分逻辑。""" + + def test_basic_clips(self, builder: VideoTreeBuilder) -> None: + """L1 区间能被 L2 步长整除。""" + clips = builder._get_l2_clips((0.0, 10.0)) + assert clips == [(0.0, 5.0), (5.0, 10.0)] + + def test_remainder_clip(self, builder: VideoTreeBuilder) -> None: + """L1 区间不能被整除时,末段截断。""" + clips = builder._get_l2_clips((0.0, 7.0)) + assert len(clips) == 2 + assert clips[0] == (0.0, 5.0) + assert clips[1] == (5.0, 7.0) + + +# --------------------------------------------------------------------------- +# 测试:_sample_representative_frames +# --------------------------------------------------------------------------- + + +class TestSampleRepresentativeFrames: + """测试帧采样逻辑。""" + + def test_fewer_frames_than_n(self) -> None: + """帧数不足时返回全部。""" + frames = [("a.jpg", 0.0), ("b.jpg", 1.0)] + result = VideoTreeBuilder._sample_representative_frames(frames, 5) + assert result == ["a.jpg", "b.jpg"] + + def test_exact_n(self) -> None: + """帧数恰等于 n 时返回全部。""" + frames = [("a.jpg", 0.0), ("b.jpg", 1.0), ("c.jpg", 2.0)] + result = VideoTreeBuilder._sample_representative_frames(frames, 3) + assert result == ["a.jpg", "b.jpg", "c.jpg"] + + def test_uniform_sampling(self) -> None: + """10 帧中采样 3 帧,应均匀分布。""" + frames = [(f"{i}.jpg", float(i)) for i in range(10)] + result = VideoTreeBuilder._sample_representative_frames(frames, 3) + # step = 10/3 = 3.33 → indices 0, 3, 6 + assert result == ["0.jpg", "3.jpg", "6.jpg"] + + def test_sampling_two_from_five(self) -> None: + """5 帧中采样 2 帧。""" + frames = [(f"{i}.jpg", float(i)) for i in range(5)] + result = VideoTreeBuilder._sample_representative_frames(frames, 2) + # step = 5/2 = 2.5 → indices 0, 2 + assert result == ["0.jpg", "2.jpg"] + + +# --------------------------------------------------------------------------- +# 测试:_parse_l3_cards +# --------------------------------------------------------------------------- + + +class TestParseL3Cards: + """测试 L3 批量 JSON 解析(保真算法 #2 的解析环节)。""" + + def test_valid_json_array(self, builder: VideoTreeBuilder) -> None: + """正常 JSON 数组解析成功。""" + raw = json.dumps([_l3_card_dict(i) for i in range(3)]) + result = builder._parse_l3_cards(raw, 3) + assert result is not None + assert len(result) == 3 + assert result[0].frame_summary == "帧 0 的描述" + assert result[2].frame_summary == "帧 2 的描述" + + def test_count_mismatch_returns_none( + self, + builder: VideoTreeBuilder, + ) -> None: + """数量不匹配 → 返回 None(触发 fallback)。""" + raw = json.dumps([_l3_card_dict(0), _l3_card_dict(1)]) + result = builder._parse_l3_cards(raw, 3) + assert result is None + + def test_missing_field_returns_none( + self, + builder: VideoTreeBuilder, + ) -> None: + """必填字段缺失 → 整批次返回 None。""" + card = _l3_card_dict(0) + del card["frame_summary"] + raw = json.dumps([card]) + result = builder._parse_l3_cards(raw, 1) + assert result is None + + def test_invalid_json_returns_none( + self, + builder: VideoTreeBuilder, + ) -> None: + """非法 JSON 返回 None。""" + result = builder._parse_l3_cards("not json at all", 1) + assert result is None + + def test_json_in_code_block(self, builder: VideoTreeBuilder) -> None: + """Markdown 代码块包裹的 JSON 也能解析。""" + inner = json.dumps([_l3_card_dict(0)]) + raw = f"```json\n{inner}\n```" + result = builder._parse_l3_cards(raw, 1) + assert result is not None + assert len(result) == 1 + + def test_non_dict_item_returns_none( + self, + builder: VideoTreeBuilder, + ) -> None: + """数组元素非 dict → 返回 None。""" + raw = json.dumps(["string_item"]) + result = builder._parse_l3_cards(raw, 1) + assert result is None + + +# --------------------------------------------------------------------------- +# 测试:_parse_l2_card +# --------------------------------------------------------------------------- + + +class TestParseL2Card: + """测试 L2 JSON 解析。""" + + def test_valid_json(self, builder: VideoTreeBuilder) -> None: + """正常 JSON 解析成功。""" + raw = json.dumps(_l2_card_dict(), ensure_ascii=False) + card = builder._parse_l2_card(raw) + assert card.event_description == "视频片段描述" + assert card.entities == ["实体A"] + assert card.state_changes is None + + def test_invalid_json_fallback( + self, + builder: VideoTreeBuilder, + ) -> None: + """JSON 解析失败 → 退化卡片。""" + card = builder._parse_l2_card("这是一段普通文字描述") + assert card.event_description == "这是一段普通文字描述" + assert card.entities == [] + + def test_with_state_changes(self, builder: VideoTreeBuilder) -> None: + """state_changes 非 null 时正常解析。""" + d = _l2_card_dict() + d["state_changes"] = "从站立到坐下" + raw = json.dumps(d, ensure_ascii=False) + card = builder._parse_l2_card(raw) + assert card.state_changes == "从站立到坐下" + + +# --------------------------------------------------------------------------- +# 测试:_parse_l1_card +# --------------------------------------------------------------------------- + + +class TestParseL1Card: + """测试 L1 JSON 解析。""" + + def test_valid_json(self, builder: VideoTreeBuilder) -> None: + """正常 JSON 解析成功。""" + raw = json.dumps(_l1_card_dict(), ensure_ascii=False) + card = builder._parse_l1_card(raw) + assert card.scene_summary == "场景摘要描述" + assert card.main_setting == "室内" + + def test_invalid_json_fallback( + self, + builder: VideoTreeBuilder, + ) -> None: + """JSON 解析失败 → 退化卡片。""" + card = builder._parse_l1_card("这是段落摘要") + assert card.scene_summary == "这是段落摘要" + assert card.key_entities == [] + + +# --------------------------------------------------------------------------- +# 测试:_parse_l3_card_single +# --------------------------------------------------------------------------- + + +class TestParseL3CardSingle: + """测试单帧 L3 JSON 解析。""" + + def test_valid_json(self, builder: VideoTreeBuilder) -> None: + """正常 JSON 解析成功。""" + raw = json.dumps(_l3_card_dict(0), ensure_ascii=False) + card = builder._parse_l3_card_single(raw) + assert card.frame_summary == "帧 0 的描述" + + def test_invalid_json_fallback( + self, + builder: VideoTreeBuilder, + ) -> None: + """JSON 解析失败 → 退化卡片。""" + card = builder._parse_l3_card_single("一帧画面的描述") + assert card.frame_summary == "一帧画面的描述" + assert card.visible_entities == [] + + +# --------------------------------------------------------------------------- +# 测试:_extract_json +# --------------------------------------------------------------------------- + + +class TestExtractJson: + """测试 JSON 提取辅助方法。""" + + def test_plain_object(self) -> None: + """直接 JSON 对象。""" + result = VideoTreeBuilder._extract_json('{"key": "value"}') + assert result == {"key": "value"} + + def test_plain_array(self) -> None: + """直接 JSON 数组。""" + result = VideoTreeBuilder._extract_json("[1, 2, 3]") + assert result == [1, 2, 3] + + def test_code_block(self) -> None: + """Markdown 代码块中的 JSON。""" + raw = '```json\n{"key": "value"}\n```' + result = VideoTreeBuilder._extract_json(raw) + assert result == {"key": "value"} + + def test_surrounding_text(self) -> None: + """JSON 前后有文字。""" + raw = 'Here is the result: {"key": "value"} end' + result = VideoTreeBuilder._extract_json(raw) + assert result == {"key": "value"} + + def test_invalid(self) -> None: + """完全无 JSON 内容。""" + result = VideoTreeBuilder._extract_json("no json here") + assert result is None + + +# --------------------------------------------------------------------------- +# 测试:完整构建流程 +# --------------------------------------------------------------------------- + + +def _mock_ffmpeg_factory(tmp_path: Path): + """创建 mock ffmpeg 帧提取函数。""" + + def _mock_extract( + video_path: str, + ts: float, + out_path: str, + ) -> bool: + Path(out_path).parent.mkdir(parents=True, exist_ok=True) + Path(out_path).write_bytes(b"FAKE_JPEG") + return True + + return _mock_extract + + +class TestBuildFullFlow: + """测试完整构建流程(mock VLM/LLM/ffmpeg)。""" + + def test_build_produces_correct_structure( + self, + mock_vlm: MockVLMProvider, + mock_llm: MockLLMProvider, + tmp_path: Path, + ) -> None: + """10 秒视频 → 1 L1 → 2 L2 → 每 L2 约 5 帧 L3。""" + config = TreeConfig( + l1_segment_duration=10.0, + l2_clip_duration=5.0, + l3_fps=1.0, + l2_representative_frames=2, + cache_dir=str(tmp_path / "cache"), + concurrency=4, + ) + builder = VideoTreeBuilder( + vlm=mock_vlm, + llm=mock_llm, + config=config, + ) + + dummy_video = tmp_path / "test_video.mp4" + dummy_video.write_bytes(b"FAKE") + + with ( + patch.object( + builder, + "_segment_video", + return_value=[(0.0, 10.0)], + ), + patch.object( + builder, + "_ffmpeg_extract_frame", + side_effect=_mock_ffmpeg_factory(tmp_path), + ), + ): + index = builder.build(str(dummy_video)) + + # 结构校验 + assert len(index.roots) == 1 + l1 = index.roots[0] + assert l1.id == "l1_0" + assert l1.card.scene_summary == "场景摘要描述" + assert l1.time_range == (0.0, 10.0) + + # 2 个 L2 clip + assert len(l1.children) == 2 + for j, l2 in enumerate(l1.children): + assert l2.id == f"l1_0_l2_{j}" + assert l2.card.event_description == "视频片段描述" + # 5 帧/clip(5 秒 * 1 fps) + assert len(l2.children) == 5 + for k, l3 in enumerate(l2.children): + assert l3.id == f"l1_0_l2_{j}_l3_{k}" + assert l3.card.frame_summary is not None + assert l3.frame_path is not None + assert l3.timestamp is not None + + # 确认 VLM/LLM 被调用 + # L2: 2 calls (one per clip) + # L3: 2 calls (one batch per clip, each batch has 5 frames) + # L1: 1 call + assert len(mock_vlm.calls) == 4 # 2 L2 + 2 L3 batches + assert len(mock_llm.calls) == 1 # 1 L1 + + # metadata 校验 + assert index.metadata.source_path == str(dummy_video) + assert index.metadata.modality == "video" + + def test_build_cleans_up_intermediate( + self, + mock_vlm: MockVLMProvider, + mock_llm: MockLLMProvider, + tmp_path: Path, + ) -> None: + """构建成功后中间文件已清理。""" + config = TreeConfig( + l1_segment_duration=10.0, + l2_clip_duration=10.0, + l3_fps=1.0, + l2_representative_frames=2, + cache_dir=str(tmp_path / "cache"), + concurrency=4, + ) + builder = VideoTreeBuilder( + vlm=mock_vlm, + llm=mock_llm, + config=config, + ) + + dummy_video = tmp_path / "test_video.mp4" + dummy_video.write_bytes(b"FAKE") + + with ( + patch.object( + builder, + "_segment_video", + return_value=[(0.0, 10.0)], + ), + patch.object( + builder, + "_ffmpeg_extract_frame", + side_effect=_mock_ffmpeg_factory(tmp_path), + ), + ): + builder.build(str(dummy_video)) + + # 中间文件应已清理 + progress_dir = tmp_path / "cache" / "progress" + inter_dir = tmp_path / "cache" / "intermediate" / "test_video" + assert not (progress_dir / "test_video.json").exists() + # intermediate 目录可能不存在或为空 + if inter_dir.exists(): + assert len(list(inter_dir.glob("l1_*.json"))) == 0 + + +# --------------------------------------------------------------------------- +# 测试:L3 fallback(保真算法 #2) +# --------------------------------------------------------------------------- + + +class MockVLMWithBatchFailure: + """模拟批量 VLM 调用失败、单帧调用成功的 VLM 提供者。""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def chat_with_images( + self, + messages: list[dict[str, Any]], + images: list[str | Path], + *, + session_id: str | None = None, + parent_call_id: str | None = None, + ) -> LLMResponse: + """batch 返回无效 JSON,single 返回有效 JSON,L2 正常。""" + self.calls.append({"n_images": len(images)}) + content = messages[0]["content"] + + if "JSON 数组" in content: + # L3 batch: 返回无效 JSON 触发 fallback + return _make_llm_response("INVALID JSON OUTPUT") + if "用一到两句话描述这帧" in content: + # L3 single fallback: 有效 JSON + return _make_llm_response( + json.dumps(_l3_card_dict(0), ensure_ascii=False), + ) + # L2: 有效 JSON + return _make_llm_response( + json.dumps(_l2_card_dict(), ensure_ascii=False), + ) + + +class TestL3Fallback: + """测试 L3 批量失败→逐帧 fallback(保真算法 #2)。""" + + def test_fallback_to_single_frame( + self, + mock_llm: MockLLMProvider, + tmp_path: Path, + ) -> None: + """批量 VLM 解析失败时,逐帧 fallback 仍能构建完整树。""" + vlm = MockVLMWithBatchFailure() + config = TreeConfig( + l1_segment_duration=10.0, + l2_clip_duration=10.0, + l3_fps=1.0, + l2_representative_frames=2, + cache_dir=str(tmp_path / "cache"), + concurrency=4, + ) + builder = VideoTreeBuilder(vlm=vlm, llm=mock_llm, config=config) + + dummy_video = tmp_path / "test_video.mp4" + dummy_video.write_bytes(b"FAKE") + + with ( + patch.object( + builder, + "_segment_video", + return_value=[(0.0, 10.0)], + ), + patch.object( + builder, + "_ffmpeg_extract_frame", + side_effect=_mock_ffmpeg_factory(tmp_path), + ), + ): + index = builder.build(str(dummy_video)) + + # 结构仍然完整 + assert len(index.roots) == 1 + l1 = index.roots[0] + assert len(l1.children) == 1 # 1 clip (10s clip) + l2 = l1.children[0] + + # 10 frames (10s * 1fps), all from single-frame fallback + assert len(l2.children) == 10 + for l3 in l2.children: + assert l3.card.frame_summary == "帧 0 的描述" + + # VLM 调用次数:1 L2 + 1 batch(fail) + 10 single = 12 + # 但 batch 分为 2 batches (10 frames / 5 per batch) + # 所以: 1 L2 + 2 batch(fail) + 10 single = 13 + assert len(vlm.calls) == 13 + + +# --------------------------------------------------------------------------- +# 测试:断点续跑(保真算法 #3) +# --------------------------------------------------------------------------- + + +class TestCheckpointResume: + """测试断点续跑机制。""" + + def test_save_and_load_progress( + self, + builder: VideoTreeBuilder, + ) -> None: + """进度文件保存和加载。""" + stem = "test_video" + builder._save_progress(stem, total_l1=3, finished_l1_ids={0, 1}) + progress = builder._load_progress(stem) + + assert progress is not None + assert progress["total_l1"] == 3 + assert sorted(progress["finished_l1_ids"]) == [0, 1] + assert "created_at" in progress + assert "updated_at" in progress + + def test_load_nonexistent_progress( + self, + builder: VideoTreeBuilder, + ) -> None: + """不存在的进度文件返回 None。""" + assert builder._load_progress("nonexistent") is None + + def test_save_and_load_l1_intermediate( + self, + builder: VideoTreeBuilder, + ) -> None: + """L1 中间结果保存和加载。""" + stem = "test_video" + l1_card = L1Card( + scene_summary="测试摘要", + main_setting="测试场景", + key_entities=["实体"], + main_actions=["动作"], + topic_keywords=["关键词"], + visible_text=[], + temporal_flow="测试流向", + ) + l2_card = L2Card( + event_description="事件描述", + entities=["实体"], + actions=["动作"], + action_subjects=["主体"], + visible_text=[], + spatial_relations="居中", + state_changes=None, + ) + l3_card = L3Card( + frame_summary="帧描述", + visible_entities=["实体"], + ongoing_actions=["动作"], + visible_text=[], + spatial_layout="居中", + visual_attributes={"lighting": "自然光"}, + ) + l3_node = L3Node( + id="l1_0_l2_0_l3_0", + card=l3_card, + timestamp=1.0, + frame_path="/tmp/frame.jpg", + ) + l2_node = L2Node( + id="l1_0_l2_0", + card=l2_card, + time_range=(0.0, 5.0), + children=[l3_node], + ) + l1_node = L1Node( + id="l1_0", + card=l1_card, + time_range=(0.0, 10.0), + children=[l2_node], + ) + + builder._save_l1_intermediate(stem, l1_node, 0) + assert builder._has_l1_intermediate(stem, 0) + assert not builder._has_l1_intermediate(stem, 1) + + loaded = builder._load_l1_intermediate(stem, 0) + assert loaded is not None + assert loaded.id == "l1_0" + assert loaded.card.scene_summary == "测试摘要" + assert len(loaded.children) == 1 + assert loaded.children[0].id == "l1_0_l2_0" + + def test_cleanup_removes_files( + self, + builder: VideoTreeBuilder, + ) -> None: + """清理函数删除进度文件和中间 JSON。""" + stem = "test_video" + builder._save_progress(stem, total_l1=1, finished_l1_ids={0}) + + # 创建一个假的中间文件 + inter_dir = builder._intermediate_dir(stem) + inter_dir.mkdir(parents=True, exist_ok=True) + (inter_dir / "l1_0.json").write_text("{}") + + builder._cleanup_intermediate_and_progress(stem) + + assert not builder._progress_path(stem).is_file() + assert not (inter_dir / "l1_0.json").is_file() + + def test_resume_skips_finished_segments( + self, + mock_vlm: MockVLMProvider, + mock_llm: MockLLMProvider, + tmp_path: Path, + ) -> None: + """断点续跑:跳过已完成的 L1 段,只构建未完成的段。""" + config = TreeConfig( + l1_segment_duration=5.0, + l2_clip_duration=5.0, + l3_fps=1.0, + l2_representative_frames=2, + cache_dir=str(tmp_path / "cache"), + concurrency=4, + ) + builder = VideoTreeBuilder( + vlm=mock_vlm, + llm=mock_llm, + config=config, + ) + + source_id = "test_resume" + + # Phase 1: 手动创建 L1_0 的中间结果(模拟已完成) + l1_card = L1Card( + scene_summary="已完成的段", + main_setting="场景A", + key_entities=["实体A"], + main_actions=["动作A"], + topic_keywords=["关键词A"], + visible_text=[], + temporal_flow="流向A", + ) + l2_card = L2Card( + event_description="已完成的片段", + entities=["实体A"], + actions=["动作A"], + action_subjects=["主体A"], + visible_text=[], + spatial_relations="居中", + state_changes=None, + ) + l3_card = L3Card( + frame_summary="已完成的帧", + visible_entities=["实体A"], + ongoing_actions=["动作A"], + visible_text=[], + spatial_layout="居中", + visual_attributes={"lighting": "自然光"}, + ) + l3_node = L3Node( + id="l1_0_l2_0_l3_0", + card=l3_card, + timestamp=1.0, + ) + l2_node = L2Node( + id="l1_0_l2_0", + card=l2_card, + time_range=(0.0, 5.0), + children=[l3_node], + ) + l1_node_0 = L1Node( + id="l1_0", + card=l1_card, + time_range=(0.0, 5.0), + children=[l2_node], + ) + builder._save_l1_intermediate(source_id, l1_node_0, 0) + + # Phase 2: 创建进度文件(标记 L1_0 完成) + builder._save_progress( + source_id, + total_l1=2, + finished_l1_ids={0}, + ) + + # Phase 3: 构建(应跳过 L1_0,只构建 L1_1) + dummy_video = tmp_path / f"{source_id}.mp4" + dummy_video.write_bytes(b"FAKE") + + with ( + patch.object( + builder, + "_segment_video", + return_value=[(0.0, 5.0), (5.0, 10.0)], + ), + patch.object( + builder, + "_ffmpeg_extract_frame", + side_effect=_mock_ffmpeg_factory(tmp_path), + ), + ): + index = builder.build(str(dummy_video)) + + # Phase 4: 验证 + assert len(index.roots) == 2 + + # L1_0 来自中间结果 + assert index.roots[0].id == "l1_0" + assert index.roots[0].card.scene_summary == "已完成的段" + + # L1_1 是新构建的 + assert index.roots[1].id == "l1_1" + assert index.roots[1].card.scene_summary == "场景摘要描述" + + # VLM 只被调用了 L1_1 的部分(1 L2 + 1 L3 batch) + assert len(mock_vlm.calls) == 2 # 1 L2 + 1 L3 + assert len(mock_llm.calls) == 1 # 1 L1 + + # 构建完成后,进度和中间文件已清理 + assert not builder._progress_path(source_id).is_file() + + +# --------------------------------------------------------------------------- +# 测试:字幕注入 +# --------------------------------------------------------------------------- + + +class TestSubtitleInjection: + """测试字幕注入功能。""" + + def test_build_subtitle_block_with_entries( + self, + builder: VideoTreeBuilder, + ) -> None: + """有匹配字幕时返回字幕文本块。""" + from app.tree.subtitle import SRTEntry + + entries = [ + SRTEntry(start=1.0, end=3.0, text="你好世界"), + SRTEntry(start=4.0, end=6.0, text="再见"), + ] + block = builder._build_subtitle_block(entries, (0.0, 5.0)) + assert "字幕信息" in block + assert "你好世界" in block + + def test_build_subtitle_block_no_match( + self, + builder: VideoTreeBuilder, + ) -> None: + """无匹配字幕时返回空字符串。""" + from app.tree.subtitle import SRTEntry + + entries = [SRTEntry(start=100.0, end=110.0, text="远处的字幕")] + block = builder._build_subtitle_block(entries, (0.0, 5.0)) + assert block == "" + + def test_build_subtitle_block_none_entries( + self, + builder: VideoTreeBuilder, + ) -> None: + """srt_entries 为 None 时返回空字符串。""" + block = builder._build_subtitle_block(None, (0.0, 5.0)) + assert block == "" + + def test_build_subtitle_block_point_range( + self, + builder: VideoTreeBuilder, + ) -> None: + """点时间范围(单帧)自动扩展窗口。""" + from app.tree.subtitle import SRTEntry + + entries = [SRTEntry(start=4.0, end=6.0, text="窗口内字幕")] + # 点时间 5.0,窗口 ±5.0 → (0.0, 10.0) + block = builder._build_subtitle_block(entries, (5.0, 5.0)) + assert "窗口内字幕" in block + + +# --------------------------------------------------------------------------- +# 测试:URL 与 stem 辅助 +# --------------------------------------------------------------------------- + + +class TestHelpers: + """测试静态辅助方法。""" + + def test_is_url_true(self) -> None: + """HTTP/HTTPS URL 识别。""" + assert VideoTreeBuilder._is_url("https://example.com/video.mp4") + assert VideoTreeBuilder._is_url("http://example.com/video.mp4") + + def test_is_url_false(self) -> None: + """本地路径不是 URL。""" + assert not VideoTreeBuilder._is_url("/path/to/video.mp4") + assert not VideoTreeBuilder._is_url("video.mp4") + + def test_source_stem_local(self) -> None: + """本地文件的 stem。""" + assert VideoTreeBuilder._source_stem("/path/to/my_video.mp4") == "my_video" + + def test_source_stem_youtube(self) -> None: + """YouTube URL 的 stem 为视频 ID。""" + stem = VideoTreeBuilder._source_stem( + "https://www.youtube.com/watch?v=dQw4w9WgXcQ", + ) + assert stem == "dQw4w9WgXcQ" + + def test_source_stem_long_name(self) -> None: + """超长文件名截断到 64 字符。""" + long_name = "a" * 100 + ".mp4" + stem = VideoTreeBuilder._source_stem(f"/path/{long_name}") + assert len(stem) == 64