527 lines
16 KiB
Python
527 lines
16 KiB
Python
"""Q&A 反向补全：基于问题答案分析，将树中缺失的事实注入节点。
|
||
|
||
通过 LLM 分析正确答案需要哪些关键事实，再检查树中是否已有，
|
||
对缺失事实执行注入。仅注入客观事实（人名、地点、得分、物体名称），
|
||
不注入情感、因果推理、时间推理等主观或高阶信息。
|
||
|
||
与 TRM4 的关键差异:
|
||
- 树结构从扁平 dict 变为 TreeIndex（L1Node → L2Node → L3Node）。
|
||
- Card 为 frozen dataclass，注入时使用 dataclasses.replace() 创建新实例。
|
||
- LLMProvider 为异步接口，返回 LLMResponse（.content 获取文本）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from dataclasses import dataclass, replace
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
from loguru import logger
|
||
|
||
if TYPE_CHECKING:
|
||
from app.tree.index import L1Node, L2Node, L3Node, TreeIndex
|
||
from core.protocols import LLMProvider
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 允许注入的类别白名单
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_ALLOWED_CATEGORIES = frozenset(
|
||
{
|
||
"person_name",
|
||
"location",
|
||
"score_number",
|
||
"object_name",
|
||
}
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 类别 → 默认注入字段映射（L2 Card 字段名）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_CATEGORY_DEFAULT_FIELD: dict[str, str] = {
|
||
"person_name": "entities",
|
||
"location": "entities",
|
||
"score_number": "entities",
|
||
"object_name": "entities",
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 统计
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class SupplementStats:
|
||
"""反向补全统计信息。
|
||
|
||
属性:
|
||
questions_analyzed: 分析的问题数量。
|
||
facts_injected: 成功注入的事实数量。
|
||
facts_skipped: 跳过的事实数量（类别不在白名单中）。
|
||
"""
|
||
|
||
questions_analyzed: int = 0
|
||
facts_injected: int = 0
|
||
facts_skipped: int = 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 去重
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def deduplicate_field(values: list[str]) -> list[str]:
|
||
"""大小写归一化去重，保留首次出现的原始形式。
|
||
|
||
参数:
|
||
values: 待去重字符串列表。
|
||
|
||
返回:
|
||
去重后的列表，保留各值首次出现时的大小写。
|
||
空字符串和纯空白字符串会被跳过。
|
||
"""
|
||
seen: set[str] = set()
|
||
result: list[str] = []
|
||
for v in values:
|
||
key = v.strip().lower()
|
||
if key and key not in seen:
|
||
seen.add(key)
|
||
result.append(v)
|
||
return result
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 节点查找
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _find_node_by_id(
|
||
index: TreeIndex,
|
||
node_id: str,
|
||
) -> tuple[L1Node | L2Node | L3Node | None, int]:
|
||
"""在 TreeIndex 中按 ID 查找节点，返回节点和所属层级。
|
||
|
||
参数:
|
||
index: 树索引。
|
||
node_id: 目标节点 ID。
|
||
|
||
返回:
|
||
(node, level) 元组。找不到时返回 (None, -1)。
|
||
level: 1=L1, 2=L2, 3=L3。
|
||
"""
|
||
for l1 in index.roots:
|
||
if l1.id == node_id:
|
||
return l1, 1
|
||
for l2 in l1.children:
|
||
if l2.id == node_id:
|
||
return l2, 2
|
||
for l3 in l2.children:
|
||
if l3.id == node_id:
|
||
return l3, 3
|
||
return None, -1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 单值注入（适配 frozen Card）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _inject_into_l2(l2: L2Node, field: str, value: str) -> bool:
|
||
"""向 L2 节点的 Card 指定字段注入一个值。
|
||
|
||
使用 dataclasses.replace() 创建新的 frozen L2Card。
|
||
仅支持 list[str] 类型字段（entities / actions / action_subjects / visible_text）
|
||
和 str 类型字段（event_description / spatial_relations / state_changes）。
|
||
|
||
参数:
|
||
l2: L2 节点（card 会被替换为新实例）。
|
||
field: 目标字段名。
|
||
value: 要注入的值。
|
||
|
||
返回:
|
||
True 表示实际注入了新内容，False 表示已存在（跳过）。
|
||
"""
|
||
card = l2.card
|
||
current = getattr(card, field, None)
|
||
|
||
if current is None:
|
||
# 字段不存在于 Card schema，跳过
|
||
logger.debug("L2Card 无字段 {}，跳过注入", field)
|
||
return False
|
||
|
||
if isinstance(current, list):
|
||
lower_set = {v.strip().lower() for v in current if isinstance(v, str)}
|
||
if value.strip().lower() in lower_set:
|
||
return False
|
||
new_list = deduplicate_field([*current, value])
|
||
l2.card = replace(card, **{field: new_list})
|
||
return True
|
||
|
||
if isinstance(current, str):
|
||
if value.strip().lower() in current.lower():
|
||
return False
|
||
new_val = current + "; " + value if current else value
|
||
l2.card = replace(card, **{field: new_val})
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def _inject_into_l3(l3: L3Node, field: str, value: str) -> bool:
|
||
"""向 L3 节点的 Card 指定字段注入一个值。
|
||
|
||
使用 dataclasses.replace() 创建新的 frozen L3Card。
|
||
|
||
参数:
|
||
l3: L3 节点（card 会被替换为新实例）。
|
||
field: 目标字段名。
|
||
value: 要注入的值。
|
||
|
||
返回:
|
||
True 表示实际注入了新内容，False 表示已存在（跳过）。
|
||
"""
|
||
card = l3.card
|
||
current = getattr(card, field, None)
|
||
|
||
if current is None:
|
||
logger.debug("L3Card 无字段 {}，跳过注入", field)
|
||
return False
|
||
|
||
if isinstance(current, list):
|
||
lower_set = {v.strip().lower() for v in current if isinstance(v, str)}
|
||
if value.strip().lower() in lower_set:
|
||
return False
|
||
new_list = deduplicate_field([*current, value])
|
||
l3.card = replace(card, **{field: new_list})
|
||
return True
|
||
|
||
if isinstance(current, str):
|
||
if value.strip().lower() in current.lower():
|
||
return False
|
||
new_val = current + "; " + value if current else value
|
||
l3.card = replace(card, **{field: new_val})
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def _inject_into_l1(l1: L1Node, field: str, value: str) -> bool:
|
||
"""向 L1 节点的 Card 指定字段注入一个值。
|
||
|
||
使用 dataclasses.replace() 创建新的 frozen L1Card。
|
||
|
||
参数:
|
||
l1: L1 节点（card 会被替换为新实例）。
|
||
field: 目标字段名。
|
||
value: 要注入的值。
|
||
|
||
返回:
|
||
True 表示实际注入了新内容，False 表示已存在（跳过）。
|
||
"""
|
||
card = l1.card
|
||
current = getattr(card, field, None)
|
||
|
||
if current is None:
|
||
logger.debug("L1Card 无字段 {}，跳过注入", field)
|
||
return False
|
||
|
||
if isinstance(current, list):
|
||
lower_set = {v.strip().lower() for v in current if isinstance(v, str)}
|
||
if value.strip().lower() in lower_set:
|
||
return False
|
||
new_list = deduplicate_field([*current, value])
|
||
l1.card = replace(card, **{field: new_list})
|
||
return True
|
||
|
||
if isinstance(current, str):
|
||
if value.strip().lower() in current.lower():
|
||
return False
|
||
new_val = current + "; " + value if current else value
|
||
l1.card = replace(card, **{field: new_val})
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 批量注入
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def apply_injections(index: TreeIndex, injections: list[dict[str, Any]]) -> SupplementStats:
|
||
"""执行一组注入指令，将事实写入树节点 Card。
|
||
|
||
每条指令格式::
|
||
|
||
{
|
||
"category": "person_name" | "location" | "score_number" | "object_name",
|
||
"inject_value": "...",
|
||
"targets": [{"node_id": "...", "field": "..."}, ...]
|
||
}
|
||
|
||
向后兼容: 若无 targets，读取 target_node_id + target_field 构造单目标。
|
||
|
||
参数:
|
||
index: TreeIndex 实例（节点 Card 会被替换为新实例）。
|
||
injections: 注入指令列表。
|
||
|
||
返回:
|
||
注入统计信息。
|
||
"""
|
||
stats = SupplementStats()
|
||
|
||
for instr in injections:
|
||
category = instr.get("category", "")
|
||
if category not in _ALLOWED_CATEGORIES:
|
||
logger.debug("拒绝非法类别: {}", category)
|
||
stats.facts_skipped += 1
|
||
continue
|
||
|
||
inject_value = instr.get("inject_value", "")
|
||
if not inject_value:
|
||
stats.facts_skipped += 1
|
||
continue
|
||
|
||
# 解析目标列表（兼容新旧格式）
|
||
targets = instr.get("targets")
|
||
if not targets:
|
||
node_id = instr.get("target_node_id", "")
|
||
field = instr.get("target_field", "")
|
||
if node_id and field:
|
||
targets = [{"node_id": node_id, "field": field}]
|
||
else:
|
||
stats.facts_skipped += 1
|
||
continue
|
||
|
||
for target in targets:
|
||
node_id = target.get("node_id", "")
|
||
field = target.get("field", "")
|
||
node, level = _find_node_by_id(index, node_id)
|
||
|
||
if node is None:
|
||
logger.debug("跳过不存在的节点: {}", node_id)
|
||
stats.facts_skipped += 1
|
||
continue
|
||
|
||
injected = False
|
||
if level == 1:
|
||
injected = _inject_into_l1(node, field, inject_value) # type: ignore[arg-type]
|
||
elif level == 2:
|
||
injected = _inject_into_l2(node, field, inject_value) # type: ignore[arg-type]
|
||
elif level == 3:
|
||
injected = _inject_into_l3(node, field, inject_value) # type: ignore[arg-type]
|
||
|
||
if injected:
|
||
stats.facts_injected += 1
|
||
else:
|
||
stats.facts_skipped += 1
|
||
|
||
return stats
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LLM Prompt
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_SUPPLEMENT_SYSTEM_PROMPT = """\
|
||
你是一个视频内容分析专家。你的任务是分析回答某个问题需要哪些关键事实，
|
||
并判断这些事实是否已存在于视频树的摘要中。
|
||
|
||
## 输出规则
|
||
|
||
1. 只输出**客观事实**，包括以下四类:
|
||
- person_name: 人物姓名
|
||
- location: 地点名称
|
||
- score_number: 比分、数字
|
||
- object_name: 关键物体名称
|
||
|
||
2. **不要**输出以下类型:
|
||
- 情感、态度、心情
|
||
- 因果推理（"因为…所以…"）
|
||
- 时间顺序推理（"先…后…"）
|
||
- 主观评价
|
||
|
||
3. 对于 person_name 类别，输出 targets 数组包含两个写入点:
|
||
- L2 节点的 entities 字段
|
||
- L3 节点的 visible_entities 字段
|
||
其他类别只写入最相关的单个节点的 entities 字段。
|
||
|
||
4. 每条 missing fact 必须包含 inject_value（要注入的值）和 targets 数组。
|
||
|
||
## 输出格式 (严格 JSON)
|
||
|
||
```json
|
||
{
|
||
"needed_facts": [
|
||
{"category": "person_name", "value": "..."}
|
||
],
|
||
"found_in_tree": [
|
||
{"category": "person_name", "value": "...", "found_at": "node_id"}
|
||
],
|
||
"missing_facts": [
|
||
{
|
||
"category": "person_name",
|
||
"inject_value": "...",
|
||
"targets": [
|
||
{"node_id": "...", "field": "entities"},
|
||
{"node_id": "...", "field": "visible_entities"}
|
||
]
|
||
}
|
||
]
|
||
}
|
||
```
|
||
|
||
只输出 JSON，不要输出其他内容。
|
||
"""
|
||
|
||
|
||
def _build_user_prompt(
|
||
question: dict[str, Any],
|
||
index: TreeIndex,
|
||
srt_text: str,
|
||
) -> str:
|
||
"""构建 supplement 分析的 user prompt。
|
||
|
||
包含: 问题 + 选项 + 正确答案 + 树 L2 摘要 + SRT 字幕（截断至 3000 字符）。
|
||
|
||
参数:
|
||
question: 包含 question/options/answer 的字典。
|
||
index: TreeIndex 实例。
|
||
srt_text: SRT 字幕文本。
|
||
|
||
返回:
|
||
拼装后的 user prompt 字符串。
|
||
"""
|
||
# 问题部分
|
||
q_text = question.get("question", "")
|
||
options = question.get("options", [])
|
||
answer = question.get("answer", "")
|
||
options_str = "\n".join(f" {chr(65 + i)}. {opt}" for i, opt in enumerate(options))
|
||
|
||
# 树 L2 摘要（从 TreeIndex 结构中提取）
|
||
l2_summaries: list[str] = []
|
||
for l1 in index.roots:
|
||
for l2 in l1.children:
|
||
description = l2.card.event_description
|
||
entities_str = ", ".join(l2.card.entities) if l2.card.entities else ""
|
||
time_str = ""
|
||
if l2.time_range:
|
||
time_str = f"{l2.time_range[0]:.1f}-{l2.time_range[1]:.1f}s: "
|
||
l2_summaries.append(
|
||
f"[{l2.id}] {time_str}{description}"
|
||
+ (f" | entities: {entities_str}" if entities_str else "")
|
||
)
|
||
|
||
l2_block = "\n".join(l2_summaries) if l2_summaries else "(无 L2 摘要)"
|
||
|
||
# SRT 截断
|
||
srt_truncated = srt_text[:3000] if srt_text else "(无字幕)"
|
||
|
||
return (
|
||
f"## 问题\n{q_text}\n\n"
|
||
f"## 选项\n{options_str}\n\n"
|
||
f"## 正确答案\n{answer}\n\n"
|
||
f"## 视频树 L2 摘要\n{l2_block}\n\n"
|
||
f"## 字幕 (前 3000 字符)\n{srt_truncated}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LLM 调用
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def analyze_question(
|
||
llm: LLMProvider,
|
||
question: dict[str, Any],
|
||
index: TreeIndex,
|
||
srt_text: str,
|
||
) -> list[dict[str, Any]]:
|
||
"""调用 LLM 分析单个问题，返回需要注入的事实列表。
|
||
|
||
参数:
|
||
llm: LLMProvider 实例（异步接口）。
|
||
question: 问题字典（含 question/options/answer）。
|
||
index: TreeIndex 实例。
|
||
srt_text: SRT 字幕文本。
|
||
|
||
返回:
|
||
missing_facts 列表，每项含 category / inject_value / targets。
|
||
解析失败时返回空列表。
|
||
"""
|
||
user_prompt = _build_user_prompt(question, index, srt_text)
|
||
messages = [
|
||
{"role": "system", "content": _SUPPLEMENT_SYSTEM_PROMPT},
|
||
{"role": "user", "content": user_prompt},
|
||
]
|
||
|
||
response = await llm.chat(messages)
|
||
raw = response.content
|
||
|
||
# 提取 JSON（兼容 markdown 代码块包裹）
|
||
text = raw.strip()
|
||
if text.startswith("```"):
|
||
lines = text.split("\n")
|
||
lines = [ln for ln in lines if not ln.strip().startswith("```")]
|
||
text = "\n".join(lines)
|
||
|
||
try:
|
||
parsed = json.loads(text)
|
||
except json.JSONDecodeError:
|
||
logger.warning("supplement LLM 返回非法 JSON，跳过。原始内容: {}", raw[:200])
|
||
return []
|
||
|
||
missing = parsed.get("missing_facts", [])
|
||
if not isinstance(missing, list):
|
||
logger.warning("missing_facts 不是列表，跳过")
|
||
return []
|
||
|
||
return missing
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 主入口
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def supplement_tree(
|
||
index: TreeIndex,
|
||
questions: list[dict[str, Any]],
|
||
llm: LLMProvider,
|
||
srt_text: str = "",
|
||
) -> SupplementStats:
|
||
"""对树索引执行 Q&A 反向补全：遍历问题，分析缺失事实，注入节点。
|
||
|
||
参数:
|
||
index: TreeIndex 实例（节点 Card 会被就地替换）。
|
||
questions: 问题列表，每项含 question/options/answer。
|
||
llm: LLMProvider 实例（异步接口）。
|
||
srt_text: SRT 字幕文本（可选，默认空字符串）。
|
||
|
||
返回:
|
||
补全统计信息。
|
||
"""
|
||
all_injections: list[dict[str, Any]] = []
|
||
|
||
for i, question in enumerate(questions):
|
||
logger.debug(
|
||
"supplement: 分析问题 {}/{}",
|
||
i + 1,
|
||
len(questions),
|
||
)
|
||
missing = await analyze_question(llm, question, index, srt_text)
|
||
all_injections.extend(missing)
|
||
|
||
stats = apply_injections(index, all_injections)
|
||
stats.questions_analyzed = len(questions)
|
||
|
||
logger.info(
|
||
"supplement_tree 完成: questions={} injections={} injected={} skipped={}",
|
||
len(questions),
|
||
len(all_injections),
|
||
stats.facts_injected,
|
||
stats.facts_skipped,
|
||
)
|
||
return stats
|