feat(tree): add get_node_text + get_children_info to TreeEnvironment
- get_node_text(node_id, anchor=False): returns raw text + optional anchor_map dict by parsing [cN]/[sN] prefixes from anchored text - get_children_info(node_id): returns structured child list with id/time_range/summary (description truncated to 120 chars) - Both methods reuse existing internal helpers (_node_full_text, _node_anchored_text, _get_children, _node_description, _format_time_range) - 9 new test cases across TestGetNodeText and TestGetChildrenInfo Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,519 @@
|
|||||||
|
"""TreeEnvironment:单棵视频树的运行时环境。
|
||||||
|
|
||||||
|
提供节点查询、字幕获取、帧路径解析和语义检索能力。
|
||||||
|
纯数据访问层——不涉及 LLM 调用,LLM 摘要逻辑属于 app/search/。
|
||||||
|
|
||||||
|
算法 #12 变更:分块 embedding → 单节点 embedding。
|
||||||
|
祖先去重 + 锚定验证逻辑保留自 TRM4。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from app.tree.index import L1Node, L2Node, L3Node, TreeIndex
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
# 节点联合类型(内部使用)
|
||||||
|
AnyNode = L1Node | L2Node | L3Node
|
||||||
|
|
||||||
|
# 各层级节点对应的主描述字段名
|
||||||
|
_LEVEL_LABEL = {
|
||||||
|
"L1": "场景层",
|
||||||
|
"L2": "事件层",
|
||||||
|
"L3": "关键帧层",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _node_level(node: AnyNode) -> str:
|
||||||
|
"""判断节点层级标签。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点实例。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
"L1" / "L2" / "L3"。
|
||||||
|
"""
|
||||||
|
if isinstance(node, L1Node):
|
||||||
|
return "L1"
|
||||||
|
if isinstance(node, L2Node):
|
||||||
|
return "L2"
|
||||||
|
return "L3"
|
||||||
|
|
||||||
|
|
||||||
|
def _node_description(node: AnyNode) -> str:
|
||||||
|
"""提取节点的主描述文本。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点实例。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
描述文本字符串。
|
||||||
|
"""
|
||||||
|
if isinstance(node, L1Node):
|
||||||
|
return node.card.scene_summary
|
||||||
|
if isinstance(node, L2Node):
|
||||||
|
return node.card.event_description
|
||||||
|
return node.card.frame_summary
|
||||||
|
|
||||||
|
|
||||||
|
def _collect_card_strings(node: AnyNode) -> list[str]:
|
||||||
|
"""从节点 card 中递归收集所有非空字符串字段。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点实例。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
字符串列表(每个非空字段值一项,含内嵌换行的按行拆分)。
|
||||||
|
"""
|
||||||
|
result: list[str] = []
|
||||||
|
_collect_from_obj(node.card, result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def _collect_from_obj(obj: object, out: list[str]) -> None:
|
||||||
|
"""递归收集任意嵌套结构中的非空字符串。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
obj: dict / list / str / 其他。
|
||||||
|
out: 收集结果列表(原地修改)。
|
||||||
|
"""
|
||||||
|
if isinstance(obj, str):
|
||||||
|
stripped = obj.strip()
|
||||||
|
if stripped:
|
||||||
|
out.append(stripped)
|
||||||
|
elif isinstance(obj, dict):
|
||||||
|
for v in obj.values():
|
||||||
|
_collect_from_obj(v, out)
|
||||||
|
elif isinstance(obj, (list, tuple)):
|
||||||
|
for item in obj:
|
||||||
|
_collect_from_obj(item, out)
|
||||||
|
elif hasattr(obj, "__dataclass_fields__"):
|
||||||
|
# frozen dataclass(Card 类型)
|
||||||
|
for field_name in obj.__dataclass_fields__:
|
||||||
|
_collect_from_obj(getattr(obj, field_name), out)
|
||||||
|
|
||||||
|
|
||||||
|
class TreeEnvironment:
|
||||||
|
"""单棵视频树的运行时环境,提供节点查询和语义检索。
|
||||||
|
|
||||||
|
纯数据访问层,不涉及 LLM 调用。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
index: 已加载的 TreeIndex 实例。
|
||||||
|
frames_dir: 帧文件目录路径(可选;未提供时使用节点自带的 frame_path)。
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
index: TreeIndex,
|
||||||
|
frames_dir: Path | None = None,
|
||||||
|
) -> None:
|
||||||
|
self._index = index
|
||||||
|
self._frames_dir = frames_dir
|
||||||
|
|
||||||
|
# O(1) 查找表:node_id → 节点实例
|
||||||
|
self._id_to_node: dict[str, AnyNode] = {}
|
||||||
|
# 父节点映射:node_id → parent_id(根节点为 None)
|
||||||
|
self._id_to_parent: dict[str, str | None] = {}
|
||||||
|
|
||||||
|
self._build_lookup_tables()
|
||||||
|
logger.debug(
|
||||||
|
"TreeEnvironment 初始化完成,节点数={}",
|
||||||
|
len(self._id_to_node),
|
||||||
|
)
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 初始化辅助
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def _build_lookup_tables(self) -> None:
|
||||||
|
"""遍历 TreeIndex 构建 _id_to_node 和 _id_to_parent 映射表。"""
|
||||||
|
for l1 in self._index.roots:
|
||||||
|
self._id_to_node[l1.id] = l1
|
||||||
|
self._id_to_parent[l1.id] = None
|
||||||
|
for l2 in l1.children:
|
||||||
|
self._id_to_node[l2.id] = l2
|
||||||
|
self._id_to_parent[l2.id] = l1.id
|
||||||
|
for l3 in l2.children:
|
||||||
|
self._id_to_node[l3.id] = l3
|
||||||
|
self._id_to_parent[l3.id] = l2.id
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 公开方法
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def view_node(self, node_id: str, *, anchor: bool = False) -> str:
|
||||||
|
"""返回节点卡片内容 + 子节点概览。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node_id: 节点 ID。
|
||||||
|
anchor: 为卡片字段添加行锚标 [c1] [s1] 供引用验证。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
格式化文本。
|
||||||
|
|
||||||
|
异常:
|
||||||
|
KeyError: 节点不存在。
|
||||||
|
"""
|
||||||
|
node = self._id_to_node.get(node_id)
|
||||||
|
if node is None:
|
||||||
|
raise KeyError(f"节点不存在: {node_id}")
|
||||||
|
|
||||||
|
level = _node_level(node)
|
||||||
|
level_label = _LEVEL_LABEL[level]
|
||||||
|
|
||||||
|
# 时间范围
|
||||||
|
time_range_str = self._format_time_range(node)
|
||||||
|
|
||||||
|
# 节点内容
|
||||||
|
content = self._node_anchored_text(node) if anchor else self._node_full_text(node)
|
||||||
|
|
||||||
|
parts = [
|
||||||
|
f"[节点] {node_id} | {level_label} | {time_range_str}",
|
||||||
|
"",
|
||||||
|
content,
|
||||||
|
]
|
||||||
|
|
||||||
|
# 子节点概览
|
||||||
|
children = self._get_children(node)
|
||||||
|
if children:
|
||||||
|
parts.append("")
|
||||||
|
parts.append(f"[子节点概览] {len(children)} 个子节点")
|
||||||
|
for child in children:
|
||||||
|
child_desc = _node_description(child)
|
||||||
|
child_time = self._format_time_range(child)
|
||||||
|
# 截断描述到 120 字符
|
||||||
|
if len(child_desc) > 120:
|
||||||
|
child_desc = child_desc[:120] + "..."
|
||||||
|
parts.append(f" - {child.id} | {child_time} | {child_desc}")
|
||||||
|
|
||||||
|
return "\n".join(parts)
|
||||||
|
|
||||||
|
def search_similar(
|
||||||
|
self,
|
||||||
|
query: str,
|
||||||
|
top_k: int = 5,
|
||||||
|
*,
|
||||||
|
embed_fn: Callable[[str | list[str]], np.ndarray] | None = None,
|
||||||
|
) -> list[tuple[str, float]]:
|
||||||
|
"""语义搜索 + 祖先去重。
|
||||||
|
|
||||||
|
算法 #12 变更:单节点 embedding(非分块),祖先去重 + 锚定验证保留。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
query: 搜索文本。
|
||||||
|
top_k: 返回数量。
|
||||||
|
embed_fn: 嵌入函数(未提供时使用 TreeIndex 已有 embedding)。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
[(node_id, score), ...] 按相似度降序。
|
||||||
|
|
||||||
|
异常:
|
||||||
|
ValueError: 节点未 embed 且未提供 embed_fn。
|
||||||
|
"""
|
||||||
|
if embed_fn is None:
|
||||||
|
raise ValueError(
|
||||||
|
"embed_fn 为必需参数:搜索 query 需要 embed_fn 来编码。请传入 embed_fn 参数。"
|
||||||
|
)
|
||||||
|
|
||||||
|
# 收集所有节点的 embedding(优先使用 TreeIndex 已有 embedding)
|
||||||
|
node_ids: list[str] = []
|
||||||
|
embeddings: list[np.ndarray] = []
|
||||||
|
|
||||||
|
if self._index.is_embedded:
|
||||||
|
# 使用已有 embedding
|
||||||
|
for nid, node in self._id_to_node.items():
|
||||||
|
if node.embedding is not None:
|
||||||
|
node_ids.append(nid)
|
||||||
|
embeddings.append(node.embedding)
|
||||||
|
else:
|
||||||
|
# 使用 embed_fn 为所有节点生成 embedding
|
||||||
|
all_ids = list(self._id_to_node.keys())
|
||||||
|
all_texts = [_node_description(self._id_to_node[nid]) for nid in all_ids]
|
||||||
|
all_embs = embed_fn(all_texts) # [N, D]
|
||||||
|
for i, nid in enumerate(all_ids):
|
||||||
|
node_ids.append(nid)
|
||||||
|
embeddings.append(all_embs[i])
|
||||||
|
|
||||||
|
if not embeddings:
|
||||||
|
return []
|
||||||
|
|
||||||
|
node_embeddings = np.stack(embeddings, axis=0) # [N, D]
|
||||||
|
# 归一化(确保余弦相似度正确)
|
||||||
|
norms = np.linalg.norm(node_embeddings, axis=1, keepdims=True)
|
||||||
|
norms = np.where(norms == 0, 1.0, norms)
|
||||||
|
node_embeddings = node_embeddings / norms
|
||||||
|
|
||||||
|
# 编码 query
|
||||||
|
query_emb = embed_fn(query) # [1, D]
|
||||||
|
|
||||||
|
if query_emb.ndim == 1:
|
||||||
|
query_emb = query_emb.reshape(1, -1)
|
||||||
|
# 归一化 query
|
||||||
|
q_norm = np.linalg.norm(query_emb)
|
||||||
|
if q_norm > 0:
|
||||||
|
query_emb = query_emb / q_norm
|
||||||
|
|
||||||
|
# 余弦相似度
|
||||||
|
scores = (node_embeddings @ query_emb.T).squeeze() # [N]
|
||||||
|
if scores.ndim == 0:
|
||||||
|
scores = scores.reshape(1)
|
||||||
|
|
||||||
|
# 按分数排序
|
||||||
|
scored_pairs = sorted(
|
||||||
|
zip(node_ids, scores.tolist(), strict=True),
|
||||||
|
key=lambda x: x[1],
|
||||||
|
reverse=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 祖先去重:如果更细粒度的子节点已入选,跳过其祖先
|
||||||
|
deduped: list[tuple[str, float]] = []
|
||||||
|
seen_prefixes: set[str] = set()
|
||||||
|
for nid, score in scored_pairs:
|
||||||
|
is_ancestor_of_seen = any(s.startswith(nid + "_") for s in seen_prefixes)
|
||||||
|
if is_ancestor_of_seen:
|
||||||
|
continue
|
||||||
|
deduped.append((nid, score))
|
||||||
|
seen_prefixes.add(nid)
|
||||||
|
if len(deduped) >= top_k:
|
||||||
|
break
|
||||||
|
|
||||||
|
return deduped
|
||||||
|
|
||||||
|
def get_node_text(
|
||||||
|
self,
|
||||||
|
node_id: str,
|
||||||
|
*,
|
||||||
|
anchor: bool = False,
|
||||||
|
) -> tuple[str, dict[str, str] | None]:
|
||||||
|
"""返回节点原始文本及可选的锚映射表。
|
||||||
|
|
||||||
|
供 SearchToolDispatcher 使用:将原始文本和锚映射传给
|
||||||
|
summarizer.summarize_node(),实现引用验证。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node_id: 节点 ID。
|
||||||
|
anchor: 若 True,返回带 [cN]/[sN] 锚标的文本并构建 anchor_map。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
(text, anchor_map) 元组。anchor=False 时 anchor_map 为 None;
|
||||||
|
anchor=True 时 anchor_map 为 {"c1": "行文本", "s1": "字幕行", ...}。
|
||||||
|
|
||||||
|
异常:
|
||||||
|
KeyError: 节点不存在。
|
||||||
|
"""
|
||||||
|
node = self._id_to_node.get(node_id)
|
||||||
|
if node is None:
|
||||||
|
raise KeyError(f"节点不存在: {node_id}")
|
||||||
|
|
||||||
|
if not anchor:
|
||||||
|
return self._node_full_text(node), None
|
||||||
|
|
||||||
|
anchored_text = self._node_anchored_text(node)
|
||||||
|
# 解析锚标行 "[c1] xxx" / "[s2] yyy" 构建映射
|
||||||
|
anchor_map: dict[str, str] = {}
|
||||||
|
anchor_pattern = re.compile(r"^\[([cs]\d+)\]\s(.+)$")
|
||||||
|
for line in anchored_text.splitlines():
|
||||||
|
m = anchor_pattern.match(line)
|
||||||
|
if m:
|
||||||
|
anchor_map[m.group(1)] = m.group(2)
|
||||||
|
|
||||||
|
return anchored_text, anchor_map
|
||||||
|
|
||||||
|
def get_children_info(self, node_id: str) -> list[dict[str, Any]]:
|
||||||
|
"""返回节点的直接子节点结构化信息。
|
||||||
|
|
||||||
|
供 SearchToolDispatcher 使用:将子节点列表传给
|
||||||
|
summarizer.summarize_children(),用于层级摘要。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node_id: 节点 ID。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
子节点信息列表,每项包含 {"id", "time_range", "summary"}。
|
||||||
|
L3 叶子节点返回空列表。
|
||||||
|
|
||||||
|
异常:
|
||||||
|
KeyError: 节点不存在。
|
||||||
|
"""
|
||||||
|
node = self._id_to_node.get(node_id)
|
||||||
|
if node is None:
|
||||||
|
raise KeyError(f"节点不存在: {node_id}")
|
||||||
|
|
||||||
|
children = self._get_children(node)
|
||||||
|
result: list[dict[str, Any]] = []
|
||||||
|
for child in children:
|
||||||
|
desc = _node_description(child)
|
||||||
|
if len(desc) > 120:
|
||||||
|
desc = desc[:120] + "..."
|
||||||
|
result.append({
|
||||||
|
"id": child.id,
|
||||||
|
"time_range": self._format_time_range(child),
|
||||||
|
"summary": desc,
|
||||||
|
})
|
||||||
|
return result
|
||||||
|
|
||||||
|
def get_subtitle(self, node_id: str) -> str:
|
||||||
|
"""返回节点字幕文本。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node_id: 节点 ID。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
字幕文本;无字幕或节点不存在时返回空字符串。
|
||||||
|
"""
|
||||||
|
node = self._id_to_node.get(node_id)
|
||||||
|
if node is None:
|
||||||
|
return ""
|
||||||
|
if isinstance(node, L3Node):
|
||||||
|
return node.subtitle or ""
|
||||||
|
return ""
|
||||||
|
|
||||||
|
def resolve_frame_paths(self, node_ids: list[str]) -> list[Path]:
|
||||||
|
"""node_id → 帧文件路径。支持 L3(直接映射)和 L2(展开为 L3 children)。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node_ids: 节点 ID 列表。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
帧文件 Path 列表。
|
||||||
|
|
||||||
|
异常:
|
||||||
|
KeyError: 节点不存在。
|
||||||
|
"""
|
||||||
|
if not node_ids:
|
||||||
|
return []
|
||||||
|
|
||||||
|
paths: list[Path] = []
|
||||||
|
for nid in node_ids:
|
||||||
|
node = self._id_to_node.get(nid)
|
||||||
|
if node is None:
|
||||||
|
raise KeyError(f"节点不存在: {nid}")
|
||||||
|
|
||||||
|
if isinstance(node, L3Node):
|
||||||
|
paths.append(self._l3_frame_path(node))
|
||||||
|
elif isinstance(node, L2Node):
|
||||||
|
# 展开为所有 L3 子节点
|
||||||
|
for l3 in node.children:
|
||||||
|
paths.append(self._l3_frame_path(l3))
|
||||||
|
else:
|
||||||
|
# L1 节点:展开为所有 L2 下的 L3
|
||||||
|
assert isinstance(node, L1Node)
|
||||||
|
for l2 in node.children:
|
||||||
|
for l3 in l2.children:
|
||||||
|
paths.append(self._l3_frame_path(l3))
|
||||||
|
|
||||||
|
return paths
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 内部辅助方法
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def _l3_frame_path(self, node: L3Node) -> Path:
|
||||||
|
"""将 L3 节点映射到帧文件路径。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: L3 节点。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
帧文件 Path。
|
||||||
|
"""
|
||||||
|
if self._frames_dir is not None:
|
||||||
|
# 从 node.id 中提取后缀(去掉 video_id 前缀)
|
||||||
|
# ID 格式: {video_id}_{L1_xxx_L2_xxx_L3_xxx}
|
||||||
|
# frame_path 格式: frames/{L1_xxx_L2_xxx_L3_xxx}.jpg
|
||||||
|
if node.frame_path:
|
||||||
|
return self._frames_dir / Path(node.frame_path).name
|
||||||
|
# fallback: 从 ID 推断
|
||||||
|
parts = node.id.split("_", 1)
|
||||||
|
suffix = parts[1] if len(parts) > 1 else node.id
|
||||||
|
return self._frames_dir / f"{suffix}.jpg"
|
||||||
|
|
||||||
|
# 无 frames_dir 时使用节点自带路径
|
||||||
|
if node.frame_path:
|
||||||
|
return Path(node.frame_path)
|
||||||
|
raise ValueError(f"L3 节点无 frame_path 且未提供 frames_dir: {node.id}")
|
||||||
|
|
||||||
|
def _node_full_text(self, node: AnyNode) -> str:
|
||||||
|
"""获取节点完整文本(card 所有字段 + subtitle)。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
拼接后的全文本。
|
||||||
|
"""
|
||||||
|
card_strings = _collect_card_strings(node)
|
||||||
|
text = "\n".join(card_strings)
|
||||||
|
if isinstance(node, L3Node) and node.subtitle:
|
||||||
|
text += f"\n字幕: {node.subtitle}"
|
||||||
|
return text
|
||||||
|
|
||||||
|
def _node_anchored_text(self, node: AnyNode) -> str:
|
||||||
|
"""获取带行号锚的节点文本。
|
||||||
|
|
||||||
|
card 字符串逐行编 [c1]..[cN],字幕逐行编 [s1]..[sM]。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
带锚文本。
|
||||||
|
"""
|
||||||
|
card_strings = _collect_card_strings(node)
|
||||||
|
# 拆分内嵌换行,确保一锚一行
|
||||||
|
card_lines: list[str] = []
|
||||||
|
for s in card_strings:
|
||||||
|
card_lines.extend(ln for ln in s.splitlines() if ln.strip())
|
||||||
|
|
||||||
|
sub_lines: list[str] = []
|
||||||
|
if isinstance(node, L3Node) and node.subtitle:
|
||||||
|
sub_lines = [ln for ln in node.subtitle.splitlines() if ln.strip()]
|
||||||
|
|
||||||
|
anchored: list[str] = []
|
||||||
|
for i, line in enumerate(card_lines, 1):
|
||||||
|
anchored.append(f"[c{i}] {line}")
|
||||||
|
for i, line in enumerate(sub_lines, 1):
|
||||||
|
anchored.append(f"[s{i}] {line}")
|
||||||
|
|
||||||
|
return "\n".join(anchored)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _format_time_range(node: AnyNode) -> str:
|
||||||
|
"""格式化节点的时间范围。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
"start-end s" 格式字符串,或 timestamp,或 "N/A"。
|
||||||
|
"""
|
||||||
|
if isinstance(node, (L1Node, L2Node)) and node.time_range:
|
||||||
|
return f"{node.time_range[0]:.1f}-{node.time_range[1]:.1f}s"
|
||||||
|
if isinstance(node, L3Node) and node.timestamp is not None:
|
||||||
|
return f"{node.timestamp:.1f}s"
|
||||||
|
return "N/A"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _get_children(node: AnyNode) -> list[AnyNode]:
|
||||||
|
"""获取节点的直接子节点列表。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
node: 树节点。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
子节点列表(L3 节点返回空列表)。
|
||||||
|
"""
|
||||||
|
if isinstance(node, L1Node):
|
||||||
|
return list(node.children)
|
||||||
|
if isinstance(node, L2Node):
|
||||||
|
return list(node.children)
|
||||||
|
return []
|
||||||
@@ -0,0 +1,324 @@
|
|||||||
|
"""TreeEnvironment 运行时单元测试。"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from app.tree.environment import TreeEnvironment
|
||||||
|
from app.tree.index import (
|
||||||
|
IndexMeta,
|
||||||
|
L1Card,
|
||||||
|
L1Node,
|
||||||
|
L2Card,
|
||||||
|
L2Node,
|
||||||
|
L3Card,
|
||||||
|
L3Node,
|
||||||
|
TreeIndex,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _make_test_index() -> TreeIndex:
|
||||||
|
"""构建测试用的三层树索引。"""
|
||||||
|
l3_0 = L3Node(
|
||||||
|
id="vid_L1_000_L2_000_L3_000",
|
||||||
|
card=L3Card(
|
||||||
|
"运动员在跑步",
|
||||||
|
["运动员"],
|
||||||
|
["跑步"],
|
||||||
|
["Nike"],
|
||||||
|
"居中",
|
||||||
|
{"lighting": "明亮"},
|
||||||
|
),
|
||||||
|
timestamp=1.0,
|
||||||
|
frame_path="frames/L1_000_L2_000_L3_000.jpg",
|
||||||
|
subtitle="he is running",
|
||||||
|
)
|
||||||
|
l3_1 = L3Node(
|
||||||
|
id="vid_L1_000_L2_000_L3_001",
|
||||||
|
card=L3Card(
|
||||||
|
"观众欢呼",
|
||||||
|
["观众"],
|
||||||
|
["欢呼"],
|
||||||
|
[],
|
||||||
|
"广角",
|
||||||
|
{},
|
||||||
|
),
|
||||||
|
timestamp=3.0,
|
||||||
|
frame_path="frames/L1_000_L2_000_L3_001.jpg",
|
||||||
|
)
|
||||||
|
l2 = L2Node(
|
||||||
|
id="vid_L1_000_L2_000",
|
||||||
|
card=L2Card(
|
||||||
|
"比赛片段",
|
||||||
|
["运动员"],
|
||||||
|
["跑步"],
|
||||||
|
["运动员"],
|
||||||
|
["Nike"],
|
||||||
|
"",
|
||||||
|
None,
|
||||||
|
),
|
||||||
|
time_range=(0.0, 10.0),
|
||||||
|
children=[l3_0, l3_1],
|
||||||
|
)
|
||||||
|
l1 = L1Node(
|
||||||
|
id="vid_L1_000",
|
||||||
|
card=L1Card(
|
||||||
|
"体育赛事",
|
||||||
|
"体育场",
|
||||||
|
["运动员"],
|
||||||
|
["比赛"],
|
||||||
|
["体育"],
|
||||||
|
["Nike"],
|
||||||
|
"从左到右",
|
||||||
|
),
|
||||||
|
time_range=(0.0, 10.0),
|
||||||
|
children=[l2],
|
||||||
|
)
|
||||||
|
return TreeIndex(metadata=IndexMeta("/test.mp4", "video"), roots=[l1])
|
||||||
|
|
||||||
|
|
||||||
|
class TestViewNode:
|
||||||
|
"""view_node 方法测试。"""
|
||||||
|
|
||||||
|
def test_l3_node(self) -> None:
|
||||||
|
"""L3 节点应显示帧描述。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
result = env.view_node("vid_L1_000_L2_000_L3_000")
|
||||||
|
assert "运动员在跑步" in result
|
||||||
|
assert "vid_L1_000_L2_000_L3_000" in result
|
||||||
|
|
||||||
|
def test_l2_node_shows_children(self) -> None:
|
||||||
|
"""L2 节点应显示子节点概览。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
result = env.view_node("vid_L1_000_L2_000")
|
||||||
|
assert "比赛片段" in result
|
||||||
|
assert "vid_L1_000_L2_000_L3_000" in result # child listed
|
||||||
|
|
||||||
|
def test_anchor_mode(self) -> None:
|
||||||
|
"""锚模式应在输出中添加 [cN] 标记。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
result = env.view_node("vid_L1_000_L2_000_L3_000", anchor=True)
|
||||||
|
assert "[c" in result # anchor markers present
|
||||||
|
|
||||||
|
def test_unknown_node_raises(self) -> None:
|
||||||
|
"""查询不存在的节点应抛出 KeyError。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
with pytest.raises(KeyError):
|
||||||
|
env.view_node("nonexistent")
|
||||||
|
|
||||||
|
|
||||||
|
class TestSearchSimilar:
|
||||||
|
"""search_similar 方法测试。"""
|
||||||
|
|
||||||
|
def test_returns_results(self) -> None:
|
||||||
|
"""使用 embed_fn 应返回搜索结果。"""
|
||||||
|
index = _make_test_index()
|
||||||
|
|
||||||
|
def fake_embed(
|
||||||
|
texts: str | list[str],
|
||||||
|
) -> np.ndarray:
|
||||||
|
if isinstance(texts, str):
|
||||||
|
texts = [texts]
|
||||||
|
return np.random.randn(len(texts), 4).astype(np.float32)
|
||||||
|
|
||||||
|
index.embed_all(fake_embed, "test", 4)
|
||||||
|
env = TreeEnvironment(index)
|
||||||
|
results = env.search_similar("运动员", top_k=3, embed_fn=fake_embed)
|
||||||
|
assert len(results) > 0
|
||||||
|
assert all(isinstance(r, tuple) and len(r) == 2 for r in results)
|
||||||
|
|
||||||
|
def test_ancestor_dedup(self) -> None:
|
||||||
|
"""祖先去重:如果 L3 已在结果中,其 L1/L2 祖先应被跳过。"""
|
||||||
|
index = _make_test_index()
|
||||||
|
# 手动设置 embedding,使 L3 节点分数高于 L1/L2
|
||||||
|
l3_0 = index.roots[0].children[0].children[0]
|
||||||
|
l3_1 = index.roots[0].children[0].children[1]
|
||||||
|
l2 = index.roots[0].children[0]
|
||||||
|
l1 = index.roots[0]
|
||||||
|
l3_0.embedding = np.array([1.0, 0, 0, 0], dtype=np.float32)
|
||||||
|
l3_1.embedding = np.array([0.9, 0.1, 0, 0], dtype=np.float32)
|
||||||
|
l2.embedding = np.array([0.5, 0.5, 0, 0], dtype=np.float32)
|
||||||
|
l1.embedding = np.array([0.3, 0.3, 0.3, 0], dtype=np.float32)
|
||||||
|
index.metadata.embed_model = "test"
|
||||||
|
index.metadata.embed_dim = 4
|
||||||
|
|
||||||
|
env = TreeEnvironment(index)
|
||||||
|
results = env.search_similar(
|
||||||
|
"运动员",
|
||||||
|
top_k=5,
|
||||||
|
embed_fn=lambda t: np.array([[1.0, 0, 0, 0]], dtype=np.float32),
|
||||||
|
)
|
||||||
|
result_ids = [r[0] for r in results]
|
||||||
|
# L3 节点应存在;其祖先应被去重跳过
|
||||||
|
assert "vid_L1_000_L2_000_L3_000" in result_ids
|
||||||
|
|
||||||
|
def test_with_embed_fn_overrides_existing(self) -> None:
|
||||||
|
"""即使已有 embedding,提供 embed_fn 时仍应用于 query 编码。"""
|
||||||
|
index = _make_test_index()
|
||||||
|
|
||||||
|
def fake_embed(
|
||||||
|
texts: str | list[str],
|
||||||
|
) -> np.ndarray:
|
||||||
|
if isinstance(texts, str):
|
||||||
|
texts = [texts]
|
||||||
|
return np.random.randn(len(texts), 4).astype(np.float32)
|
||||||
|
|
||||||
|
index.embed_all(fake_embed, "test", 4)
|
||||||
|
env = TreeEnvironment(index)
|
||||||
|
|
||||||
|
def query_embed(
|
||||||
|
texts: str | list[str],
|
||||||
|
) -> np.ndarray:
|
||||||
|
if isinstance(texts, str):
|
||||||
|
texts = [texts]
|
||||||
|
return np.ones((len(texts), 4), dtype=np.float32) * 0.5
|
||||||
|
|
||||||
|
results = env.search_similar("test", top_k=2, embed_fn=query_embed)
|
||||||
|
assert len(results) > 0
|
||||||
|
|
||||||
|
def test_no_embed_fn_raises(self) -> None:
|
||||||
|
"""未提供 embed_fn 时应报错。"""
|
||||||
|
index = _make_test_index()
|
||||||
|
env = TreeEnvironment(index)
|
||||||
|
with pytest.raises(ValueError, match="embed_fn"):
|
||||||
|
env.search_similar("test", top_k=3)
|
||||||
|
|
||||||
|
|
||||||
|
class TestGetSubtitle:
|
||||||
|
"""get_subtitle 方法测试。"""
|
||||||
|
|
||||||
|
def test_existing_subtitle(self) -> None:
|
||||||
|
"""有字幕的节点应返回字幕文本。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
assert env.get_subtitle("vid_L1_000_L2_000_L3_000") == "he is running"
|
||||||
|
|
||||||
|
def test_no_subtitle(self) -> None:
|
||||||
|
"""无字幕的节点应返回空字符串。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
assert env.get_subtitle("vid_L1_000_L2_000_L3_001") == ""
|
||||||
|
|
||||||
|
def test_unknown_node(self) -> None:
|
||||||
|
"""不存在的节点应返回空字符串。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
assert env.get_subtitle("nonexistent") == ""
|
||||||
|
|
||||||
|
|
||||||
|
class TestResolveFramePaths:
|
||||||
|
"""resolve_frame_paths 方法测试。"""
|
||||||
|
|
||||||
|
def test_l3_nodes(self) -> None:
|
||||||
|
"""L3 节点应映射到帧文件路径。"""
|
||||||
|
env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames"))
|
||||||
|
paths = env.resolve_frame_paths(["vid_L1_000_L2_000_L3_000"])
|
||||||
|
assert len(paths) == 1
|
||||||
|
assert "L1_000_L2_000_L3_000" in str(paths[0])
|
||||||
|
|
||||||
|
def test_l2_expands_to_children(self) -> None:
|
||||||
|
"""L2 节点应展开为其所有 L3 子节点的帧路径。"""
|
||||||
|
env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames"))
|
||||||
|
paths = env.resolve_frame_paths(["vid_L1_000_L2_000"])
|
||||||
|
assert len(paths) == 2 # 2 L3 children
|
||||||
|
|
||||||
|
def test_no_frames_dir_uses_node_path(self) -> None:
|
||||||
|
"""未提供 frames_dir 时应使用节点自带的 frame_path。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
paths = env.resolve_frame_paths(["vid_L1_000_L2_000_L3_000"])
|
||||||
|
assert len(paths) == 1
|
||||||
|
assert "L1_000_L2_000_L3_000" in str(paths[0])
|
||||||
|
|
||||||
|
def test_empty_list_returns_empty(self) -> None:
|
||||||
|
"""空列表应返回空结果。"""
|
||||||
|
env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames"))
|
||||||
|
paths = env.resolve_frame_paths([])
|
||||||
|
assert paths == []
|
||||||
|
|
||||||
|
|
||||||
|
class TestGetNodeText:
|
||||||
|
"""get_node_text 方法测试。"""
|
||||||
|
|
||||||
|
def test_normal_mode_returns_full_text(self) -> None:
|
||||||
|
"""默认模式应返回完整文本和 None anchor_map。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
text, anchor_map = env.get_node_text("vid_L1_000_L2_000_L3_000")
|
||||||
|
assert "运动员在跑步" in text
|
||||||
|
assert anchor_map is None
|
||||||
|
|
||||||
|
def test_anchor_mode_returns_anchored_text_and_map(self) -> None:
|
||||||
|
"""锚模式应返回带锚文本和 anchor_map 字典。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
text, anchor_map = env.get_node_text(
|
||||||
|
"vid_L1_000_L2_000_L3_000", anchor=True,
|
||||||
|
)
|
||||||
|
# 锚文本包含 [cN] 标记
|
||||||
|
assert "[c1]" in text
|
||||||
|
# anchor_map 非空,键为锚标(如 "c1"),值为对应行文本
|
||||||
|
assert anchor_map is not None
|
||||||
|
assert len(anchor_map) > 0
|
||||||
|
assert "c1" in anchor_map
|
||||||
|
# 字幕行也应在 anchor_map 中(该节点有 subtitle)
|
||||||
|
assert any(k.startswith("s") for k in anchor_map)
|
||||||
|
|
||||||
|
def test_nonexistent_node_raises(self) -> None:
|
||||||
|
"""查询不存在的节点应抛出 KeyError。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
with pytest.raises(KeyError):
|
||||||
|
env.get_node_text("nonexistent")
|
||||||
|
|
||||||
|
def test_node_without_subtitle_no_s_anchors(self) -> None:
|
||||||
|
"""无字幕的 L3 节点锚模式不应产生 [sN] 锚。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
text, anchor_map = env.get_node_text(
|
||||||
|
"vid_L1_000_L2_000_L3_001", anchor=True,
|
||||||
|
)
|
||||||
|
assert anchor_map is not None
|
||||||
|
assert not any(k.startswith("s") for k in anchor_map)
|
||||||
|
|
||||||
|
|
||||||
|
class TestGetChildrenInfo:
|
||||||
|
"""get_children_info 方法测试。"""
|
||||||
|
|
||||||
|
def test_l1_has_children(self) -> None:
|
||||||
|
"""L1 节点应返回其 L2 子节点信息列表。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
children = env.get_children_info("vid_L1_000")
|
||||||
|
assert len(children) == 1
|
||||||
|
child = children[0]
|
||||||
|
assert child["id"] == "vid_L1_000_L2_000"
|
||||||
|
assert "time_range" in child
|
||||||
|
assert "summary" in child
|
||||||
|
assert isinstance(child["summary"], str)
|
||||||
|
|
||||||
|
def test_l2_has_children(self) -> None:
|
||||||
|
"""L2 节点应返回其 L3 子节点信息列表。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
children = env.get_children_info("vid_L1_000_L2_000")
|
||||||
|
assert len(children) == 2
|
||||||
|
ids = [c["id"] for c in children]
|
||||||
|
assert "vid_L1_000_L2_000_L3_000" in ids
|
||||||
|
assert "vid_L1_000_L2_000_L3_001" in ids
|
||||||
|
|
||||||
|
def test_l3_has_no_children(self) -> None:
|
||||||
|
"""L3 叶子节点应返回空列表。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
children = env.get_children_info("vid_L1_000_L2_000_L3_000")
|
||||||
|
assert children == []
|
||||||
|
|
||||||
|
def test_nonexistent_node_raises(self) -> None:
|
||||||
|
"""查询不存在的节点应抛出 KeyError。"""
|
||||||
|
env = TreeEnvironment(_make_test_index())
|
||||||
|
with pytest.raises(KeyError):
|
||||||
|
env.get_children_info("nonexistent")
|
||||||
|
|
||||||
|
def test_summary_truncation(self) -> None:
|
||||||
|
"""超过 120 字符的描述应被截断。"""
|
||||||
|
index = _make_test_index()
|
||||||
|
# 修改 L2 的事件描述为超长文本
|
||||||
|
l2 = index.roots[0].children[0]
|
||||||
|
long_desc = "A" * 200
|
||||||
|
object.__setattr__(l2.card, "event_description", long_desc)
|
||||||
|
env = TreeEnvironment(index)
|
||||||
|
children = env.get_children_info("vid_L1_000")
|
||||||
|
assert len(children[0]["summary"]) == 123 # 120 + "..."
|
||||||
Reference in New Issue
Block a user