main.py

"""統一進入點 CLI：

    python main.py ingest [--rebuild]    # 索引文件
    python main.py ask "你的問題"        # 單次問答
    python main.py chat                  # 互動式 CLI
    python main.py ui                    # 啟動 Streamlit UI
    python main.py eval                  # 跑 RAGAS 評估
"""

from __future__ import annotations

import argparse
import subprocess
import sys


def cmd_ingest(args: argparse.Namespace) -> None:
    from src.ingest import ingest

    ingest(rebuild=args.rebuild)


def cmd_ask(args: argparse.Namespace) -> None:
    from src.rag import build_chain

    chain, retriever = build_chain()
    print(f"\n問題：{args.question}\n")
    print("回答：", end="", flush=True)
    for token in chain.stream(args.question):
        print(token, end="", flush=True)
    print("\n\n--- 引用來源 ---")
    for i, doc in enumerate(retriever.invoke(args.question), 1):
        src = doc.metadata.get("source", "?")
        page = doc.metadata.get("page", "-")
        print(f"  [{i}] {src} (頁 {page})")


def cmd_chat(_: argparse.Namespace) -> None:
    from src.rag import cli

    cli()


def cmd_ui(_: argparse.Namespace) -> None:
    print("啟動 Streamlit UI...")
    subprocess.run(
        [sys.executable, "-m", "streamlit", "run", "src/app.py"],
        check=False,
    )


def cmd_eval(_: argparse.Namespace) -> None:
    from src.evaluate import evaluate

    evaluate()


def cmd_smoke(_: argparse.Namespace) -> None:
    from tests.test_smoke import main_cli

    main_cli()


def main() -> None:
    parser = argparse.ArgumentParser(prog="rag", description="RAG AI 專案 CLI")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_ingest = sub.add_parser("ingest", help="索引 data/ 下的文件")
    p_ingest.add_argument("--rebuild", action="store_true", help="清空後重建")
    p_ingest.set_defaults(func=cmd_ingest)

    p_ask = sub.add_parser("ask", help="問一個問題並印出答案")
    p_ask.add_argument("question", help="你的問題")
    p_ask.set_defaults(func=cmd_ask)

    p_chat = sub.add_parser("chat", help="互動式 CLI")
    p_chat.set_defaults(func=cmd_chat)

    p_ui = sub.add_parser("ui", help="啟動 Streamlit Web UI")
    p_ui.set_defaults(func=cmd_ui)

    p_eval = sub.add_parser("eval", help="跑 RAGAS 評估")
    p_eval.set_defaults(func=cmd_eval)

    p_smoke = sub.add_parser("smoke", help="跑 smoke test")
    p_smoke.set_defaults(func=cmd_smoke)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()