scripts/ask.py

"""Ask a question against the indexed RL corpus.

Usage:
    uv run python scripts/ask.py "What is the main idea of PPO?"
    uv run python scripts/ask.py --k 8 "How does Rainbow combine Double DQN and PER?"
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel

from researchpath.embeddings import Embedder
from researchpath.index import load_index, search
from researchpath.rag import answer
from researchpath.retrieval import HybridRetriever, Reranker

console = Console()
ROOT = Path(__file__).resolve().parents[1]
INDEX_PATH = ROOT / "data" / "index.faiss"


def main() -> int:
    parser = argparse.ArgumentParser(description="Ask the RL corpus a question.")
    parser.add_argument("question", help="Question to ask.")
    parser.add_argument("--k", type=int, default=5, help="Top-k chunks to retrieve (default: 5).")
    parser.add_argument("--hybrid", action="store_true", help="Use BM25+FAISS hybrid retrieval.")
    parser.add_argument("--rerank", action="store_true", help="Apply cross-encoder reranking on retrieved candidates.")
    parser.add_argument(
        "--show-sources",
        action="store_true",
        help="Print the retrieved chunks before the answer.",
    )
    args = parser.parse_args()

    if not INDEX_PATH.exists():
        console.print(f"[red]No index at {INDEX_PATH}. Run scripts/build_index.py first.[/red]")
        return 1

    if args.rerank:
        mode = "hybrid+rerank (BM25+FAISS+CrossEncoder)" if args.hybrid else "dense+rerank (FAISS+CrossEncoder)"
    else:
        mode = "hybrid (BM25+FAISS)" if args.hybrid else "dense (FAISS)"
    console.print(f"[bold cyan]Q:[/bold cyan] {args.question}  [dim][retrieval: {mode}][/dim]\n")

    index, chunks = load_index(INDEX_PATH)
    embedder = Embedder()
    if args.rerank:
        base = HybridRetriever(index, chunks, embedder) if args.hybrid else None
        if base is None:
            # Wrap dense FAISS in a hybrid-compatible interface for the reranker
            base = HybridRetriever(index, chunks, embedder)
        reranker = Reranker(base)
        hits = reranker.search(args.question, k=args.k)
    elif args.hybrid:
        retriever = HybridRetriever(index, chunks, embedder)
        hits = retriever.search(args.question, k=args.k)
    else:
        hits = search(index, chunks, embedder, args.question, k=args.k)

    if args.show_sources:
        for i, h in enumerate(hits, 1):
            console.print(
                Panel(
                    h.text,
                    title=f"#{i}  [{h.arxiv_id}, p{h.page}]  score={h.score:.3f}",
                    border_style="dim",
                )
            )
        console.print()

    result = answer(args.question, hits)

    console.print(Panel(Markdown(result.answer), title="Answer", border_style="green"))
    console.print(
        f"\n[dim]Retrieved {len(hits)} chunks  |  "
        f"model: {result.llm.model}  |  "
        f"tokens: {result.llm.input_tokens} in / {result.llm.output_tokens} out[/dim]"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())