Video-Tree-TRM5/app/harness/factory.py

"""推理依赖工厂 — 组装一次推理所需的全套依赖。

将 TreeIndex 加载、TreeEnvironment 构建、SkillRegistry 发现、
SearchToolDispatcher 装配、PromptManager 初始化等步骤封装为
单一工厂函数 ``build_inference_deps``，返回不可变的 ``InferenceDeps``。

调用方（runner / inference）只需传入配置参数，无需了解内部装配逻辑。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from loguru import logger

from app.search.prompt import PromptManager
from app.search.skills import discover_skills
from app.search.tools import SearchToolDispatcher
from app.tree.environment import TreeEnvironment
from app.tree.index import TreeIndex

if TYPE_CHECKING:
    from collections.abc import Callable
    from pathlib import Path

    from app.ports import EmbeddingProvider, OCRProvider
    from core.protocols import LLMProvider, VLMProvider
    from core.types import GeneratedQuestion


@dataclass(frozen=True)
class InferenceDeps:
    """跑一次推理所需的全套依赖（不含 HarnessLog，其生命周期由调用方管理）。

    属性:
        llm: LLM 端口实例。
        tool_dispatch_fn: SearchToolDispatcher.dispatch 的绑定方法。
        prompt_builder: (GeneratedQuestion) -> (system_prompt, user_prompt)。
    """

    llm: LLMProvider
    tool_dispatch_fn: Callable[..., Any]
    prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]]


def build_inference_deps(
    *,
    store_dir: Path,
    video_id: str,
    prompts_dir: Path,
    skills_dir: Path | None,
    skill_mode: str,
    embed_provider: EmbeddingProvider,
    llm: LLMProvider,
    vlm: VLMProvider,
    ocr: OCRProvider | None,
    verify_vision: bool,
    anchor: bool,
    assemble_mode: str,
) -> InferenceDeps:
    """组装一次推理所需的全套依赖。

    参数:
        store_dir: store 根目录（包含 videos/{video_id}/tree.json）。
        video_id: 视频标识。
        prompts_dir: prompt 文件目录。
        skills_dir: skill 文件目录（None 则不加载 skill）。
        skill_mode: skill 模式（"auto"/"manual"/"none"）。
        embed_provider: 嵌入端口实例。
        llm: LLM 端口实例。
        vlm: VLM 端口实例。
        ocr: OCR 端口实例（None 不启用）。
        verify_vision: observe_frame 是否执行验证轮。
        anchor: view_node 是否启用行号锚模式。
        assemble_mode: 锚模式装配形态。

    返回:
        InferenceDeps 实例。

    异常:
        FileNotFoundError: tree.json 不存在。
    """
    # Phase 1: 加载 TreeIndex
    tree_path = store_dir / "videos" / video_id / "tree.json"
    if not tree_path.exists():
        raise FileNotFoundError(f"树索引文件不存在: {tree_path}")
    tree_index = TreeIndex.load_json(str(tree_path))
    logger.info("已加载 TreeIndex: video_id={}, L1 节点数={}", video_id, len(tree_index.roots))

    # Phase 2: 构建 TreeEnvironment
    frames_dir = store_dir / "videos" / video_id / "frames"
    env = TreeEnvironment(index=tree_index, frames_dir=frames_dir)

    # Phase 3: 构建 SkillRegistry
    skills = None
    always_skills_text = ""
    task_skill_map: dict[str, str] = {}
    catalog_text = ""
    if skills_dir is not None:
        always_skills_text, task_skill_map, catalog_text, skills = discover_skills(skills_dir)
        logger.info(
            "已发现 skills: always={} 字符, task_map={} 项",
            len(always_skills_text),
            len(task_skill_map),
        )

    # Phase 4: 构建 SearchToolDispatcher
    dispatcher = SearchToolDispatcher(
        env,
        tool_llm=llm,
        vlm=vlm,
        ocr=ocr,
        prompts_dir=prompts_dir,
        skills=skills,
        embed_fn=embed_provider.embed,
        verify_vision=verify_vision,
        anchor=anchor,
        assemble_mode=assemble_mode,
    )

    # Phase 5: 构建 PromptManager + _prompt_builder 闭包
    pm = PromptManager(prompts_dir)
    l1_ids = [root.id for root in tree_index.roots]

    def _prompt_builder(qa: GeneratedQuestion) -> tuple[str, str]:
        """为单条题目生成 (system_prompt, user_prompt)。

        参数:
            qa: 生成的题目实例。

        返回:
            (system_prompt, user_prompt) 二元组。
        """
        system = pm.build_inference_prompt(
            skill_mode,
            qa.task_type,
            always_skills_text,
            task_skill_map,
            catalog_text,
        )
        user = pm.format_user_prompt(
            qa.question,
            list(qa.options),
            l1_ids,
            qa.task_type,
        )
        return system, user

    logger.info("InferenceDeps 组装完成: video_id={}, skill_mode={}", video_id, skill_mode)
    return InferenceDeps(
        llm=llm,
        tool_dispatch_fn=dispatcher.dispatch,
        prompt_builder=_prompt_builder,
    )