Video-Tree-TRM5/app/question_gen/strategy.py

"""题型出题策略 — Clean Architecture 的 Strategy 层。

将出题管线从 5 个粗粒度 QuestionFamilySpec 替换为 12 个题型级别的
TaskTypeStrategy。BaseTaskTypeStrategy 封装现有 family 行为，
特化策略（如 ActionRecognitionStrategy）自包含。

典型用法::

    strategy = get_strategy("Action Reasoning")
    level = strategy.sampling_level
    constraint = strategy.sampling_constraint
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

from app.question_gen.families import (
    ENUMERATION_FAMILY,
    REASONING_FAMILY,
    RETRIEVAL_FAMILY,
    SPATIAL_FAMILY,
    VISUAL_FAMILY,
    QuestionFamilySpec,
    SamplingConstraint,
)
from app.question_gen.sampler_v2 import _TASK_TYPE_TO_LEVEL

if TYPE_CHECKING:
    import random


# ---------------------------------------------------------------------------
# SubPattern
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class SubPattern:
    """出题子模式 — 靶向特定失败机制。

    属性:
        name: 子模式标识名。
        weight: 选择权重（按比例分配）。
        sampling_level_override: 覆盖策略默认采样层级（None = 不覆盖）。
        constraint_override: 覆盖策略默认采样约束（None = 不覆盖）。
        instruction: 核心出题指令（1-3 句话）。
        positive_examples: VME 原题 few-shot 示范。
        negative_examples: 反面示例。
        distractor_rules: 干扰项构造规则。
    """

    name: str
    weight: float
    sampling_level_override: int | None
    constraint_override: SamplingConstraint | None
    instruction: str
    positive_examples: list[dict] = field(default_factory=list)
    negative_examples: list[dict] = field(default_factory=list)
    distractor_rules: str = ""


# ---------------------------------------------------------------------------
# Protocol
# ---------------------------------------------------------------------------


@runtime_checkable
class TaskTypeStrategy(Protocol):
    """题型出题策略 — pipeline 的唯一接口。"""

    @property
    def task_type(self) -> str: ...

    @property
    def sampling_level(self) -> int: ...

    @property
    def sampling_constraint(self) -> SamplingConstraint: ...

    @property
    def prompt_template(self) -> str: ...

    @property
    def strategy_name(self) -> str: ...

    @property
    def skill_target(self) -> str: ...

    @property
    def leak_probe_template(self) -> str: ...

    def select_sub_pattern(self, rng: random.Random) -> SubPattern | None: ...

    def build_prompt_context(self, material: Any, sub_pattern: SubPattern | None) -> dict: ...

    def extra_gates(self, candidate: Any) -> list: ...


# ---------------------------------------------------------------------------
# BaseTaskTypeStrategy
# ---------------------------------------------------------------------------

# 消歧绑定表：多归属题型确定性绑定到一个 family
_TASK_TYPE_TO_FAMILY: dict[str, QuestionFamilySpec] = {
    "Action Recognition": VISUAL_FAMILY,  # Plan A 临时绑定；Plan B 替换为特化策略
    "Object Recognition": RETRIEVAL_FAMILY,
    "Object Reasoning": REASONING_FAMILY,
    "Action Reasoning": REASONING_FAMILY,
    "Attribute Perception": VISUAL_FAMILY,
    "OCR Problems": VISUAL_FAMILY,
    "Counting Problem": ENUMERATION_FAMILY,
    "Information Synopsis": REASONING_FAMILY,
    "Temporal Reasoning": ENUMERATION_FAMILY,
    "Temporal Perception": ENUMERATION_FAMILY,
    "Spatial Reasoning": SPATIAL_FAMILY,
    "Spatial Perception": SPATIAL_FAMILY,
}


class BaseTaskTypeStrategy:
    """封装现有 QuestionFamilySpec 行为的默认策略。

    所有属性委托给绑定的 family，确保未特化题型的行为不变。

    参数:
        task_type: 题型名。
        family: 绑定的 QuestionFamilySpec。
    """

    def __init__(self, task_type: str, family: QuestionFamilySpec) -> None:
        self._task_type = task_type
        self._family = family

    @property
    def task_type(self) -> str:
        """返回题型名。"""
        return self._task_type

    @property
    def sampling_level(self) -> int:
        """返回采样层级（从 _TASK_TYPE_TO_LEVEL 查询）。"""
        return _TASK_TYPE_TO_LEVEL[self._task_type]

    @property
    def sampling_constraint(self) -> SamplingConstraint:
        """返回采样约束（委托给绑定的 family）。"""
        return self._family.sampling

    @property
    def prompt_template(self) -> str:
        """返回 prompt 模板文件名（委托给绑定的 family）。"""
        return self._family.prompt_template

    @property
    def strategy_name(self) -> str:
        """返回策略名（即 family.name）。"""
        return self._family.name

    @property
    def skill_target(self) -> str:
        """返回目标失败机制编号（委托给绑定的 family）。"""
        return self._family.skill_target

    @property
    def leak_probe_template(self) -> str:
        """返回泄漏探测模板文件名（委托给绑定的 family.leak_profile）。"""
        return self._family.leak_profile.probe_template

    def select_sub_pattern(self, rng: random.Random) -> SubPattern | None:
        """BaseTaskTypeStrategy 无子模式。"""
        return None

    def build_prompt_context(self, material: Any, sub_pattern: SubPattern | None) -> dict:
        """返回基础 prompt 上下文（family_name + prompt_template）。"""
        return {
            "family_name": self._family.name,
            "prompt_template": self._family.prompt_template,
        }

    def extra_gates(self, candidate: Any) -> list:
        """BaseTaskTypeStrategy 无额外 gate。"""
        return []


# ---------------------------------------------------------------------------
# 注册表
# ---------------------------------------------------------------------------

_STRATEGY_REGISTRY: dict[str, TaskTypeStrategy] = {}


def register_strategy(strategy: TaskTypeStrategy) -> None:
    """注册一个题型策略。同一 task_type 重复注册会覆盖。

    参数:
        strategy: 实现 TaskTypeStrategy 接口的策略实例。
    """
    _STRATEGY_REGISTRY[strategy.task_type] = strategy


def get_strategy(task_type: str) -> TaskTypeStrategy:
    """获取题型策略。未注册的自动创建 BaseTaskTypeStrategy。

    参数:
        task_type: 题型名。

    返回:
        TaskTypeStrategy 实例。

    异常:
        KeyError: task_type 不在消歧绑定表和注册表中。
    """
    if task_type in _STRATEGY_REGISTRY:
        return _STRATEGY_REGISTRY[task_type]
    return _build_default_strategy(task_type)


def _build_default_strategy(task_type: str) -> BaseTaskTypeStrategy:
    """根据消歧绑定表自动构造 BaseTaskTypeStrategy。

    参数:
        task_type: 题型名。

    返回:
        BaseTaskTypeStrategy 实例。

    异常:
        KeyError: task_type 不在消歧绑定表中。
    """
    family = _TASK_TYPE_TO_FAMILY[task_type]
    return BaseTaskTypeStrategy(task_type=task_type, family=family)