522 lines
16 KiB
Python
522 lines
16 KiB
Python
"""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 []
|