feat(repair): regenerator + supplement 防御修复 + 迁移脚本
- 新增 app/tree/repair/regenerator.py（VLM 重生成 + 级联修复） - supplement.py: deduplicate_field str() 防御 + inject_value strip - patch.py: ruff format 格式化 - repair_trees.sh: conda source 激活修复 - 新增 migrate_from_trm4.sh 迁移工具 - enhance/__init__.py → repair/__init__.py 重命名 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,513 @@
|
||||
"""树修复重生成器：VLM 重新描述问题节点 + 底向上级联。
|
||||
|
||||
底向上修复流程:
|
||||
1. 收集需修复的 L3 节点 → VLM 重新描述帧
|
||||
2. 收集受影响的 L2 → LLM 从 L3 children 聚合
|
||||
3. 收集受影响的 L1 → LLM 从 L2 children 聚合
|
||||
|
||||
仅处理 issue_type == "empty_field" 且 level == 3 的问题节点。
|
||||
帧文件不存在时跳过该节点（不中断整体修复流程）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from app.tree.index import (
|
||||
L1Card,
|
||||
L1Node,
|
||||
L2Card,
|
||||
L2Node,
|
||||
L3Card,
|
||||
L3Node,
|
||||
TreeIndex,
|
||||
)
|
||||
from app.tree.subtitle import extract_subtitle_for_range
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from app.tree.repair.detector import NodeIssue
|
||||
from app.tree.subtitle import SRTEntry
|
||||
from core.protocols import LLMProvider, VLMProvider
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt 常量（与 VideoTreeBuilder 保持一致风格）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_L3_REPAIR_PROMPT = (
|
||||
'该片段的整体内容: "{l2_description}"\n'
|
||||
"用一到两句话描述这帧画面的具体内容。"
|
||||
"重点关注: 动作、物体变化、文字信息、人物表情。\n"
|
||||
"{subtitle_block}"
|
||||
"返回 JSON 对象，包含以下字段：\n"
|
||||
"- frame_summary: 画面描述\n"
|
||||
"- visible_entities: 可见实体列表\n"
|
||||
"- ongoing_actions: 动作列表\n"
|
||||
"- visible_text: 可见文字列表\n"
|
||||
"- spatial_layout: 空间布局\n"
|
||||
'- visual_attributes: {{"lighting": "...", "dominant_colors": [...], "camera_angle": "..."}}\n'
|
||||
"只返回 JSON 对象，不要其他内容。"
|
||||
)
|
||||
|
||||
_L2_REGEN_PROMPT = (
|
||||
"以下是一个视频片段中各帧的描述:\n{l3_texts}\n"
|
||||
"用1-2句话描述该片段的核心内容。\n"
|
||||
"返回 JSON 对象，包含以下字段：\n"
|
||||
"- event_description: 1-2句片段描述\n"
|
||||
"- entities: 可见实体列表\n"
|
||||
"- actions: 动作列表\n"
|
||||
"- action_subjects: 动作主体列表\n"
|
||||
"- visible_text: 画面中可见文字列表\n"
|
||||
"- spatial_relations: 空间关系描述\n"
|
||||
"- state_changes: 状态变化描述（无则 null）\n"
|
||||
"只返回 JSON 对象，不要其他内容。"
|
||||
)
|
||||
|
||||
_L1_REGEN_PROMPT = (
|
||||
"以下是一个视频段落中各片段的描述:\n{l2_texts}\n"
|
||||
"用2-3句话总结该段落的整体内容，涵盖所有片段的主题。\n"
|
||||
"返回 JSON 对象，包含以下字段：\n"
|
||||
"- scene_summary: 2-3句段落摘要\n"
|
||||
"- main_setting: 主要场景\n"
|
||||
"- key_entities: 关键实体列表\n"
|
||||
"- main_actions: 主要动作列表\n"
|
||||
"- topic_keywords: 主题关键词列表\n"
|
||||
"- visible_text: 出现的文字列表\n"
|
||||
"- temporal_flow: 时间流向描述\n"
|
||||
"只返回 JSON 对象，不要其他内容。"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 统计数据类
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class RepairStats:
|
||||
"""修复统计信息。
|
||||
|
||||
属性:
|
||||
l3_repaired: 修复的 L3 节点数。
|
||||
l2_regenerated: 重生成的 L2 节点数。
|
||||
l1_regenerated: 重生成的 L1 节点数。
|
||||
"""
|
||||
|
||||
l3_repaired: int = 0
|
||||
l2_regenerated: int = 0
|
||||
l1_regenerated: int = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON 解析辅助（复用 VideoTreeBuilder 的解析逻辑）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_json(raw: str) -> Any:
|
||||
"""从 VLM/LLM 原始输出中提取 JSON（处理 markdown 代码块包裹）。
|
||||
|
||||
参数:
|
||||
raw: 原始返回字符串。
|
||||
|
||||
返回:
|
||||
解析后的 Python 对象（dict/list），解析失败返回 None。
|
||||
"""
|
||||
raw = raw.strip()
|
||||
# Phase 1: 尝试提取 markdown 代码块中的 JSON
|
||||
code_match = re.search(
|
||||
r"```(?:json)?\s*([\[{].*?[\]}])\s*```",
|
||||
raw,
|
||||
re.DOTALL,
|
||||
)
|
||||
if code_match:
|
||||
raw = code_match.group(1)
|
||||
|
||||
# Phase 2: 直接解析
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Phase 3: 尝试提取裸 JSON 对象/数组
|
||||
json_match = re.search(r"[\[{].*[\]}]", raw, re.DOTALL)
|
||||
if json_match:
|
||||
try:
|
||||
return json.loads(json_match.group())
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _parse_l3_card(raw: str) -> L3Card | None:
|
||||
"""解析 VLM 输出为 L3Card。解析失败返回 None。
|
||||
|
||||
参数:
|
||||
raw: VLM 原始返回字符串。
|
||||
|
||||
返回:
|
||||
L3Card 实例或 None（解析失败时）。
|
||||
"""
|
||||
data = _extract_json(raw)
|
||||
if isinstance(data, dict):
|
||||
try:
|
||||
return L3Card(
|
||||
frame_summary=str(data["frame_summary"]),
|
||||
visible_entities=list(data["visible_entities"]),
|
||||
ongoing_actions=list(data["ongoing_actions"]),
|
||||
visible_text=list(data["visible_text"]),
|
||||
spatial_layout=str(data["spatial_layout"]),
|
||||
visual_attributes=dict(data["visual_attributes"]),
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _parse_l2_card(raw: str) -> L2Card | None:
|
||||
"""解析 LLM 输出为 L2Card。解析失败返回 None。
|
||||
|
||||
参数:
|
||||
raw: LLM 原始返回字符串。
|
||||
|
||||
返回:
|
||||
L2Card 实例或 None（解析失败时）。
|
||||
"""
|
||||
data = _extract_json(raw)
|
||||
if isinstance(data, dict):
|
||||
try:
|
||||
state_changes = data.get("state_changes")
|
||||
if state_changes is not None:
|
||||
state_changes = str(state_changes)
|
||||
return L2Card(
|
||||
event_description=str(data["event_description"]),
|
||||
entities=list(data["entities"]),
|
||||
actions=list(data["actions"]),
|
||||
action_subjects=list(data["action_subjects"]),
|
||||
visible_text=list(data["visible_text"]),
|
||||
spatial_relations=str(data["spatial_relations"]),
|
||||
state_changes=state_changes,
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _parse_l1_card(raw: str) -> L1Card | None:
|
||||
"""解析 LLM 输出为 L1Card。解析失败返回 None。
|
||||
|
||||
参数:
|
||||
raw: LLM 原始返回字符串。
|
||||
|
||||
返回:
|
||||
L1Card 实例或 None（解析失败时）。
|
||||
"""
|
||||
data = _extract_json(raw)
|
||||
if isinstance(data, dict):
|
||||
try:
|
||||
return L1Card(
|
||||
scene_summary=str(data["scene_summary"]),
|
||||
main_setting=str(data["main_setting"]),
|
||||
key_entities=list(data["key_entities"]),
|
||||
main_actions=list(data["main_actions"]),
|
||||
topic_keywords=list(data["topic_keywords"]),
|
||||
visible_text=list(data["visible_text"]),
|
||||
temporal_flow=str(data["temporal_flow"]),
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 节点查找辅助
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_node_lookup(
|
||||
index: TreeIndex,
|
||||
) -> tuple[
|
||||
dict[str, L3Node],
|
||||
dict[str, L2Node],
|
||||
dict[str, L1Node],
|
||||
dict[str, L2Node],
|
||||
dict[str, L1Node],
|
||||
]:
|
||||
"""构建节点 ID 到节点的查找表 + 子节点到父节点的映射。
|
||||
|
||||
参数:
|
||||
index: 树索引。
|
||||
|
||||
返回:
|
||||
(l3_by_id, l2_by_id, l1_by_id, l3_parent_l2, l2_parent_l1)
|
||||
- l3_by_id: L3 节点 ID → L3Node
|
||||
- l2_by_id: L2 节点 ID → L2Node
|
||||
- l1_by_id: L1 节点 ID → L1Node
|
||||
- l3_parent_l2: L3 节点 ID → 其父 L2Node
|
||||
- l2_parent_l1: L2 节点 ID → 其父 L1Node
|
||||
"""
|
||||
l3_by_id: dict[str, L3Node] = {}
|
||||
l2_by_id: dict[str, L2Node] = {}
|
||||
l1_by_id: dict[str, L1Node] = {}
|
||||
l3_parent_l2: dict[str, L2Node] = {}
|
||||
l2_parent_l1: dict[str, L1Node] = {}
|
||||
|
||||
for l1 in index.roots:
|
||||
l1_by_id[l1.id] = l1
|
||||
for l2 in l1.children:
|
||||
l2_by_id[l2.id] = l2
|
||||
l2_parent_l1[l2.id] = l1
|
||||
for l3 in l2.children:
|
||||
l3_by_id[l3.id] = l3
|
||||
l3_parent_l2[l3.id] = l2
|
||||
|
||||
return l3_by_id, l2_by_id, l1_by_id, l3_parent_l2, l2_parent_l1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 字幕辅助
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_subtitle_block(
|
||||
srt_entries: list[SRTEntry] | None,
|
||||
timestamp: float | None,
|
||||
) -> str:
|
||||
"""构建字幕注入文本块。
|
||||
|
||||
参数:
|
||||
srt_entries: SRT 字幕条目列表。
|
||||
timestamp: 帧时间戳（秒）。
|
||||
|
||||
返回:
|
||||
字幕文本块字符串（无匹配时返回空字符串）。
|
||||
"""
|
||||
if not srt_entries or timestamp is None:
|
||||
return ""
|
||||
window = 2.0
|
||||
start = max(0.0, timestamp - window)
|
||||
end = timestamp + window
|
||||
text = extract_subtitle_for_range(srt_entries, (start, end))
|
||||
if not text:
|
||||
return ""
|
||||
return f"字幕信息：\n{text}\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 主修复函数
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def repair_tree(
|
||||
index: TreeIndex,
|
||||
issues: list[NodeIssue],
|
||||
vlm: VLMProvider,
|
||||
llm: LLMProvider,
|
||||
frames_dir: Path,
|
||||
srt_entries: list[SRTEntry] | None = None,
|
||||
) -> RepairStats:
|
||||
"""修复有问题的节点，底向上级联。
|
||||
|
||||
流程:
|
||||
1. 收集需修复的 L3 节点 → VLM 重新描述帧
|
||||
2. 收集受影响的 L2 → LLM 从 L3 children 聚合
|
||||
3. 收集受影响的 L1 → LLM 从 L2 children 聚合
|
||||
|
||||
参数:
|
||||
index: 待修复的 TreeIndex（原地修改）。
|
||||
issues: detect_issues() 返回的问题列表。
|
||||
vlm: VLM 调用端口。
|
||||
llm: LLM 调用端口。
|
||||
frames_dir: 帧文件根目录。
|
||||
srt_entries: 字幕条目列表（可选）。
|
||||
|
||||
返回:
|
||||
RepairStats 统计。
|
||||
"""
|
||||
stats = RepairStats()
|
||||
|
||||
if not issues:
|
||||
logger.info("无修复任务，跳过")
|
||||
return stats
|
||||
|
||||
# 构建查找表
|
||||
l3_by_id, l2_by_id, l1_by_id, l3_parent_l2, l2_parent_l1 = _build_node_lookup(index)
|
||||
|
||||
# Step 1: 修复 L3 节点（仅处理 empty_field + level 3）
|
||||
l3_issues = [
|
||||
issue for issue in issues if issue.issue_type == "empty_field" and issue.level == 3
|
||||
]
|
||||
|
||||
affected_l2_ids: set[str] = set()
|
||||
|
||||
for issue in l3_issues:
|
||||
l3_node = l3_by_id.get(issue.node_id)
|
||||
if l3_node is None:
|
||||
logger.warning(
|
||||
"L3 节点 ID 未在树中找到，跳过",
|
||||
node_id=issue.node_id,
|
||||
)
|
||||
continue
|
||||
|
||||
# 查找帧文件
|
||||
if l3_node.frame_path is None:
|
||||
logger.warning(
|
||||
"L3 节点无 frame_path，跳过",
|
||||
node_id=issue.node_id,
|
||||
)
|
||||
continue
|
||||
|
||||
frame_file = frames_dir / l3_node.frame_path
|
||||
if not frame_file.exists():
|
||||
logger.warning(
|
||||
"L3 帧文件不存在，跳过修复",
|
||||
node_id=issue.node_id,
|
||||
frame_path=str(frame_file),
|
||||
)
|
||||
continue
|
||||
|
||||
# 获取 L2 父节点描述作为上下文
|
||||
parent_l2 = l3_parent_l2.get(issue.node_id)
|
||||
l2_description = parent_l2.card.event_description if parent_l2 else ""
|
||||
|
||||
# 构建字幕块
|
||||
subtitle_block = _build_subtitle_block(srt_entries, l3_node.timestamp)
|
||||
|
||||
# VLM 重新描述帧
|
||||
prompt = _L3_REPAIR_PROMPT.format(
|
||||
l2_description=l2_description,
|
||||
subtitle_block=subtitle_block,
|
||||
)
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
try:
|
||||
response = await vlm.chat_with_images(messages, [str(frame_file)])
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"L3 修复 VLM 调用失败，跳过: {}",
|
||||
exc,
|
||||
node_id=issue.node_id,
|
||||
)
|
||||
continue
|
||||
|
||||
new_card = _parse_l3_card(response.content)
|
||||
if new_card is None:
|
||||
logger.warning(
|
||||
"L3 修复 VLM 输出解析失败，跳过",
|
||||
node_id=issue.node_id,
|
||||
raw_preview=response.content[:200],
|
||||
)
|
||||
continue
|
||||
|
||||
# 原地替换 card（L3Node.card 不是 frozen dataclass 的限制字段）
|
||||
l3_node.card = new_card
|
||||
stats.l3_repaired += 1
|
||||
|
||||
# 标记受影响的 L2 父节点
|
||||
if parent_l2 is not None:
|
||||
affected_l2_ids.add(parent_l2.id)
|
||||
|
||||
logger.debug(
|
||||
"L3 节点修复完成",
|
||||
node_id=issue.node_id,
|
||||
frame_summary=new_card.frame_summary[:50],
|
||||
)
|
||||
|
||||
# Step 2: 重生成受影响的 L2 节点
|
||||
affected_l1_ids: set[str] = set()
|
||||
|
||||
for l2_id in affected_l2_ids:
|
||||
l2_node = l2_by_id.get(l2_id)
|
||||
if l2_node is None:
|
||||
continue
|
||||
|
||||
# 从 L3 children 聚合描述
|
||||
l3_texts = "\n".join(f"- {l3.card.frame_summary}" for l3 in l2_node.children)
|
||||
prompt = _L2_REGEN_PROMPT.format(l3_texts=l3_texts)
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
try:
|
||||
response = await llm.chat(messages)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"L2 重生成 LLM 调用失败，跳过: {}",
|
||||
exc,
|
||||
l2_id=l2_id,
|
||||
)
|
||||
continue
|
||||
|
||||
new_card = _parse_l2_card(response.content)
|
||||
if new_card is None:
|
||||
logger.warning(
|
||||
"L2 重生成 LLM 输出解析失败，跳过",
|
||||
l2_id=l2_id,
|
||||
raw_preview=response.content[:200],
|
||||
)
|
||||
continue
|
||||
|
||||
l2_node.card = new_card
|
||||
stats.l2_regenerated += 1
|
||||
|
||||
# 标记受影响的 L1 父节点
|
||||
parent_l1 = l2_parent_l1.get(l2_id)
|
||||
if parent_l1 is not None:
|
||||
affected_l1_ids.add(parent_l1.id)
|
||||
|
||||
logger.debug(
|
||||
"L2 节点重生成完成",
|
||||
l2_id=l2_id,
|
||||
event_description=new_card.event_description[:50],
|
||||
)
|
||||
|
||||
# Step 3: 重生成受影响的 L1 节点
|
||||
for l1_id in affected_l1_ids:
|
||||
l1_node = l1_by_id.get(l1_id)
|
||||
if l1_node is None:
|
||||
continue
|
||||
|
||||
# 从 L2 children 聚合描述
|
||||
l2_texts = "\n".join(f"- {l2.card.event_description}" for l2 in l1_node.children)
|
||||
prompt = _L1_REGEN_PROMPT.format(l2_texts=l2_texts)
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
try:
|
||||
response = await llm.chat(messages)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"L1 重生成 LLM 调用失败，跳过: {}",
|
||||
exc,
|
||||
l1_id=l1_id,
|
||||
)
|
||||
continue
|
||||
|
||||
new_card = _parse_l1_card(response.content)
|
||||
if new_card is None:
|
||||
logger.warning(
|
||||
"L1 重生成 LLM 输出解析失败，跳过",
|
||||
l1_id=l1_id,
|
||||
raw_preview=response.content[:200],
|
||||
)
|
||||
continue
|
||||
|
||||
l1_node.card = new_card
|
||||
stats.l1_regenerated += 1
|
||||
|
||||
logger.debug(
|
||||
"L1 节点重生成完成",
|
||||
l1_id=l1_id,
|
||||
scene_summary=new_card.scene_summary[:50],
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"树修复完成",
|
||||
l3_repaired=stats.l3_repaired,
|
||||
l2_regenerated=stats.l2_regenerated,
|
||||
l1_regenerated=stats.l1_regenerated,
|
||||
)
|
||||
return stats
|
||||
Reference in New Issue
Block a user