src/agentic_search/cli.py

"""CLI for corpus preparation, retrieval-only evaluation, reports, and trace replay."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Annotated

import httpx
import typer

from . import __version__
from .config import AppConfig, load_config, load_yaml_mapping
from .data.combined import sha256_file, write_benchmark_manifests, write_combined_corpus
from .data.corpus import write_documents
from .data.hotpotqa import HotpotQAAdapter
from .data.musique import MuSiQueAdapter
from .data.schemas import BenchmarkExample
from .data.twowiki import TwoWikiAdapter
from .evaluation.aggregate import aggregate_records, latest_records
from .evaluation.report import write_report
from .evaluation.runner import BenchmarkRunner
from .exceptions import ConfigurationError
from .logging import configure_logging, get_logger, redact
from .retrieval import DenseRetriever, DocumentStore, SparseRetriever
from .tracing.replay import frozen_replay

app = typer.Typer(
    name="agentic-search",
    help="Retrieval-only benchmark: FIXED vs Agentic vs iterative Search-as-Code.",
    invoke_without_command=True,
    no_args_is_help=False,
)
data_app = typer.Typer(help="Collect all benchmarks and build a shared offline corpus.", no_args_is_help=True)
index_app = typer.Typer(help="Build auditable sparse or dense index metadata.", no_args_is_help=True)
run_app = typer.Typer(help="Run one retrieval query and show retrieved context.", no_args_is_help=True)
benchmark_app = typer.Typer(help="Run a resumable retrieval-only benchmark.", no_args_is_help=True)
report_app = typer.Typer(help="Aggregate retrieval-only result records.", no_args_is_help=True)
replay_app = typer.Typer(help="Read a frozen trajectory without model or index calls.", no_args_is_help=True)
ui_app = typer.Typer(help="Terminal trace viewer for inspectable benchmark demos.", no_args_is_help=True)
config_app = typer.Typer(help="Inspect and validate project configuration.", no_args_is_help=True)

DEFAULT_HOTPOT_PARQUET_URL = (
    "https://huggingface.co/datasets/hotpotqa/hotpot_qa/resolve/"
    "refs%2Fconvert%2Fparquet/distractor/validation/0000.parquet"
)
DEFAULT_TWOWIKI_PARQUET_URL = (
    "https://huggingface.co/datasets/framolfese/2WikiMultihopQA/resolve/main/"
    "data/validation-00000-of-00001.parquet"
)
DEFAULT_MUSIQUE_JSON_URL = "https://huggingface.co/datasets/fladhak/musique/resolve/main/dev.json"


def _public_config(config: AppConfig) -> dict[str, object]:
    output = config.model_dump(mode="json")
    output["model"]["base_url_configured"] = bool(config.resolved_model_base_urls())
    output["model"]["replica_count"] = len(config.resolved_model_base_urls())
    output["model"]["api_key_configured"] = config.resolved_model_api_key() is not None
    return redact(output)


def _examples(path: Path, max_examples: int) -> list[BenchmarkExample]:
    if path.suffix == ".jsonl":
        examples = [
            BenchmarkExample.model_validate_json(line)
            for line in path.read_text(encoding="utf-8").splitlines()
            if line.strip()
        ]
    else:
        examples = list(HotpotQAAdapter(path).load_examples())
    return examples[:max_examples] if max_examples > 0 else examples


def _download_stream(url: str, output: Path) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(output.suffix + ".tmp")
    with httpx.stream("GET", url, timeout=300, follow_redirects=True) as response:
        response.raise_for_status()
        with temporary.open("wb") as destination:
            for chunk in response.iter_bytes():
                destination.write(chunk)
    temporary.replace(output)


@data_app.command("download-hotpot")
def download_hotpot(
    output: Annotated[Path, typer.Option("--output", help="Destination HotpotQA JSON file.")] = Path("data/raw/hotpotqa/dev_distractor.json"),
    parquet_url: Annotated[str, typer.Option("--parquet-url", help="Hugging Face validation parquet URL.")] = DEFAULT_HOTPOT_PARQUET_URL,
    limit: Annotated[int, typer.Option("--limit", help="0 fetches the complete validation split.")] = 0,
) -> None:
    """Download the official HF HotpotQA distractor validation split as original-style JSON."""
    output = output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    parquet_path = output.with_suffix(".parquet.tmp")
    with httpx.stream("GET", parquet_url, timeout=180, follow_redirects=True) as response:
        response.raise_for_status()
        with parquet_path.open("wb") as destination:
            for chunk in response.iter_bytes():
                destination.write(chunk)
    try:
        import pyarrow.parquet as parquet
    except ImportError as error:
        raise RuntimeError("pyarrow is required to normalize the official HotpotQA parquet") from error
    rows = parquet.read_table(parquet_path).to_pylist()
    parquet_path.unlink(missing_ok=True)
    if limit:
        rows = rows[:limit]
    records = [
        {
            "_id": row["id"],
            "question": row["question"],
            "answer": row["answer"],
            "type": row.get("type"),
            "level": row.get("level"),
            "supporting_facts": [
                (title, sent_id)
                for title, sent_id in zip(
                    row["supporting_facts"]["title"], row["supporting_facts"]["sent_id"], strict=True
                )
            ],
            "context": [
                (title, sentences)
                for title, sentences in zip(
                    row["context"]["title"], row["context"]["sentences"], strict=True
                )
            ],
        }
        for row in rows
    ]
    temporary = output.with_suffix(output.suffix + ".tmp")
    temporary.write_text(json.dumps(records, ensure_ascii=False), encoding="utf-8")
    temporary.replace(output)
    typer.echo(json.dumps({"output": str(output), "examples": len(records), "parquet_url": parquet_url}))


@data_app.command("build-hotpot-corpus")
def build_hotpot_corpus(
    source: Annotated[Path, typer.Option("--source", help="Raw HotpotQA JSON file.")],
    output: Annotated[Path, typer.Option("--output", help="Corpus JSONL output.")] = Path("data/processed/hotpotqa_context_union.jsonl"),
) -> None:
    """Build a stable corpus from the union of offline HotpotQA contexts."""
    manifest = write_documents(HotpotQAAdapter(source).build_corpus(), output.resolve())
    manifest_path = output.resolve().with_suffix(".manifest.json")
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    typer.echo(json.dumps({**manifest, "manifest": str(manifest_path)}, ensure_ascii=False))


@data_app.command("download-2wiki")
def download_twowiki(
    output: Annotated[Path, typer.Option("--output", help="2Wiki validation parquet destination.")] = Path(
        "data/raw/2wikimultihopqa/validation.parquet"
    ),
    parquet_url: Annotated[str, typer.Option("--parquet-url")] = DEFAULT_TWOWIKI_PARQUET_URL,
) -> None:
    """Download the HotpotQA-schema-compatible 2WikiMultiHopQA validation split."""
    output = output.resolve()
    _download_stream(parquet_url, output)
    typer.echo(json.dumps({"output": str(output), "sha256": sha256_file(output), "url": parquet_url}))


@data_app.command("download-musique")
def download_musique(
    output: Annotated[Path, typer.Option("--output", help="MuSiQue dev JSON destination.")] = Path(
        "data/raw/musique/dev.json"
    ),
    json_url: Annotated[str, typer.Option("--json-url")] = DEFAULT_MUSIQUE_JSON_URL,
) -> None:
    """Download the MuSiQue development set with paragraph-level supporting labels."""
    output = output.resolve()
    _download_stream(json_url, output)
    typer.echo(json.dumps({"output": str(output), "sha256": sha256_file(output), "url": json_url}))


@data_app.command("download-all")
def download_all() -> None:
    """Collect every retrieval benchmark source used by the combined corpus."""
    hotpot = Path("data/raw/hotpotqa/dev_distractor.json")
    if not hotpot.is_file():
        download_hotpot(output=hotpot)
    twowiki = Path("data/raw/2wikimultihopqa/validation.parquet")
    if not twowiki.is_file():
        download_twowiki(output=twowiki)
    musique = Path("data/raw/musique/dev.json")
    if not musique.is_file():
        download_musique(output=musique)
    typer.echo(
        json.dumps(
            {
                "hotpotqa": str(hotpot.resolve()),
                "2wikimultihopqa": str(twowiki.resolve()),
                "musique": str(musique.resolve()),
            }
        )
    )


@data_app.command("build-combined-corpus")
def build_combined_corpus(
    hotpot: Annotated[Path, typer.Option("--hotpot")]=Path("data/raw/hotpotqa/dev_distractor.json"),
    twowiki: Annotated[Path, typer.Option("--twowiki")]=Path("data/raw/2wikimultihopqa/validation.parquet"),
    musique: Annotated[Path, typer.Option("--musique")]=Path("data/raw/musique/dev.json"),
    output: Annotated[Path, typer.Option("--output")]=Path("data/processed/multihop_context_union.jsonl"),
) -> None:
    """Build one global deduplicated corpus from all three benchmark context collections."""
    inputs = {"hotpotqa": hotpot.resolve(), "2wikimultihopqa": twowiki.resolve(), "musique": musique.resolve()}
    missing = [str(path) for path in inputs.values() if not path.is_file()]
    if missing:
        raise typer.BadParameter(f"missing benchmark sources: {', '.join(missing)}")
    manifest = write_combined_corpus(
        [HotpotQAAdapter(inputs["hotpotqa"]), TwoWikiAdapter(inputs["2wikimultihopqa"]), MuSiQueAdapter(inputs["musique"])],
        output.resolve(),
    )
    manifest["source_checksums"] = {dataset: sha256_file(path) for dataset, path in inputs.items()}
    manifest_path = output.resolve().with_suffix(".manifest.json")
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    typer.echo(json.dumps({**manifest, "manifest": str(manifest_path)}, ensure_ascii=False))


@data_app.command("build-benchmark-manifests")
def build_benchmark_manifests(
    hotpot: Annotated[Path, typer.Option("--hotpot")] = Path("data/raw/hotpotqa/dev_distractor.json"),
    twowiki: Annotated[Path, typer.Option("--twowiki")] = Path("data/raw/2wikimultihopqa/validation.parquet"),
    musique: Annotated[Path, typer.Option("--musique")] = Path("data/raw/musique/dev.json"),
    output_dir: Annotated[Path, typer.Option("--output-dir")] = Path("data/processed/benchmarks"),
) -> None:
    """Export source-partitioned questions and gold IDs for later retrieval-only evaluation."""
    inputs = {"hotpotqa": hotpot.resolve(), "2wikimultihopqa": twowiki.resolve(), "musique": musique.resolve()}
    missing = [str(path) for path in inputs.values() if not path.is_file()]
    if missing:
        raise typer.BadParameter(f"missing benchmark sources: {', '.join(missing)}")
    manifest = write_benchmark_manifests(
        [
            HotpotQAAdapter(inputs["hotpotqa"]),
            TwoWikiAdapter(inputs["2wikimultihopqa"]),
            MuSiQueAdapter(inputs["musique"]),
        ],
        output_dir.resolve(),
    )
    manifest["source_checksums"] = {dataset: sha256_file(path) for dataset, path in inputs.items()}
    manifest_path = output_dir.resolve() / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    typer.echo(json.dumps(manifest, ensure_ascii=False))


@index_app.command("sparse")
def build_sparse_index(
    corpus: Annotated[Path, typer.Option("--corpus", help="Corpus JSONL path.")],
    output: Annotated[Path, typer.Option("--output", help="Sparse index metadata JSON.")] = Path("data/indices/sparse.metadata.json"),
) -> None:
    """Build and audit the portable BM25 index statistics."""
    retriever = SparseRetriever(DocumentStore.from_jsonl(corpus).all())
    retriever.save_metadata(output.resolve())
    typer.echo(str(output.resolve()))


@index_app.command("dense")
def build_dense_index(
    corpus: Annotated[Path, typer.Option("--corpus", help="Corpus JSONL path.")],
    index_dir: Annotated[Path | None, typer.Option("--index-dir", help="Persistent BGE/FAISS index directory.")] = None,
    config_path: Annotated[Path | None, typer.Option("--config", "-c")] = None,
) -> None:
    """Encode every corpus passage with BGE and persist FAISS, vectors, and doc-id mapping."""
    config = load_config(config_path)
    documents = DocumentStore.from_jsonl(corpus).all()
    if config.retrieval.dense.backend != "sentence_transformers":
        raise typer.BadParameter("index dense requires dense.backend=sentence_transformers")
    output = (index_dir or config.retrieval.dense.index_dir).resolve()
    manifest = DenseRetriever.build_persisted_index(
        documents,
        model_name=config.retrieval.dense.model_name,
        device=config.retrieval.dense.device,
        batch_size=config.retrieval.dense.batch_size,
        output_dir=output,
    )
    typer.echo(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True))


@run_app.command("query")
def run_query(
    question: Annotated[str, typer.Argument(help="Question to retrieve evidence for.")],
    corpus: Annotated[Path, typer.Option("--corpus")],
    system: Annotated[str, typer.Option("--system", help="fixed, agentic, sac, or a controller baseline.")] = "fixed",
    config_path: Annotated[Path | None, typer.Option("--config", "-c")] = None,
) -> None:
    """Run one system and print only its retrieved contexts plus cost/latency."""
    runner = BenchmarkRunner(load_config(config_path), corpus_path=corpus.resolve())
    try:
        record = runner.run_one(BenchmarkExample(example_id="adhoc", question=question, answers=[]), system)
        typer.echo(json.dumps(record, ensure_ascii=False, indent=2, sort_keys=True))
    finally:
        runner.close()


@benchmark_app.command("run")
def benchmark_run(
    corpus: Annotated[Path, typer.Option("--corpus", help="Offline corpus JSONL.")],
    dataset: Annotated[
        Path | None,
        typer.Option("--dataset", help="Legacy raw HotpotQA JSON input; prefer --examples."),
    ] = None,
    examples: Annotated[
        Path | None,
        typer.Option("--examples", help="Exported source-partitioned BenchmarkExample JSONL."),
    ] = None,
    experiment: Annotated[Path, typer.Option("--experiment", "-e", help="Experiment YAML.")] = Path("configs/experiments/smoke.yaml"),
    config_path: Annotated[Path | None, typer.Option("--config", "-c")] = None,
    run_id: Annotated[str | None, typer.Option("--run-id", help="Resume this run ID when supplied.")] = None,
    systems_override: Annotated[
        str | None,
        typer.Option("--systems", help="Comma-separated systems; overrides the experiment matrix."),
    ] = None,
    max_examples: Annotated[int | None, typer.Option("--max-examples")] = None,
    concurrency: Annotated[int | None, typer.Option("--concurrency")] = None,
    warmup_examples: Annotated[
        int,
        typer.Option("--warmup-examples", help="Untimed examples per system used only to warm models."),
    ] = 0,
    use_cache: Annotated[
        bool | None,
        typer.Option("--cache/--no-cache", help="Override configured cache behavior for this run."),
    ] = None,
    force: Annotated[bool, typer.Option("--force", help="Re-run completed example/system pairs.")] = False,
) -> None:
    """Evaluate context coverage, retrieval evidence, latency, and token cost only."""
    if (dataset is None) == (examples is None):
        raise typer.BadParameter("provide exactly one of --dataset or --examples")
    config = load_config(config_path)
    if use_cache is not None:
        config = config.model_copy(
            update={"cache": config.cache.model_copy(update={"enabled": use_cache})}
        )
    experiment_values = load_yaml_mapping(experiment)
    settings = experiment_values.get("experiment", {})
    systems = (
        [item.strip() for item in systems_override.split(",") if item.strip()]
        if systems_override is not None
        else [str(item) for item in settings.get("systems", ["fixed", "agentic", "sac"])]
    )
    if not systems:
        raise typer.BadParameter("--systems must name at least one system")
    limit = max_examples if max_examples is not None else int(settings.get("max_examples", 20))
    workers = concurrency if concurrency is not None else int(settings.get("concurrency", config.model.request_concurrency))
    runner = BenchmarkRunner(config, corpus_path=corpus.resolve(), run_id=run_id)
    try:
        source = (examples or dataset).resolve()
        examples_for_run = _examples(source, limit)
        runner.warmup(examples_for_run, systems, count=warmup_examples)
        records_path, records = runner.run(examples_for_run, systems, concurrency=workers, force=force)
        report_rows = aggregate_records(
            records,
            bootstrap_samples=config.evaluation.bootstrap_samples,
            seed=config.project.seed,
        )
        report_dir = config.project.artifacts_dir / "reports" / runner.run_id
        json_path, csv_path = write_report(report_rows, report_dir)
        typer.echo(json.dumps({"run_id": runner.run_id, "records": str(records_path), "report_json": str(json_path), "report_csv": str(csv_path)}, indent=2))
    finally:
        runner.close()


@report_app.command("build")
def build_report(
    records: Annotated[Path, typer.Option("--records", help="results.jsonl path.")],
    output_dir: Annotated[Path, typer.Option("--output-dir")],
    config_path: Annotated[Path | None, typer.Option("--config", "-c")] = None,
) -> None:
    """Rebuild a report from recorded results without rerunning retrieval or model calls."""
    config = load_config(config_path)
    rows = latest_records(
        [json.loads(line) for line in records.read_text(encoding="utf-8").splitlines() if line]
    )
    report = aggregate_records(rows, bootstrap_samples=config.evaluation.bootstrap_samples, seed=config.project.seed)
    json_path, csv_path = write_report(report, output_dir)
    typer.echo(json.dumps({"json": str(json_path), "csv": str(csv_path)}))


@replay_app.command("show")
def replay_show(trace: Annotated[Path, typer.Option("--trace", help="Leaf trace directory.")]) -> None:
    """Display persisted events exactly as recorded, without live model calls."""
    typer.echo(json.dumps(frozen_replay(trace), ensure_ascii=False, indent=2))


@ui_app.command("trace")
def ui_trace(trace: Annotated[Path, typer.Option("--trace", help="Leaf trace directory.")]) -> None:
    """Render a concise terminal trace viewer for demos and failure inspection."""
    for event in frozen_replay(trace):
        typer.echo(
            f"{event['step_id']} {event['component']}:{event['event_type']} "
            f"latency_ms={event['latency_ms']:.2f}"
        )


@config_app.command("show")
def show_config(
    config_path: Annotated[Path | None, typer.Option("--config", "-c", help="YAML config path.")] = None,
) -> None:
    """Print resolved configuration without exposing secret values."""
    try:
        typer.echo(json.dumps(_public_config(load_config(config_path)), ensure_ascii=False, indent=2, sort_keys=True))
    except ConfigurationError as error:
        typer.echo(str(error), err=True)
        raise typer.Exit(code=2) from error


@app.callback()
def main_callback(
    context: typer.Context,
    version: Annotated[bool, typer.Option("--version", help="Print package version.")] = False,
    log_level: Annotated[str | None, typer.Option("--log-level", help="Logging severity.")] = None,
) -> None:
    configure_logging(log_level)
    if version:
        typer.echo(__version__)
        raise typer.Exit()
    if context.invoked_subcommand is None:
        typer.echo(context.get_help())


app.add_typer(data_app, name="data")
app.add_typer(index_app, name="index")
app.add_typer(run_app, name="run")
app.add_typer(benchmark_app, name="benchmark")
app.add_typer(report_app, name="report")
app.add_typer(replay_app, name="replay")
app.add_typer(ui_app, name="ui")
app.add_typer(config_app, name="config")


def main() -> None:
    get_logger("cli").debug("starting CLI", extra={"fields": {"version": __version__}})
    app()