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