iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(question_gen): add v2 pipeline with retry loop and heavy check

Browse Source 
- PipelineConfig: YAML-driven configuration with family_ratios, retry,
  concurrency, dedup threshold, and heavy sampling rate
- _assign_slots: deterministic round-robin slot assignment across videos
  with per-family weighted random selection
- _process_one_slot: full retry loop (generate → postprocess → gates →
  dedup) with reject-reason feedback to VLM on retry
- _heavy_check_one: blind LLM agent trial-answer for difficulty_steps
- run_pipeline_v2: orchestration with semaphore-bounded concurrency,
  progress/resume support, and store integration
- is_duplicate: cosine similarity dedup against embedding pool

Tests: 11 integration tests covering slot assignment, retry behavior,
max-retries exhaustion, full pipeline flow, progress resume, heavy
sampling, and store record completeness.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
iomgaa
2026-07-11 23:54:21 -04:00
parent 6d6eb8e3a3
commit 206c553143
2 changed files with 1385 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/test_pipeline_v2.py
+650
View File
@@ -0,0 +1,650 @@
"""v2 出题管线集成测试 — 覆盖 slot 分配、重出循环、重量抽检与完整编排。
测试策略：使用受控 mock VLM/LLM 返回，验证管线逻辑正确性。
"""
from __future__ import annotations
import asyncio
import json
import random
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from pathlib import Path
import numpy as np
import pytest
from app.question_gen.families import ALL_FAMILIES
from app.question_gen.pipeline_v2 import (
PipelineConfig,
PipelineResult,
SlotAssignment,
_assign_slots,
_process_one_slot,
run_pipeline_v2,
)
from app.tree.index import (
IndexMeta,
L1Card,
L1Node,
L2Card,
L2Node,
L3Card,
L3Node,
TreeIndex,
)
from core.types import LLMResponse
# ---------------------------------------------------------------------------
# 测试用 fixtures
# ---------------------------------------------------------------------------
def _make_llm_response(content: str) -> LLMResponse:
"""构造标准 LLMResponse。"""
return LLMResponse(
content=content,
thinking="",
model="test-model",
provider="test",
prompt_tokens=10,
completion_tokens=20,
latency_ms=100,
ttft_ms=50.0,
max_inter_token_ms=10.0,
cache_hit=False,
call_id="call-001",
)
def _make_gate_pass_response() -> str:
"""门控全部通过的 JSON 响应。"""
return json.dumps({"verdict": "pass", "reason": "looks good"})
def _make_gate_fail_response(reason: str = "quality issue") -> str:
"""门控失败的 JSON 响应。"""
return json.dumps({"verdict": "fail", "reason": reason})
def _make_candidate_json(
question: str = "What happened next?",
answer: str = "A",
) -> str:
"""构造 VLM 返回的候选题 JSON。"""
return json.dumps(
{
"question": question,
"options": [
"A. The cat jumped",
"B. The dog ran",
"C. Nothing happened",
"D. It rained",
],
"answer": answer,
"difficulty": "medium",
}
)
def _make_tree() -> TreeIndex:
"""构造最小合法三层树。"""
l3_nodes = [
L3Node(
id=f"vid_L1_000_L2_000_L3_{i:03d}",
card=L3Card(
frame_summary=f"Frame {i} shows activity",
visible_entities=["person", "object"],
ongoing_actions=["walking"],
visible_text=[],
spatial_layout="center",
visual_attributes={},
subtitle=f"Subtitle sentence {i} with some unique content here",
),
timestamp=float(i * 2),
frame_path=f"frames/L1_000_L2_000_L3_{i:03d}.jpg",
)
for i in range(5)
]
l2 = L2Node(
id="vid_L1_000_L2_000",
card=L2Card(
event_description="A person walks through the park",
entities=["person", "park"],
actions=["walking"],
action_subjects=["person"],
visible_text=[],
spatial_relations="person in center of park",
state_changes=None,
subtitle="Person walking in park doing activities",
),
children=l3_nodes,
)
l2_b = L2Node(
id="vid_L1_000_L2_001",
card=L2Card(
event_description="A dog runs across the field",
entities=["dog", "field"],
actions=["running"],
action_subjects=["dog"],
visible_text=[],
spatial_relations="dog in the field",
state_changes=None,
subtitle="Dog running across the field",
),
children=[
L3Node(
id=f"vid_L1_000_L2_001_L3_{i:03d}",
card=L3Card(
frame_summary=f"Dog frame {i}",
visible_entities=["dog"],
ongoing_actions=["running"],
visible_text=[],
spatial_layout="wide",
visual_attributes={},
subtitle=f"Dog subtitle {i}",
),
timestamp=float(10 + i * 2),
frame_path=f"frames/L1_000_L2_001_L3_{i:03d}.jpg",
)
for i in range(4)
],
)
l1 = L1Node(
id="vid_L1_000",
card=L1Card(
scene_summary="Outdoor activities in a park",
main_setting="outdoor park",
key_entities=["person", "dog"],
main_actions=["walking", "running"],
topic_keywords=["outdoor", "activity"],
visible_text=[],
temporal_flow="sequential activities",
),
children=[l2, l2_b],
)
meta = IndexMeta(source_path="test_video.mp4", modality="video")
return TreeIndex(metadata=meta, roots=[l1])
class MockVLM:
"""受控 VLM mock — 每次调用返回候选题 JSON。"""
def __init__(self, responses: list[str] | None = None) -> None:
self._responses = responses or [_make_candidate_json()]
self._call_count = 0
async def chat_with_images(
self,
messages: list[dict[str, Any]],
images: list[str | Path],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
idx = min(self._call_count, len(self._responses) - 1)
self._call_count += 1
return _make_llm_response(self._responses[idx])
class MockLLM:
"""受控 LLM mock — 支持配置门控 pass/fail 序列。"""
def __init__(self, responses: list[str] | None = None) -> None:
self._responses = responses or [_make_gate_pass_response()]
self._call_count = 0
async def chat(
self,
messages: list[dict[str, Any]],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
idx = min(self._call_count, len(self._responses) - 1)
self._call_count += 1
return _make_llm_response(self._responses[idx])
def _mock_embed_fn(text: str) -> np.ndarray:
"""确定性 embedding：基于文本 hash 生成向量。"""
rng = np.random.default_rng(hash(text) % (2**32))
vec = rng.standard_normal(64).astype(np.float32)
return vec / np.linalg.norm(vec)
@pytest.fixture
def tree() -> TreeIndex:
return _make_tree()
@pytest.fixture
def default_config(tmp_path: Path) -> PipelineConfig:
return PipelineConfig(
family_ratios={
"RETRIEVAL": 0.30,
"REASONING": 0.25,
"ENUMERATION": 0.20,
"VISUAL": 0.15,
"SPATIAL": 0.10,
},
per_type=2,
retry_limit=3,
heavy_sample_rate=0.15,
dedup_threshold=0.85,
concurrency=2,
seed=42,
output_dir=tmp_path / "output",
gate_models={
"blind_answer_model": "gpt-4.1-mini",
"leak_test_model": "gpt-4.1-mini",
"key_verify_model": "gpt-4.1-mini",
"multi_true_model": "gpt-4.1-mini",
},
heavy_agent_model="gpt-4.1-mini",
)
@pytest.fixture
def store(tmp_path: Path):
from app.question_gen.run_store import QuestionGenStore
db_path = tmp_path / "test_qgen.db"
s = QuestionGenStore(db_path=db_path)
yield s
s.close()
# ---------------------------------------------------------------------------
# TestSlotAssignment
# ---------------------------------------------------------------------------
class TestSlotAssignment:
"""_assign_slots 单元测试。"""
def test_per_type_count(self):
"""验证生成的 slot 总数 = len(task_types) * per_type。"""
video_ids = ["vid_001", "vid_002"]
task_types = ["Action Recognition", "Object Recognition", "Causal Reasoning"]
per_type = 4
family_ratios = {"RETRIEVAL": 0.5, "VISUAL": 0.5}
rng = random.Random(42)
slots = _assign_slots(video_ids, task_types, per_type, family_ratios, rng)
assert len(slots) == len(task_types) * per_type
def test_family_distribution(self):
"""验证家族分配来自 get_family_for_slot（合法 family 与 task_type 匹配）。"""
video_ids = ["vid_001", "vid_002", "vid_003"]
task_types = ["Action Recognition"]
per_type = 20
family_ratios = {"RETRIEVAL": 0.5, "VISUAL": 0.5}
rng = random.Random(42)
slots = _assign_slots(video_ids, task_types, per_type, family_ratios, rng)
for slot in slots:
assert slot.task_type == "Action Recognition"
# Family 必须是 task_type 合法的家族之一
family = slot.family
assert slot.task_type in family.legal_task_types
def test_round_robin_across_videos(self):
"""验证 slot 在视频间轮转分配。"""
video_ids = ["vid_A", "vid_B"]
task_types = ["Action Recognition"]
per_type = 4
family_ratios = {"RETRIEVAL": 1.0}
rng = random.Random(42)
slots = _assign_slots(video_ids, task_types, per_type, family_ratios, rng)
video_assignments = [s.video_id for s in slots]
# 应该轮转分配
assert video_assignments.count("vid_A") == 2
assert video_assignments.count("vid_B") == 2
def test_deterministic_with_seed(self):
"""相同 seed 产出相同 slot 序列。"""
video_ids = ["vid_001", "vid_002"]
task_types = ["Action Recognition", "Causal Reasoning"]
per_type = 3
family_ratios = {"RETRIEVAL": 0.5, "REASONING": 0.5}
rng1 = random.Random(99)
slots1 = _assign_slots(video_ids, task_types, per_type, family_ratios, rng1)
rng2 = random.Random(99)
slots2 = _assign_slots(video_ids, task_types, per_type, family_ratios, rng2)
for s1, s2 in zip(slots1, slots2, strict=True):
assert s1.slot_id == s2.slot_id
assert s1.video_id == s2.video_id
assert s1.family.name == s2.family.name
# ---------------------------------------------------------------------------
# TestProcessOneSlot
# ---------------------------------------------------------------------------
class TestProcessOneSlot:
"""_process_one_slot 集成测试。"""
@pytest.mark.asyncio
async def test_happy_path(self, tree, default_config, store, tmp_path):
"""首次生成即通过门控 → 返回 GeneratedQuestion。"""
vlm = MockVLM()
llm = MockLLM([_make_gate_pass_response()] * 4)
sem = asyncio.Semaphore(2)
used_node_ids: set[str] = set()
rng = random.Random(42)
family = ALL_FAMILIES[0] # RETRIEVAL
slot = SlotAssignment(
slot_id="slot_001",
video_id="vid_L1_000",
task_type="Action Recognition",
family=family,
seq=1,
)
run_id = "test-run-001"
store.record_run_start(run_id, "abc123", "{}")
result = await _process_one_slot(
slot=slot,
tree=tree,
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
embed_pool=[],
store=store,
config=default_config,
used_node_ids=used_node_ids,
rng=rng,
sem=sem,
session_id="sess-001",
run_id=run_id,
)
assert result is not None
assert result.question_id
assert result.skill_target == family.skill_target
@pytest.mark.asyncio
async def test_retry_on_fail(self, tree, default_config, store, tmp_path):
"""第一次门控失败，第二次通过 → 重出成功。"""
vlm = MockVLM(
[
_make_candidate_json("First question?"),
_make_candidate_json("Second better question?"),
]
)
# 第一轮 4 个门有一个 fail，第二轮 4 个门全 pass
llm_responses = [
_make_gate_fail_response("answer not grounded"),
_make_gate_pass_response(),
_make_gate_pass_response(),
_make_gate_pass_response(),
# 第二轮
_make_gate_pass_response(),
_make_gate_pass_response(),
_make_gate_pass_response(),
_make_gate_pass_response(),
]
llm = MockLLM(llm_responses)
sem = asyncio.Semaphore(2)
used_node_ids: set[str] = set()
rng = random.Random(42)
family = ALL_FAMILIES[0] # RETRIEVAL
slot = SlotAssignment(
slot_id="slot_002",
video_id="vid_L1_000",
task_type="Action Recognition",
family=family,
seq=2,
)
run_id = "test-run-002"
store.record_run_start(run_id, "abc123", "{}")
result = await _process_one_slot(
slot=slot,
tree=tree,
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
embed_pool=[],
store=store,
config=default_config,
used_node_ids=used_node_ids,
rng=rng,
sem=sem,
session_id="sess-002",
run_id=run_id,
)
assert result is not None
# 确认第二个问题被接受
assert "Second" in result.question or result.question_id is not None
@pytest.mark.asyncio
async def test_max_retries_none(self, tree, default_config, store, tmp_path):
"""所有重试均失败 → 返回 None。"""
vlm = MockVLM([_make_candidate_json()] * 5)
# 所有门控均失败
llm = MockLLM([_make_gate_fail_response("always fails")] * 20)
sem = asyncio.Semaphore(2)
used_node_ids: set[str] = set()
rng = random.Random(42)
family = ALL_FAMILIES[0] # RETRIEVAL
slot = SlotAssignment(
slot_id="slot_003",
video_id="vid_L1_000",
task_type="Action Recognition",
family=family,
seq=3,
)
run_id = "test-run-003"
store.record_run_start(run_id, "abc123", "{}")
result = await _process_one_slot(
slot=slot,
tree=tree,
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
embed_pool=[],
store=store,
config=default_config,
used_node_ids=used_node_ids,
rng=rng,
sem=sem,
session_id="sess-003",
run_id=run_id,
)
assert result is None
# ---------------------------------------------------------------------------
# TestPipelineV2
# ---------------------------------------------------------------------------
class TestPipelineV2:
"""run_pipeline_v2 完整流程测试。"""
@pytest.mark.asyncio
async def test_full_flow(self, tree, default_config, store, tmp_path):
"""完整管线运行，产出 PipelineResult。"""
vlm = MockVLM([_make_candidate_json(f"Question {i}?") for i in range(50)])
llm = MockLLM([_make_gate_pass_response()] * 200)
config = PipelineConfig(
family_ratios=default_config.family_ratios,
per_type=2,
retry_limit=2,
heavy_sample_rate=0.5, # 高比例便于测试
dedup_threshold=0.85,
concurrency=2,
seed=42,
output_dir=tmp_path / "out",
gate_models=default_config.gate_models,
heavy_agent_model="gpt-4.1-mini",
)
# 只用一个 task_type 确保树能满足采样
result = await run_pipeline_v2(
video_ids=["vid_L1_000"],
trees={"vid_L1_000": tree},
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
store=store,
config=config,
task_types=["Action Recognition"],
)
assert isinstance(result, PipelineResult)
assert result.rejected_count >= 0
# 至少一题被接受（VLM 和 LLM 全部正常返回）
assert len(result.accepted) > 0
@pytest.mark.asyncio
async def test_progress_resume(self, tree, default_config, store, tmp_path):
"""传入 progress dict → 已处理 slot 被跳过。"""
llm = MockLLM([_make_gate_pass_response()] * 100)
config = PipelineConfig(
family_ratios=default_config.family_ratios,
per_type=2,
retry_limit=2,
heavy_sample_rate=0.0, # 不做 heavy check
dedup_threshold=0.85,
concurrency=2,
seed=42,
output_dir=tmp_path / "out",
gate_models=default_config.gate_models,
heavy_agent_model="gpt-4.1-mini",
)
# 用相同 seed 计算 slot_ids（_assign_slots 是确定性的）
rng_preview = random.Random(config.seed)
preview_slots = _assign_slots(
["vid_L1_000"],
["Action Recognition"],
config.per_type,
config.family_ratios,
rng_preview,
)
# 构造 progress 标记所有 slot 已完成
progress = {s.slot_id: "accepted" for s in preview_slots}
# 用 progress 跑管线 → 所有 slot 被跳过
vlm2 = MockVLM([_make_candidate_json()] * 20)
result2 = await run_pipeline_v2(
video_ids=["vid_L1_000"],
trees={"vid_L1_000": tree},
vlm=vlm2,
llm=llm,
embed_fn=_mock_embed_fn,
store=store,
config=config,
task_types=["Action Recognition"],
progress=progress,
)
# 所有 slot 在 progress 中 → VLM 零调用
assert vlm2._call_count == 0
assert len(result2.accepted) == 0
@pytest.mark.asyncio
async def test_heavy_check_samples(self, tree, default_config, store, tmp_path):
"""heavy_sample_rate > 0 时有题被抽检。"""
# LLM 响应：前面是 gate pass，后面增加 heavy check 的步骤响应
heavy_response = json.dumps(
{
"steps": [
{"thought": "step 1"},
{"thought": "step 2"},
{"thought": "step 3"},
],
"answer": "A",
}
)
llm = MockLLM([_make_gate_pass_response()] * 100 + [heavy_response] * 20)
vlm = MockVLM([_make_candidate_json(f"Q{i}?") for i in range(20)])
config = PipelineConfig(
family_ratios=default_config.family_ratios,
per_type=2,
retry_limit=2,
heavy_sample_rate=1.0, # 100% 抽检
dedup_threshold=0.85,
concurrency=2,
seed=42,
output_dir=tmp_path / "out",
gate_models=default_config.gate_models,
heavy_agent_model="gpt-4.1-mini",
)
result = await run_pipeline_v2(
video_ids=["vid_L1_000"],
trees={"vid_L1_000": tree},
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
store=store,
config=config,
task_types=["Action Recognition"],
)
# 100% 抽检 → heavy_sampled 等于 accepted 数
if result.accepted:
assert len(result.heavy_sampled) == len(result.accepted)
@pytest.mark.asyncio
async def test_store_records_all(self, tree, default_config, store, tmp_path):
"""验证 store 中记录了每道题的生成与门控结果。"""
vlm = MockVLM([_make_candidate_json()] * 10)
llm = MockLLM([_make_gate_pass_response()] * 50)
config = PipelineConfig(
family_ratios=default_config.family_ratios,
per_type=2,
retry_limit=2,
heavy_sample_rate=0.0,
dedup_threshold=0.85,
concurrency=1,
seed=42,
output_dir=tmp_path / "out",
gate_models=default_config.gate_models,
heavy_agent_model="gpt-4.1-mini",
)
result = await run_pipeline_v2(
video_ids=["vid_L1_000"],
trees={"vid_L1_000": tree},
vlm=vlm,
llm=llm,
embed_fn=_mock_embed_fn,
store=store,
config=config,
task_types=["Action Recognition"],
)
# 查询 store 中的 items
cursor = store._conn.execute("SELECT COUNT(*) FROM question_gen_items")
item_count = cursor.fetchone()[0]
# 至少有 accepted 数量的记录
assert item_count >= len(result.accepted)