"""三层树索引核心数据结构。 定义 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: self._embed_l2_subtree(l2, embed_fn) self.metadata.embed_model = model_name self.metadata.embed_dim = embed_dim logger.info( "embed_all 完成", model=model_name, embed_dim=embed_dim, ) def _embed_l2_subtree( self, l2: L2Node, embed_fn: Callable[[str | list[str]], np.ndarray], ) -> None: """对单个 L2 节点及其 L3 子节点执行 embedding（仅处理 embedding 为 None 的节点）。 参数: l2: 待 embed 的 L2 节点。 embed_fn: EmbeddingModel.embed 方法，接受 str 或 List[str]， 返回 [N, D] ndarray。 """ 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 # ------------------------------------------------------------------ # # 嵌入矩阵提取 # ------------------------------------------------------------------ # 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