761 lines
26 KiB
Python
761 lines
26 KiB
Python
"""三层树索引核心数据结构。
|
||
|
||
定义 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
|