From 22ad01497321744c9e4d6e60c2b2c57c2ccec1d8 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Tue, 7 Jul 2026 01:35:10 -0400 Subject: [PATCH] =?UTF-8?q?feat(tree):=20TreeIndex=20=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E7=BB=93=E6=9E=84=20=E2=80=94=20Card=20=E4=BD=93=E7=B3=BB=20+?= =?UTF-8?q?=20=E8=8A=82=E7=82=B9=20+=20=E5=BA=8F=E5=88=97=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增三级 frozen Card dataclass: L3Card(6字段), L2Card(7字段), L1Card(7字段) - 节点重构: L3Node/L2Node/L1Node 使用 Card 替代原始字符串字段 - 添加 @property 兼容层: description/summary 代理到 Card 字段 - L3Node 新增 subtitle 字段（字幕集成预留） - JSON 序列化/反序列化支持 Card 结构 + embedding base64 编解码 - load_json 新增 ID 唯一性校验（重复 ID 抛 ValueError） - 移除 pickle 序列化（仅保留 JSON） - 日志从 log_msg 迁移到 loguru - 17 个单元测试全部通过 Co-Authored-By: Claude Opus 4.6 (1M context) --- app/tree/index.py | 746 ++++++++++++++++++++++++++++++++++ tests/unit/test_tree_index.py | 227 +++++++++++ 2 files changed, 973 insertions(+) create mode 100644 app/tree/index.py create mode 100644 tests/unit/test_tree_index.py diff --git a/app/tree/index.py b/app/tree/index.py new file mode 100644 index 0000000..7e0e8e2 --- /dev/null +++ b/app/tree/index.py @@ -0,0 +1,746 @@ +"""三层树索引核心数据结构。 + +定义 Video-Tree-TRM 的三层树状索引结构，是所有后续模块 +（builder、retriever、harness、search）的基础依赖。 + +数据结构层次:: + + TreeIndex + └─ List[L1Node] 全局叙事节点 + └─ List[L2Node] 片段级语义节点 + └─ List[L3Node] 帧/细节级节点 + +与参考项目 (TRM4) 的关键区别: + - Card 体系：每层节点的描述信息封装为 frozen dataclass（L1Card/L2Card/L3Card）， + 字段来自 VLM 结构化输出，保证不可变。 + - 序列化方式：仅保留 JSON（移除 pickle）。 + - 统一嵌入空间：所有 embedding 均来自 text_embed()，无跨模态问题。 +""" + +from __future__ import annotations + +import base64 +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import TYPE_CHECKING, Any + +import numpy as np +from loguru import logger + +if TYPE_CHECKING: + from collections.abc import Callable + +# --------------------------------------------------------------------------- +# Embedding 序列化辅助函数 +# --------------------------------------------------------------------------- + + +def _embed_to_str(arr: np.ndarray | None) -> str | None: + """float32 ndarray -> base64 字符串（用于 JSON 序列化）。 + + 参数: + arr: float32 数组，形状任意。 + + 返回: + base64 编码字符串，或 None（输入为 None 时）。 + """ + if arr is None: + return None + return base64.b64encode(arr.astype(np.float32).tobytes()).decode() + + +def _embed_from_str(s: str | None) -> np.ndarray | None: + """base64 字符串 -> float32 ndarray（用于 JSON 反序列化）。 + + 参数: + s: base64 编码字符串。 + + 返回: + float32 数组，或 None（输入为 None/空时）。 + """ + if s is None or s == "": + return None + return np.frombuffer(base64.b64decode(s), dtype=np.float32) + + +# --------------------------------------------------------------------------- +# Card 数据结构（frozen，来自 VLM 结构化输出） +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class L3Card: + """L3 帧级语义卡片（不可变）。 + + 封装 VLM 对单帧的结构化描述输出。 + + 属性: + frame_summary: 帧内容摘要。 + visible_entities: 可见实体列表。 + ongoing_actions: 正在进行的动作列表。 + visible_text: 画面中可见的文字列表。 + spatial_layout: 空间布局描述。 + visual_attributes: 视觉属性字典（如光照、色调等）。 + """ + + frame_summary: str + visible_entities: list[str] + ongoing_actions: list[str] + visible_text: list[str] + spatial_layout: str + visual_attributes: dict[str, Any] + + +@dataclass(frozen=True) +class L2Card: + """L2 事件级语义卡片（不可变）。 + + 封装 VLM 对一个事件片段的结构化描述输出。 + + 属性: + event_description: 事件描述。 + entities: 参与实体列表。 + actions: 动作列表。 + action_subjects: 动作主体列表。 + visible_text: 片段中可见的文字列表。 + spatial_relations: 空间关系描述。 + state_changes: 状态变化描述（可选）。 + """ + + event_description: str + entities: list[str] + actions: list[str] + action_subjects: list[str] + visible_text: list[str] + spatial_relations: str + state_changes: str | None + + +@dataclass(frozen=True) +class L1Card: + """L1 场景级语义卡片（不可变）。 + + 封装 VLM 对一个完整场景的结构化描述输出。 + + 属性: + scene_summary: 场景摘要。 + main_setting: 主要场景设定（如"室内"、"户外"等）。 + key_entities: 关键实体列表。 + main_actions: 主要动作列表。 + topic_keywords: 主题关键词列表。 + visible_text: 场景中可见的文字列表。 + temporal_flow: 时间流描述。 + """ + + scene_summary: str + main_setting: str + key_entities: list[str] + main_actions: list[str] + topic_keywords: list[str] + visible_text: list[str] + temporal_flow: str + + +# --------------------------------------------------------------------------- +# 元数据 +# --------------------------------------------------------------------------- + + +@dataclass +class IndexMeta: + """树索引元数据。 + + 属性: + source_path: 原始数据路径（视频文件或文本文件）。 + modality: 数据模态，"text" 或 "video"。 + embed_model: 嵌入模型名称（建树时为 None，embed_all 后填充）。 + embed_dim: 嵌入向量维度（建树时为 None，embed_all 后填充）。 + created_at: 创建时间（ISO 格式字符串）。 + """ + + source_path: str + modality: str + embed_model: str | None = None + embed_dim: int | None = None + created_at: str = field(default_factory=lambda: datetime.now().isoformat()) + + +# --------------------------------------------------------------------------- +# 节点数据结构 +# --------------------------------------------------------------------------- + + +@dataclass +class L3Node: + """L3 帧/细节级节点（叶子层）。 + + 代表最细粒度的语义单元，对应一个具体的帧描述。 + + 属性: + id: 节点唯一标识。 + card: 帧级语义卡片（VLM 结构化输出）。 + embedding: 文本嵌入向量，形状 [D]，float32。 + timestamp: 对应的时间戳（秒，可选）。 + frame_path: 关联的帧图像路径（可选，仅视频模态）。 + subtitle: 该帧对应的字幕文本（可选）。 + """ + + id: str + card: L3Card + embedding: np.ndarray | None = None + timestamp: float | None = None + frame_path: str | None = None + subtitle: str | None = None + + @property + def description(self) -> str: + """帧描述文本（取自 card.frame_summary）。""" + return self.card.frame_summary + + +@dataclass +class L2Node: + """L2 片段级语义节点（中间层）。 + + 连接 L1 宏观叙事与 L3 细节描述。 + + 属性: + id: 节点唯一标识。 + card: 事件级语义卡片（VLM 结构化输出）。 + embedding: 文本嵌入向量，形状 [D]，float32。 + time_range: 时间范围 (start, end)（秒，可选）。 + children: 所属的 L3 子节点列表。 + """ + + id: str + card: L2Card + embedding: np.ndarray | None = None + time_range: tuple[float, float] | None = None + children: list[L3Node] = field(default_factory=list) + + @property + def description(self) -> str: + """事件描述文本（取自 card.event_description）。""" + return self.card.event_description + + +@dataclass +class L1Node: + """L1 全局叙事节点（根层）。 + + 代表最粗粒度的语义单元，包含宏观场景摘要。 + + 属性: + id: 节点唯一标识。 + card: 场景级语义卡片（VLM 结构化输出）。 + embedding: 文本嵌入向量，形状 [D]，float32。 + time_range: 时间范围 (start, end)（秒，可选）。 + children: 所属的 L2 子节点列表。 + """ + + id: str + card: L1Card + embedding: np.ndarray | None = None + time_range: tuple[float, float] | None = None + children: list[L2Node] = field(default_factory=list) + + @property + def summary(self) -> str: + """场景摘要文本（取自 card.scene_summary）。""" + return self.card.scene_summary + + # ------------------------------------------------------------------ + # JSON 辅助方法（单个 L1 段的轻量序列化） + # ------------------------------------------------------------------ + + def to_dict(self, include_embedding: bool = False) -> dict[str, Any]: + """将当前 L1 节点（及其全部 L2/L3 子树）序列化为纯 dict。 + + 参数: + include_embedding: 若 True，将 embedding 向量序列化为 base64 字符串。 + + 返回: + 包含 id/card/time_range/children 的字典，可选包含 embedding。 + """ + + def l3_to_dict(n: L3Node) -> dict[str, Any]: + d: dict[str, Any] = { + "id": n.id, + "card": { + "frame_summary": n.card.frame_summary, + "visible_entities": n.card.visible_entities, + "ongoing_actions": n.card.ongoing_actions, + "visible_text": n.card.visible_text, + "spatial_layout": n.card.spatial_layout, + "visual_attributes": n.card.visual_attributes, + }, + "timestamp": n.timestamp, + "frame_path": n.frame_path, + "subtitle": n.subtitle, + } + if include_embedding: + d["embedding"] = _embed_to_str(n.embedding) + return d + + def l2_to_dict(n: L2Node) -> dict[str, Any]: + d: dict[str, Any] = { + "id": n.id, + "card": { + "event_description": n.card.event_description, + "entities": n.card.entities, + "actions": n.card.actions, + "action_subjects": n.card.action_subjects, + "visible_text": n.card.visible_text, + "spatial_relations": n.card.spatial_relations, + "state_changes": n.card.state_changes, + }, + "time_range": list(n.time_range) if n.time_range else None, + "children": [l3_to_dict(c) for c in n.children], + } + if include_embedding: + d["embedding"] = _embed_to_str(n.embedding) + return d + + d: dict[str, Any] = { + "id": self.id, + "card": { + "scene_summary": self.card.scene_summary, + "main_setting": self.card.main_setting, + "key_entities": self.card.key_entities, + "main_actions": self.card.main_actions, + "topic_keywords": self.card.topic_keywords, + "visible_text": self.card.visible_text, + "temporal_flow": self.card.temporal_flow, + }, + "time_range": list(self.time_range) if self.time_range else None, + "children": [l2_to_dict(c) for c in self.children], + } + if include_embedding: + d["embedding"] = _embed_to_str(self.embedding) + return d + + @staticmethod + def from_dict(d: dict[str, Any]) -> L1Node: + """从 dict 反序列化单个 L1 节点（支持 embedding 恢复）。 + + 参数: + d: to_dict() 输出的字典，可包含 embedding 字段。 + + 返回: + L1Node 实例（embedding 自动从 base64 恢复，若无则为 None）。 + """ + l2_nodes: list[L2Node] = [] + for l2d in d.get("children", []): + l3_nodes: list[L3Node] = [] + for l3d in l2d.get("children", []): + l3_card = L3Card( + frame_summary=l3d["card"]["frame_summary"], + visible_entities=l3d["card"]["visible_entities"], + ongoing_actions=l3d["card"]["ongoing_actions"], + visible_text=l3d["card"]["visible_text"], + spatial_layout=l3d["card"]["spatial_layout"], + visual_attributes=l3d["card"]["visual_attributes"], + ) + l3_nodes.append( + L3Node( + id=l3d["id"], + card=l3_card, + embedding=_embed_from_str(l3d.get("embedding")), + timestamp=l3d.get("timestamp"), + frame_path=l3d.get("frame_path"), + subtitle=l3d.get("subtitle"), + ) + ) + l2_card = L2Card( + event_description=l2d["card"]["event_description"], + entities=l2d["card"]["entities"], + actions=l2d["card"]["actions"], + action_subjects=l2d["card"]["action_subjects"], + visible_text=l2d["card"]["visible_text"], + spatial_relations=l2d["card"]["spatial_relations"], + state_changes=l2d["card"]["state_changes"], + ) + tr2 = l2d.get("time_range") + l2_nodes.append( + L2Node( + id=l2d["id"], + card=l2_card, + embedding=_embed_from_str(l2d.get("embedding")), + time_range=tuple(tr2) if tr2 else None, + children=l3_nodes, + ) + ) + l1_card = L1Card( + scene_summary=d["card"]["scene_summary"], + main_setting=d["card"]["main_setting"], + key_entities=d["card"]["key_entities"], + main_actions=d["card"]["main_actions"], + topic_keywords=d["card"]["topic_keywords"], + visible_text=d["card"]["visible_text"], + temporal_flow=d["card"]["temporal_flow"], + ) + tr1 = d.get("time_range") + return L1Node( + id=d["id"], + card=l1_card, + embedding=_embed_from_str(d.get("embedding")), + time_range=tuple(tr1) if tr1 else None, + children=l2_nodes, + ) + + +# --------------------------------------------------------------------------- +# 树索引容器 +# --------------------------------------------------------------------------- + + +@dataclass +class TreeIndex: + """三层树索引容器。 + + 组织和管理三层节点结构，提供嵌入矩阵提取、节点访问、 + 以及 JSON 序列化/反序列化接口。 + + 典型工作流:: + + # 1. 构建索引 + index = TreeIndex(metadata=meta, roots=[l1_node_1, l1_node_2]) + + # 2. 批量 embed（首次检索前） + index.embed_all(embed_fn, "model-name", 768) + + # 3. 提取嵌入矩阵（用于检索） + M_L1 = index.l1_embeddings() + M_L2 = index.l2_embeddings_of(l1_idx=0) + M_L3 = index.l3_embeddings_of(0, 1) + + # 4. 序列化 + index.save_json("cache/my_index.json") + loaded = TreeIndex.load_json("cache/my_index.json") + + 属性: + metadata: 索引元数据。 + roots: L1 节点列表。 + """ + + metadata: IndexMeta + roots: list[L1Node] = field(default_factory=list) + + # ------------------------------------------------------------------ # + # 嵌入状态检查 + # ------------------------------------------------------------------ # + + @property + def is_embedded(self) -> bool: + """检查所有节点是否已填充嵌入向量。 + + 返回: + True 表示所有 L1/L2/L3 节点的 embedding 均非 None； + False 表示尚未 embed。 + """ + for l1 in self.roots: + if l1.embedding is None: + return False + for l2 in l1.children: + if l2.embedding is None: + return False + for l3 in l2.children: + if l3.embedding is None: + return False + return True + + # ------------------------------------------------------------------ # + # 批量嵌入 + # ------------------------------------------------------------------ # + + def embed_all( + self, + embed_fn: Callable[[str | list[str]], np.ndarray], + model_name: str, + embed_dim: int, + ) -> None: + """对所有节点批量执行 embedding，更新 metadata。 + + 建树阶段不调用此方法（embedding=None）。 + 首次检索前由 Pipeline 调用，结果缓存在节点上。 + + 参数: + embed_fn: EmbeddingModel.embed 方法，接受 str 或 List[str]， + 返回 [N, D] ndarray。 + model_name: 嵌入模型名称，写入 metadata。 + embed_dim: 嵌入维度，写入 metadata。 + + 实现细节: + - L3 节点按 L2 分组批量 embed（一次调用），减少 API 开销。 + - L1/L2 各单独 embed（数量少，不值得合并）。 + - 仅对 embedding 为 None 的节点执行（支持增量更新）。 + """ + assert len(self.roots) > 0, "embed_all: 树为空，无节点可 embed" + for l1 in self.roots: + if l1.embedding is None: + l1.embedding = embed_fn(l1.summary)[0].astype(np.float32) + for l2 in l1.children: + if l2.embedding is None: + l2.embedding = embed_fn(l2.description)[0].astype(np.float32) + # L3 批量 embed + need_embed = [l3 for l3 in l2.children if l3.embedding is None] + if need_embed: + texts = [l3.description for l3 in need_embed] + embs = embed_fn(texts).astype(np.float32) # [N, D] + for l3, emb in zip(need_embed, embs, strict=True): + l3.embedding = emb + self.metadata.embed_model = model_name + self.metadata.embed_dim = embed_dim + logger.info( + "embed_all 完成", + model=model_name, + embed_dim=embed_dim, + ) + + # ------------------------------------------------------------------ # + # 嵌入矩阵提取 + # ------------------------------------------------------------------ # + + def l1_embeddings(self) -> np.ndarray: + """返回所有 L1 节点的嵌入矩阵。 + + 返回: + 形状 [N1, D] 的 float32 矩阵。空树返回 [0, D]。 + + 异常: + AssertionError: 节点 embedding 尚未计算（请先调用 embed_all）。 + """ + assert self.is_embedded, "L1 embedding 尚未计算，请先调用 tree.embed_all()" + if not self.roots: + return np.zeros((0, self.metadata.embed_dim), dtype=np.float32) + return np.stack([r.embedding for r in self.roots], axis=0).astype(np.float32) + + def l2_embeddings_of(self, l1_idx: int) -> np.ndarray: + """返回指定 L1 节点下所有 L2 子节点的嵌入矩阵。 + + 参数: + l1_idx: L1 节点索引。 + + 返回: + 形状 [N2, D] 的 float32 矩阵。 + + 异常: + IndexError: l1_idx 越界。 + AssertionError: embedding 尚未计算。 + """ + assert self.is_embedded, "L2 embedding 尚未计算，请先调用 tree.embed_all()" + if not (0 <= l1_idx < len(self.roots)): + raise IndexError(f"l1_idx={l1_idx} 越界，L1 节点数={len(self.roots)}") + children = self.roots[l1_idx].children + if not children: + return np.zeros((0, self.metadata.embed_dim), dtype=np.float32) + return np.stack([c.embedding for c in children], axis=0).astype(np.float32) + + def l3_embeddings_of(self, l1_idx: int, l2_idx: int) -> np.ndarray: + """返回指定 L2 节点下所有 L3 子节点的嵌入矩阵。 + + 参数: + l1_idx: L1 节点索引。 + l2_idx: L2 节点索引（相对于 L1）。 + + 返回: + 形状 [N3, D] 的 float32 矩阵。 + + 异常: + IndexError: 索引越界。 + AssertionError: embedding 尚未计算。 + """ + assert self.is_embedded, "L3 embedding 尚未计算，请先调用 tree.embed_all()" + if not (0 <= l1_idx < len(self.roots)): + raise IndexError(f"l1_idx={l1_idx} 越界，L1 节点数={len(self.roots)}") + l2_children = self.roots[l1_idx].children + if not (0 <= l2_idx < len(l2_children)): + raise IndexError(f"l2_idx={l2_idx} 越界，L2 节点数={len(l2_children)}") + l3_children = l2_children[l2_idx].children + if not l3_children: + return np.zeros((0, self.metadata.embed_dim), dtype=np.float32) + return np.stack([c.embedding for c in l3_children], axis=0).astype(np.float32) + + # ------------------------------------------------------------------ # + # 节点访问 + # ------------------------------------------------------------------ # + + def get_node(self, l1: int, l2: int, l3: int) -> L3Node: + """按三级路径索引获取 L3 节点。 + + 参数: + l1: L1 节点索引。 + l2: L2 节点索引。 + l3: L3 节点索引。 + + 返回: + 目标 L3Node。 + + 异常: + IndexError: 任意层级索引越界。 + """ + if l1 < 0 or l1 >= len(self.roots): + raise IndexError(f"l1={l1} 越界，L1 节点数={len(self.roots)}") + l2_children = self.roots[l1].children + if l2 < 0 or l2 >= len(l2_children): + raise IndexError(f"l2={l2} 越界，L2 节点数={len(l2_children)}") + l3_children = l2_children[l2].children + if l3 < 0 or l3 >= len(l3_children): + raise IndexError(f"l3={l3} 越界，L3 节点数={len(l3_children)}") + return l3_children[l3] + + # ------------------------------------------------------------------ # + # JSON 序列化 + # ------------------------------------------------------------------ # + + def to_dict(self, include_embedding: bool = False) -> dict[str, Any]: + """将树索引序列化为纯 Python dict。 + + 参数: + include_embedding: 若 True，将所有节点的 embedding 向量序列化为 base64。 + + 返回: + 可直接 json.dump 的字典，结构为 {metadata, roots[...]}。 + """ + metadata_dict: dict[str, Any] = { + "source_path": self.metadata.source_path, + "modality": self.metadata.modality, + "created_at": self.metadata.created_at, + } + if include_embedding: + metadata_dict["embed_model"] = self.metadata.embed_model + metadata_dict["embed_dim"] = self.metadata.embed_dim + + return { + "metadata": metadata_dict, + "roots": [r.to_dict(include_embedding=include_embedding) for r in self.roots], + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> TreeIndex: + """从 dict 反序列化为 TreeIndex（支持 embedding 恢复）。 + + 参数: + d: to_dict() 的输出或等价结构，可包含 embedding 字段。 + + 返回: + TreeIndex 实例。 + + 异常: + ValueError: 存在重复的节点 ID。 + """ + meta = IndexMeta( + source_path=d["metadata"]["source_path"], + modality=d["metadata"]["modality"], + embed_model=d["metadata"].get("embed_model"), + embed_dim=d["metadata"].get("embed_dim"), + created_at=d["metadata"].get("created_at", datetime.now().isoformat()), + ) + + roots: list[L1Node] = [] + for r in d["roots"]: + roots.append(L1Node.from_dict(r)) + + return cls(metadata=meta, roots=roots) + + def _validate_id_uniqueness(self) -> None: + """校验树中所有节点 ID 的唯一性。 + + 异常: + ValueError: 存在重复的节点 ID。 + """ + seen: set[str] = set() + for l1 in self.roots: + if l1.id in seen: + raise ValueError(f"重复的节点 ID: {l1.id}") + seen.add(l1.id) + for l2 in l1.children: + if l2.id in seen: + raise ValueError(f"重复的节点 ID: {l2.id}") + seen.add(l2.id) + for l3 in l2.children: + if l3.id in seen: + raise ValueError(f"重复的节点 ID: {l3.id}") + seen.add(l3.id) + + def save_json(self, path: str, include_embedding: bool = False) -> None: + """将树索引以 JSON 格式保存到磁盘。 + + 参数: + path: 保存文件路径（推荐 .json 后缀）。 + include_embedding: 若 True，将所有节点的 embedding 向量保存到 JSON。 + """ + with open(path, "w", encoding="utf-8") as f: + json.dump( + self.to_dict(include_embedding=include_embedding), + f, + ensure_ascii=False, + indent=2, + ) + logger.info( + "树索引（JSON）已保存至 {}", + path, + n_l1=len(self.roots), + include_embedding=include_embedding, + ) + + @classmethod + def load_json(cls, path: str) -> TreeIndex: + """从 JSON 文件加载树索引（自动检测并恢复 embedding）。 + + 参数: + path: JSON 文件路径。 + + 返回: + TreeIndex 实例。若 JSON 中包含 embedding 字段，自动反序列化填充； + 否则 embedding=None（向后兼容旧格式）。 + + 异常: + FileNotFoundError: 文件不存在。 + ValueError: 存在重复的节点 ID。 + """ + with open(path, encoding="utf-8") as f: + d = json.load(f) + obj = cls.from_dict(d) + obj._validate_id_uniqueness() + logger.info( + "树索引（JSON）已从 {} 加载", + path, + n_l1=len(obj.roots), + is_embedded=obj.is_embedded, + ) + return obj + + +# --------------------------------------------------------------------------- +# 单 L1 段的轻量序列化（用于断点续跑） +# --------------------------------------------------------------------------- + + +def save_l1_json(path: str, l1_node: L1Node) -> None: + """将单个 L1 节点（及其子树）以 JSON 形式保存到磁盘。 + + 参数: + path: 目标文件路径。 + l1_node: 待序列化的 L1 节点。 + """ + with open(path, "w", encoding="utf-8") as f: + json.dump(l1_node.to_dict(), f, ensure_ascii=False, indent=2) + logger.info("L1 中间结果已保存", path=path, l1_id=l1_node.id) + + +def load_l1_json(path: str) -> L1Node: + """从 JSON 文件加载单个 L1 节点（embedding=None）。 + + 参数: + path: JSON 文件路径。 + + 返回: + L1Node 实例。 + """ + with open(path, encoding="utf-8") as f: + data = json.load(f) + node = L1Node.from_dict(data) + logger.info("L1 中间结果已加载", path=path, l1_id=node.id) + return node diff --git a/tests/unit/test_tree_index.py b/tests/unit/test_tree_index.py new file mode 100644 index 0000000..4802e27 --- /dev/null +++ b/tests/unit/test_tree_index.py @@ -0,0 +1,227 @@ +"""TreeIndex 数据结构单元测试。""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from app.tree.index import ( + IndexMeta, + L1Card, + L1Node, + L2Card, + L2Node, + L3Card, + L3Node, + TreeIndex, +) + + +def _make_l3(idx: int = 0) -> L3Node: + return L3Node( + id=f"l1_0_l2_0_l3_{idx}", + card=L3Card( + frame_summary=f"帧{idx}描述", + visible_entities=["实体A"], + ongoing_actions=["动作A"], + visible_text=["文字A"], + spatial_layout="居中构图", + visual_attributes={"lighting": "明亮"}, + ), + timestamp=idx * 2.0, + frame_path=f"frames/l1_0_l2_0_l3_{idx}.jpg", + ) + + +def _make_l2(n_l3: int = 2) -> L2Node: + return L2Node( + id="l1_0_l2_0", + card=L2Card( + event_description="事件描述", + entities=["实体B"], + actions=["动作B"], + action_subjects=["主体B"], + visible_text=["文字B"], + spatial_relations="左右排列", + state_changes=None, + ), + time_range=(0.0, 60.0), + children=[_make_l3(i) for i in range(n_l3)], + ) + + +def _make_l1(n_l2: int = 1, n_l3: int = 2) -> L1Node: + return L1Node( + id="l1_0", + card=L1Card( + scene_summary="场景摘要", + main_setting="室内", + key_entities=["实体C"], + main_actions=["动作C"], + topic_keywords=["关键词"], + visible_text=["文字C"], + temporal_flow="从左到右", + ), + time_range=(0.0, 600.0), + children=[_make_l2(n_l3) for _ in range(n_l2)], + ) + + +def _make_index(n_l1: int = 1) -> TreeIndex: + meta = IndexMeta(source_path="/test/video.mp4", modality="video") + return TreeIndex(metadata=meta, roots=[_make_l1() for _ in range(n_l1)]) + + +class TestCards: + def test_l3_card_frozen(self): + card = L3Card( + frame_summary="desc", + visible_entities=[], + ongoing_actions=[], + visible_text=[], + spatial_layout="", + visual_attributes={}, + ) + with pytest.raises(AttributeError): + card.frame_summary = "changed" + + def test_l2_card_fields(self): + card = L2Card( + event_description="evt", + entities=[], + actions=[], + action_subjects=[], + visible_text=[], + spatial_relations="", + state_changes=None, + ) + assert card.event_description == "evt" + assert card.state_changes is None + + def test_l1_card_fields(self): + card = L1Card( + scene_summary="scene", + main_setting="outdoor", + key_entities=["e"], + main_actions=["a"], + topic_keywords=["k"], + visible_text=["t"], + temporal_flow="flow", + ) + assert card.scene_summary == "scene" + + +class TestNodes: + def test_l3_description_property(self): + node = _make_l3() + assert node.description == node.card.frame_summary + + def test_l2_description_property(self): + node = _make_l2() + assert node.description == node.card.event_description + + def test_l1_summary_property(self): + node = _make_l1() + assert node.summary == node.card.scene_summary + + def test_l3_default_embedding_none(self): + node = _make_l3() + assert node.embedding is None + + def test_l3_subtitle_default_none(self): + node = _make_l3() + assert node.subtitle is None + + +class TestTreeIndex: + def test_is_embedded_false_by_default(self): + index = _make_index() + assert not index.is_embedded + + def test_embed_all(self): + index = _make_index() + + def fake_embed(texts): + if isinstance(texts, str): + texts = [texts] + return np.random.randn(len(texts), 4).astype(np.float32) + + index.embed_all(fake_embed, "test-model", 4) + assert index.is_embedded + assert index.metadata.embed_model == "test-model" + assert index.metadata.embed_dim == 4 + + def test_l1_embeddings_shape(self): + index = _make_index(n_l1=2) + + def fake_embed(texts): + if isinstance(texts, str): + texts = [texts] + return np.random.randn(len(texts), 4).astype(np.float32) + + index.embed_all(fake_embed, "test-model", 4) + m = index.l1_embeddings() + assert m.shape == (2, 4) + + def test_get_node(self): + index = _make_index() + node = index.get_node(0, 0, 1) + assert node.id == "l1_0_l2_0_l3_1" + + def test_get_node_out_of_bounds(self): + index = _make_index() + with pytest.raises(IndexError): + index.get_node(99, 0, 0) + + +class TestSerialization: + def test_json_roundtrip(self, tmp_path): + index = _make_index() + path = tmp_path / "tree.json" + index.save_json(str(path)) + loaded = TreeIndex.load_json(str(path)) + assert len(loaded.roots) == 1 + assert loaded.roots[0].id == "l1_0" + assert loaded.roots[0].card.scene_summary == "场景摘要" + assert loaded.roots[0].children[0].children[0].card.frame_summary == "帧0描述" + + def test_json_roundtrip_with_embedding(self, tmp_path): + index = _make_index() + + def fake_embed(texts): + if isinstance(texts, str): + texts = [texts] + return np.random.randn(len(texts), 4).astype(np.float32) + + index.embed_all(fake_embed, "test-model", 4) + path = tmp_path / "tree_emb.json" + index.save_json(str(path), include_embedding=True) + loaded = TreeIndex.load_json(str(path)) + assert loaded.is_embedded + np.testing.assert_array_almost_equal( + loaded.roots[0].embedding, index.roots[0].embedding, decimal=5 + ) + + def test_l1_json_roundtrip(self, tmp_path): + from app.tree.index import load_l1_json, save_l1_json + + l1 = _make_l1() + path = tmp_path / "l1_0.json" + save_l1_json(str(path), l1) + loaded = load_l1_json(str(path)) + assert loaded.id == "l1_0" + assert len(loaded.children) == 1 + assert len(loaded.children[0].children) == 2 + + def test_id_uniqueness_validation(self, tmp_path): + """重复 ID 在反序列化时应报错。""" + index = _make_index() + d = index.to_dict() + d["roots"].append(d["roots"][0]) + path = tmp_path / "dup.json" + with open(path, "w") as f: + json.dump(d, f) + with pytest.raises(ValueError, match="重复"): + TreeIndex.load_json(str(path))