lean-bench/leanbench/cli.py

"""Command line: ``lean-bench <command>``.

    lean-bench doctor                     is everything wired up?
    lean-bench ingest book.pdf --pages 30-45
    lean-bench extract book.pdf --pages 90 --compare   which extractor should I use?
    lean-bench detect book.pdf --pages 30-31 --out detections.jsonl
    lean-bench search --type Int --role input --kind problem
    lean-bench export datasets/items.jsonl
    lean-bench drop --source <hash> --pages 30-45     redo part of a book
"""
from __future__ import annotations

import argparse
import json
import sys

from .config import Settings
from .db import Database
from .detect import detect_document, write_jsonl
from .document import Document
from .llm import LLMError, build_client
from .pipeline import build_checker, build_ocr, ingest


def _pages(value: str | None) -> tuple[int | None, int | None]:
    if not value:
        return None, None
    text = value.strip()
    if "-" in text:
        first, _, last = text.partition("-")
        return (int(first) if first.strip() else None,
                int(last) if last.strip() else None)
    page = int(text)
    return page, page


def _say(message: str) -> None:
    print(message, flush=True)


def _add_common(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--db", help="SQLite path (default: $LEANBENCH_DATABASE)")
    parser.add_argument("--env", default=".env", help="dotenv file to load (default: .env)")


def _add_extractor(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--extractor", choices=["text", "auto", "ocr"],
        help="where page text comes from: the PDF's own text, OCR, or OCR only for "
             "pages whose embedded text looks broken (auto)")
    parser.add_argument("--ocr-model", help="vision model for OCR (default qwen-vl-max)")
    parser.add_argument("--ocr-dpi", type=int, help="render resolution for OCR (default 200)")
    parser.add_argument("--ocr-backend", choices=["vision", "tesseract"])
    parser.add_argument("--ocr-cache", help="directory for cached transcriptions")


def _apply_extractor(settings, args) -> None:
    if getattr(args, "extractor", None):
        settings.ocr.mode = args.extractor
    if getattr(args, "ocr_model", None):
        settings.ocr.model = args.ocr_model
    if getattr(args, "ocr_dpi", None):
        settings.ocr.dpi = args.ocr_dpi
    if getattr(args, "ocr_backend", None):
        settings.ocr.backend = args.ocr_backend
    if getattr(args, "ocr_cache", None):
        settings.ocr.cache_dir = args.ocr_cache


# -- commands --------------------------------------------------------------


def cmd_doctor(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    print(f"database      {settings.database}")
    print(f"llm model     {settings.llm.model}")
    print(f"llm endpoint  {settings.llm.base_url}")
    print(f"llm key env   {settings.llm.api_key_env}")

    try:
        import pymupdf  # noqa: F401
        print("pdf backend   pymupdf")
    except ImportError:
        try:
            import pypdf  # noqa: F401
            print("pdf backend   pypdf (install pymupdf for better text)")
        except ImportError:
            print("pdf backend   MISSING - pip install pymupdf")

    checker = build_checker(settings)
    available = getattr(checker, "available", lambda: False)()
    print(f"lean          {'ready: ' + ' '.join(settings.lean.command) if available else 'not available (verification will be skipped)'}")
    if settings.lean.project:
        print(f"lean project  {settings.lean.project}")

    if args.skip_llm:
        return 0
    try:
        client = build_client(settings.llm)
        reply = client.complete(
            [{"role": "user", "content": "Reply with the single word: ready"}],
            max_tokens=16,
        )
        print(f"llm           ready ({reply.strip()[:40]})")
    except LLMError as error:
        print(f"llm           FAILED: {error}")
        return 1
    return 0


def cmd_ingest(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    _apply_extractor(settings, args)
    if args.db:
        settings.database = args.db
    if args.window:
        settings.max_window_pages = args.window
    first, last = _pages(args.pages)

    database = Database(settings.database)
    try:
        result = ingest(
            args.source,
            settings=settings,
            database=database,
            checker=build_checker(settings, enabled=not args.no_verify),
            first_page=first,
            last_page=last,
            title=args.title or "",
            replace=not args.keep_existing,
            store_results=not args.no_store,
            progress=None if args.quiet else _say,
        )
        summary = result.summary()
        print(json.dumps(summary, indent=2))
        if args.detections:
            write_jsonl(args.detections, result.report.detections)
            print(f"detections -> {args.detections}")
        if args.export:
            count = database.export_jsonl(args.export, source_hash=result.document.sha256)
            print(f"exported {count} item(s) -> {args.export}")
    finally:
        database.close()
    return 0


def cmd_detect(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    _apply_extractor(settings, args)
    first, last = _pages(args.pages)
    client = build_client(settings.llm)
    document = Document.load(
        args.source,
        mode=settings.ocr.mode,
        ocr=build_ocr(settings, client),
        threshold=settings.ocr.threshold,
        first_page=first, last_page=last,
        progress=None if args.quiet else _say,
    )
    report = detect_document(
        client,
        document,
        first_page=first,
        last_page=last,
        max_window_pages=args.window or settings.max_window_pages,
        progress=None if args.quiet else _say,
    )
    print(report.summary())
    if args.out:
        write_jsonl(args.out, report.detections)
        print(f"-> {args.out}")
    else:
        for detection in report.detections:
            print(json.dumps(detection.as_dict(), ensure_ascii=False))
    return 0


def cmd_extract(args: argparse.Namespace) -> int:
    """Show what the pipeline would actually read. The decision tool for --extractor."""
    from .ocr import text_quality

    settings = Settings.from_env(args.env)
    _apply_extractor(settings, args)
    first, last = _pages(args.pages)

    plain = Document.load(args.source)
    pages = plain.page_range(first, last)
    if not pages:
        print("no pages in that range")
        return 1

    if not args.compare and settings.ocr.mode == "text":
        for page in pages:
            print(f"----- page {page.number}  (quality {text_quality(page.text)}) -----")
            print(page.text)
        return 0

    reader = build_ocr(settings if settings.ocr.mode != "text" else _forced_ocr(settings),
                       build_client(settings.llm))
    for page in pages:
        embedded = page.text
        print(f"===== page {page.number} =====")
        print(f"--- embedded text  (quality {text_quality(embedded)}) ---")
        print(embedded if args.full else embedded[:800])
        transcribed = reader.page_text(plain.path, plain.sha256, page.number)
        print(f"--- {reader.backend.name} ---")
        if transcribed is None:
            print(f"FAILED: {reader.failures.get(page.number)}")
        else:
            print(transcribed if args.full else transcribed[:800])
    return 0


def _forced_ocr(settings):
    settings.ocr.mode = "ocr"
    return settings


def cmd_search(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    database = Database(args.db or settings.database)
    try:
        rows = database.search(
            types=args.type or (),
            roles=args.role or (),
            kind=args.kind,
            source_hash=args.source,
            page=args.page,
            verified=args.verified,
            proved=args.proved,
            text=args.text,
            match_all_types=args.all_types,
            limit=args.limit,
        )
        if args.json:
            print(json.dumps(rows, indent=2, ensure_ascii=False))
            return 0
        if not rows:
            print("no matches")
            return 0
        for row in rows:
            types = ", ".join(f"{t['name']}({t['role']})" for t in row["types"]) or "-"
            flag = "proved" if row["proved"] else "verified" if row["verified"] else "unverified"
            print(f"{row['id']}  {row['kind']:<9} p{row['start_page']:<4} {flag:<10} "
                  f"{row['label'] or row['lean_name']}")
            print(f"    types: {types}")
            if args.code:
                for line in (row["lean_code"] or "").splitlines():
                    print("    " + line)
    finally:
        database.close()
    return 0


def cmd_sources(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    database = Database(args.db or settings.database)
    try:
        rows = database.sources()
        if args.json:
            print(json.dumps(rows, indent=2, ensure_ascii=False))
            return 0
        for row in rows:
            print(f"{row['hash'][:16]}  {row['items']:>5} items  p{row['pages']:<5} "
                  f"{row['title']}  ({row['extractor']}, {row['ingested_at']})")
    finally:
        database.close()
    return 0


def cmd_stats(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    database = Database(args.db or settings.database)
    try:
        stats = database.stats()
        stats["top_types"] = database.type_counts(limit=args.top)
        print(json.dumps(stats, indent=2))
    finally:
        database.close()
    return 0


def cmd_export(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    database = Database(args.db or settings.database)
    try:
        count = database.export_jsonl(
            args.out, source_hash=args.source, include_source_text=args.with_text
        )
        print(f"exported {count} item(s) -> {args.out}")
    finally:
        database.close()
    return 0


def cmd_drop(args: argparse.Namespace) -> int:
    settings = Settings.from_env(args.env)
    database = Database(args.db or settings.database)
    try:
        first, last = _pages(args.pages)
        if first or last:
            removed = database.delete_pages(args.source, first or 1, last or 10**9)
        else:
            removed = database.delete_source(args.source)
        print(f"removed {removed} item(s)")
    finally:
        database.close()
    return 0


# -- wiring ----------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="lean-bench", description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    subparsers = parser.add_subparsers(dest="command", required=True)

    doctor = subparsers.add_parser("doctor", help="check LLM, Lean and PDF wiring")
    _add_common(doctor)
    doctor.add_argument("--skip-llm", action="store_true", help="do not call the model")
    doctor.set_defaults(func=cmd_doctor)

    ingest_parser = subparsers.add_parser("ingest", help="run the whole pipeline on a source")
    _add_common(ingest_parser)
    ingest_parser.add_argument("source", help="PDF (or .txt with form-feed page breaks)")
    ingest_parser.add_argument("--pages", help="page or range, e.g. 30-45")
    ingest_parser.add_argument("--title", help="human title for this source")
    ingest_parser.add_argument("--window", type=int, help="max pages per detection window")
    ingest_parser.add_argument("--no-verify", action="store_true", help="skip Lean")
    ingest_parser.add_argument("--no-store", action="store_true", help="do not write the database")
    ingest_parser.add_argument("--keep-existing", action="store_true",
                               help="do not remove this source's previous items first")
    ingest_parser.add_argument("--detections", help="also write raw detections to this JSONL")
    ingest_parser.add_argument("--export", help="export this source's items to JSONL afterwards")
    ingest_parser.add_argument("--quiet", action="store_true")
    _add_extractor(ingest_parser)
    ingest_parser.set_defaults(func=cmd_ingest)

    detect_parser = subparsers.add_parser("detect", help="detection only, no Lean, no database")
    _add_common(detect_parser)
    detect_parser.add_argument("source")
    detect_parser.add_argument("--pages")
    detect_parser.add_argument("--window", type=int)
    detect_parser.add_argument("--out", help="write detections to this JSONL")
    detect_parser.add_argument("--quiet", action="store_true")
    _add_extractor(detect_parser)
    detect_parser.set_defaults(func=cmd_detect)

    extract = subparsers.add_parser(
        "extract", help="show the page text the pipeline would read (and compare extractors)")
    _add_common(extract)
    _add_extractor(extract)
    extract.add_argument("source")
    extract.add_argument("--pages")
    extract.add_argument("--compare", action="store_true",
                         help="show embedded text and OCR side by side")
    extract.add_argument("--full", action="store_true", help="do not truncate")
    extract.set_defaults(func=cmd_extract)

    search = subparsers.add_parser("search", help="find items by the types they are about")
    _add_common(search)
    search.add_argument("--type", action="append", help="type name; repeatable")
    search.add_argument("--role", action="append", choices=["input", "goal", "own"])
    search.add_argument("--kind", choices=["statement", "problem"])
    search.add_argument("--source", help="source hash")
    search.add_argument("--page", type=int)
    search.add_argument("--text", help="full-text query")
    search.add_argument("--all-types", action="store_true", help="require every --type")
    search.add_argument("--verified", action="store_true", default=None)
    search.add_argument("--proved", action="store_true", default=None)
    search.add_argument("--code", action="store_true", help="print the Lean too")
    search.add_argument("--json", action="store_true")
    search.add_argument("--limit", type=int, default=25)
    search.set_defaults(func=cmd_search)

    sources = subparsers.add_parser("sources", help="list ingested books")
    _add_common(sources)
    sources.add_argument("--json", action="store_true")
    sources.set_defaults(func=cmd_sources)

    stats = subparsers.add_parser("stats", help="counts, and the most common types")
    _add_common(stats)
    stats.add_argument("--top", type=int, default=20)
    stats.set_defaults(func=cmd_stats)

    export = subparsers.add_parser("export", help="dump items as JSONL")
    _add_common(export)
    export.add_argument("out")
    export.add_argument("--source", help="only this source hash")
    export.add_argument("--with-text", action="store_true", help="include source passages")
    export.set_defaults(func=cmd_export)

    drop = subparsers.add_parser("drop", help="remove a source, or some of its pages")
    _add_common(drop)
    drop.add_argument("--source", required=True, help="source hash")
    drop.add_argument("--pages", help="page or range; omit for the whole source")
    drop.set_defaults(func=cmd_drop)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        return args.func(args)
    except FileNotFoundError as error:
        print(f"not found: {error}", file=sys.stderr)
        return 2
    except LLMError as error:
        print(f"llm error: {error}", file=sys.stderr)
        return 3
    except KeyboardInterrupt:
        print("interrupted", file=sys.stderr)
        return 130


if __name__ == "__main__":
    raise SystemExit(main())