iomgaa
4e0e05210d
view_node 按题两轮摘要（summarize_node）会吞掉 entities/visible_text 字段信号，Agent 站在证据节点上仍漏读实体（benchmark 错题 M1，案例 786-2、872-3、750-1）。dispatcher 侧在摘要后确定性追加 [实体]/[画面 文字] 原文区块，LLM 无法吞掉。 - TreeEnvironment.node_entity_fields：按层级取 card 实体字段原文， 去空白、去重、分号拼接；空字段省键；未知节点抛 KeyError - _handle_view_node Phase 2.5：摘要后、子节点概览前追加实体区块 - 附带 ruff format 修正 test_tree_environment.py 两处既有格式 算法 #11（树环境语义搜索）数据访问层扩展，不改搜索算法本身。
417 lines
14 KiB
Python
417 lines
14 KiB
Python
"""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": "明亮"},
|
||
subtitle="he is running",
|
||
),
|
||
timestamp=1.0,
|
||
frame_path="frames/L1_000_L2_000_L3_000.jpg",
|
||
)
|
||
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 + "..."
|
||
|
||
|
||
# ── node_entity_fields 测试（Spec-1 B）───────────────────────
|
||
|
||
|
||
def _make_entity_test_index() -> TreeIndex:
|
||
"""带实体字段的最小三层树（含一个空字段 L2）。"""
|
||
l3 = L3Node(
|
||
id="vid_L1_000_L2_000_L3_000",
|
||
card=L3Card(
|
||
frame_summary="一名男子戴耳机",
|
||
visible_entities=["Bluetooth headset (both ears)", "man"],
|
||
ongoing_actions=["talking"],
|
||
visible_text=["EARPHONE BOTTLE OPENER"],
|
||
spatial_layout="man center",
|
||
visual_attributes={},
|
||
),
|
||
timestamp=10.0,
|
||
)
|
||
l2 = L2Node(
|
||
id="vid_L1_000_L2_000",
|
||
card=L2Card(
|
||
event_description="产品评测",
|
||
entities=["Bluetooth headset (both ears)", "reviewer"],
|
||
actions=["reviewing"],
|
||
action_subjects=["reviewer"],
|
||
visible_text=["$9.99"],
|
||
spatial_relations="",
|
||
state_changes=None,
|
||
),
|
||
time_range=(0.0, 60.0),
|
||
children=[l3],
|
||
)
|
||
l2_empty = L2Node(
|
||
id="vid_L1_000_L2_001",
|
||
card=L2Card(
|
||
event_description="空镜",
|
||
entities=[],
|
||
actions=[],
|
||
action_subjects=[],
|
||
visible_text=[],
|
||
spatial_relations="",
|
||
state_changes=None,
|
||
),
|
||
time_range=(60.0, 120.0),
|
||
)
|
||
l1 = L1Node(
|
||
id="vid_L1_000",
|
||
card=L1Card(
|
||
scene_summary="评测场景",
|
||
main_setting="室内",
|
||
key_entities=["reviewer"],
|
||
main_actions=["评测"],
|
||
topic_keywords=["数码"],
|
||
visible_text=[],
|
||
temporal_flow="线性",
|
||
),
|
||
time_range=(0.0, 120.0),
|
||
children=[l2, l2_empty],
|
||
)
|
||
return TreeIndex(metadata=IndexMeta("/test.mp4", "video"), roots=[l1])
|
||
|
||
|
||
class TestNodeEntityFields:
|
||
"""node_entity_fields 方法测试（B 修复：dispatcher 追加原文）。"""
|
||
|
||
def test_l2_entities_and_visible_text(self) -> None:
|
||
"""L2 节点应返回 entities 和 visible_text 原文。"""
|
||
env = TreeEnvironment(_make_entity_test_index())
|
||
fields = env.node_entity_fields("vid_L1_000_L2_000")
|
||
assert "Bluetooth headset (both ears)" in fields["实体"]
|
||
assert "$9.99" in fields["画面文字"]
|
||
|
||
def test_l3_visible_entities(self) -> None:
|
||
"""L3 节点应返回 visible_entities 和 visible_text 原文。"""
|
||
env = TreeEnvironment(_make_entity_test_index())
|
||
fields = env.node_entity_fields("vid_L1_000_L2_000_L3_000")
|
||
assert "Bluetooth headset (both ears)" in fields["实体"]
|
||
assert "EARPHONE BOTTLE OPENER" in fields["画面文字"]
|
||
|
||
def test_empty_fields_omitted(self) -> None:
|
||
"""实体/画面文字均为空时应返回空字典（不含空键）。"""
|
||
env = TreeEnvironment(_make_entity_test_index())
|
||
assert env.node_entity_fields("vid_L1_000_L2_001") == {}
|
||
|
||
def test_unknown_node_raises(self) -> None:
|
||
"""查询不存在的节点应抛出 KeyError。"""
|
||
env = TreeEnvironment(_make_entity_test_index())
|
||
with pytest.raises(KeyError):
|
||
env.node_entity_fields("nonexistent")
|