Video-Tree-TRM5/reference/main.py

"""
推理/演示入口
=============
从原始文档（文本或视频）构建 TreeIndex，执行问答。
与 ``train.py``（训练入口）配对，是 Video-Tree-TRM 的端到端演示入口。

子命令
------
- ``index``:  仅构建并缓存 TreeIndex，不执行问答。
- ``query``:  构建索引（或复用缓存）后执行问答。

使用示例::

    # 仅构建文本索引
    python main.py index --source data/doc.txt --modality text

    # 单次问答（视频）
    python main.py query --source data/video.mp4 --modality video \\
        --question "视频的主要内容是什么？"

    # 交互式多轮问答（文本）
    python main.py query --source data/doc.txt --modality text --interactive

    # 自定义配置
    python main.py query --source data/doc.txt --modality text \\
        --config config/default.yaml --env .env \\
        --question "文档结论是什么？"
"""

from __future__ import annotations

import argparse
import sys

from utils.logger_system import log_msg
from video_tree_trm.config import Config
from video_tree_trm.pipeline import Pipeline


# ---------------------------------------------------------------------------
# CLI 解析
# ---------------------------------------------------------------------------


def _parse_args() -> argparse.Namespace:
    """解析 CLI 参数，返回 Namespace。

    返回:
        包含子命令及所有选项的 Namespace 对象。
    """
    parser = argparse.ArgumentParser(
        prog="main.py",
        description="Video-Tree-TRM 推理/演示入口",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # 公共参数
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument(
        "--config",
        default="config/default.yaml",
        metavar="YAML",
        help="配置文件路径（默认: config/default.yaml）",
    )
    common.add_argument(
        "--env",
        default=".env",
        metavar="ENV",
        help="环境变量文件路径（默认: .env）",
    )
    common.add_argument(
        "--source",
        required=True,
        metavar="PATH",
        help="原始文件路径（文本文件或视频文件）",
    )
    common.add_argument(
        "--modality",
        required=True,
        choices=["text", "video"],
        help="文件模态：text 或 video",
    )

    subparsers = parser.add_subparsers(dest="command", title="子命令")
    subparsers.required = True

    # ── index 子命令 ──
    subparsers.add_parser(
        "index",
        parents=[common],
        help="构建并缓存 TreeIndex（不执行问答）",
        description="从原始文件构建三层树索引并保存到缓存目录。",
    )

    # ── query 子命令 ──
    query_parser = subparsers.add_parser(
        "query",
        parents=[common],
        help="构建索引并执行问答",
        description="构建索引后执行单次问答或交互式多轮问答。",
    )
    mode_group = query_parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument(
        "--question",
        metavar="TEXT",
        help="单次问答：问题字符串",
    )
    mode_group.add_argument(
        "--interactive",
        action="store_true",
        help="交互式多轮问答模式（输入 quit 退出）",
    )

    return parser.parse_args()


# ---------------------------------------------------------------------------
# 配置加载
# ---------------------------------------------------------------------------


def _load_config(args: argparse.Namespace) -> Config:
    """从 YAML + .env 加载配置。

    参数:
        args: CLI 解析结果，需含 config（YAML 路径）和 env（.env 路径）字段。

    返回:
        完整的 Config 实例。

    异常:
        FileNotFoundError: 配置文件不存在。
        TypeError: YAML 缺少必需字段。
    """
    log_msg("INFO", "加载配置", yaml_path=args.config, env_path=args.env)
    return Config.load(
        yaml_path=args.config,
        env_path=args.env,
    )


# ---------------------------------------------------------------------------
# 子命令实现
# ---------------------------------------------------------------------------


def cmd_index(args: argparse.Namespace) -> None:
    """index 子命令：构建并缓存 TreeIndex，打印缓存路径。

    参数:
        args: 含 source、modality、config、env 字段的 Namespace。
    """
    config = _load_config(args)
    pipeline = Pipeline(config)

    log_msg("INFO", "开始构建索引", source=args.source, modality=args.modality)
    tree = pipeline.build_index(args.source, args.modality)

    print(f"索引构建完成。模态={args.modality}，来源={args.source}")
    print(f"节点数量: L1={len(tree.roots)}")
    log_msg("INFO", "索引构建完成", l1_count=len(tree.roots))


def cmd_query(args: argparse.Namespace) -> None:
    """query 子命令：构建索引 → 单次或交互式问答。

    参数:
        args: 含 source、modality、question/interactive、config、env 字段的 Namespace。

    实现细节:
        - 先调用 build_index（缓存命中时直接加载，无额外开销）。
        - --question 模式：单次问答后退出。
        - --interactive 模式：循环读取用户输入，输入 'quit' 退出。
    """
    config = _load_config(args)
    pipeline = Pipeline(config)

    log_msg("INFO", "构建/加载索引", source=args.source, modality=args.modality)
    tree = pipeline.build_index(args.source, args.modality)

    if args.question:
        # Phase 1: 单次问答
        answer = pipeline.query(args.question, tree)
        print(f"\n问题: {args.question}")
        print(f"答案: {answer}\n")
        log_msg("INFO", "问答完成", question=args.question[:50])
    else:
        # Phase 2: 交互式多轮问答
        print("\n已进入交互式问答模式。输入 'quit' 退出。\n")
        while True:
            try:
                question = input("问题 (输入 'quit' 退出): ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\n再见。")
                break

            if question.lower() in ("quit", "exit", "q"):
                print("再见。")
                break

            if not question:
                continue

            answer = pipeline.query(question, tree)
            print(f"答案: {answer}\n")
            log_msg("INFO", "交互式问答", question=question[:50])


# ---------------------------------------------------------------------------
# 主入口
# ---------------------------------------------------------------------------


def main() -> None:
    """主入口：解析参数 → 分发子命令。

    退出码:
        0: 正常退出。
        1: 参数错误或运行时异常。
    """
    args = _parse_args()

    try:
        if args.command == "index":
            cmd_index(args)
        elif args.command == "query":
            cmd_query(args)
    except (FileNotFoundError, TypeError, ValueError) as exc:
        print(f"错误: {exc}", file=sys.stderr)
        log_msg("ERROR", f"main.py 异常退出: {exc}")
        sys.exit(1)


if __name__ == "__main__":
    main()