arborist/aborist/cli.py

"""Aborist CLI: ingest / search / verify / stats."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from aborist import __version__
from aborist.ingest import ingest_source, verify_random_sample
from aborist.search import FTS5Backend
from aborist.sources import WikipediaCurDump
from aborist.store import DEFAULT_DB_PATH, connect, stats


def _cmd_ingest(args: argparse.Namespace) -> int:
    if args.source in ("wikipedia_cur", "wikipedia_old"):
        if not args.path:
            print(f"--path is required for {args.source}", file=sys.stderr)
            return 2
        from aborist.sources import WikipediaSqlDump

        table = "cur" if args.source == "wikipedia_cur" else "old"
        src = WikipediaSqlDump(path=args.path, table=table)
    elif args.source == "html":
        try:
            from aborist.sources import HtmlPageSource
        except ImportError:
            print(
                "html source requires extras: pip install 'aborist[html]'",
                file=sys.stderr,
            )
            return 2
        urls: list[str] = list(args.url or [])
        if args.urls_from:
            urls.extend(
                line.strip()
                for line in Path(args.urls_from).read_text(encoding="utf-8").splitlines()
                if line.strip() and not line.lstrip().startswith("#")
            )
        if not urls:
            print("html source needs --url or --urls-from", file=sys.stderr)
            return 2
        src = HtmlPageSource(urls, respect_robots=not args.no_robots)
    else:
        print(f"unknown source: {args.source}", file=sys.stderr)
        return 2

    conn = connect(args.db)
    try:
        result = ingest_source(conn, src, chunker_name=args.chunker, limit=args.limit)
    finally:
        conn.close()
    print(json.dumps(result.__dict__, indent=2))
    return 0


def _cmd_search(args: argparse.Namespace) -> int:
    conn = connect(args.db)
    try:
        backend = FTS5Backend(conn)
        hits = backend.search(args.query, limit=args.limit)
    finally:
        conn.close()
    if args.json:
        print(
            json.dumps(
                [
                    {
                        "document_root": h.document_root,
                        "document_uri": h.document_uri,
                        "chunk_idx": h.chunk_idx,
                        "snippet": h.snippet,
                        "score": h.score,
                        "audit_mode": h.audit_mode.value,
                        "title": h.title,
                    }
                    for h in hits
                ],
                indent=2,
            )
        )
    else:
        for h in hits:
            print(f"[{h.audit_mode.value}] {h.score:7.3f}  {h.title or h.document_uri}")
            print(f"    chunk {h.chunk_idx}  root={h.document_root[:16]}…")
            print(f"    {h.snippet}")
            print()
    return 0


def _cmd_verify(args: argparse.Namespace) -> int:
    conn = connect(args.db)
    try:
        result = verify_random_sample(conn, n=args.n)
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0 if result["failed"] == 0 else 1


def _cmd_distill(args: argparse.Namespace) -> int:
    from aborist.distill import get_distiller
    from aborist.distill.runner import distill_existing

    try:
        distiller = get_distiller(args.process)
    except ValueError as e:
        print(str(e), file=sys.stderr)
        return 2
    conn = connect(args.db)
    try:
        result = distill_existing(
            conn,
            distiller,
            kind=args.kind,
            source_type=args.source_type,
            limit=args.limit,
            chunker_name=args.chunker,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def _cmd_ask(args: argparse.Namespace) -> int:
    import os

    from aborist.qa import ask
    from aborist.qa.client import OpenAICompatibleClient, StubClient

    base_url = args.endpoint or os.environ.get(
        "ABORIST_LLM_ENDPOINT", "https://uncloseai.com/v1"
    )
    model = args.model or os.environ.get(
        "ABORIST_LLM_MODEL",
        "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
    )
    revision = os.environ.get("ABORIST_LLM_REVISION", "")
    quantization = os.environ.get("ABORIST_LLM_QUANTIZATION", "fp8-dynamic")
    api_key = os.environ.get("ABORIST_LLM_API_KEY")

    client: object
    if args.dry_run:
        client = StubClient(
            answer=f"[STUB] would have answered '{args.question}' against root {args.document_root[:16]}…"
        )
    else:
        client = OpenAICompatibleClient(base_url=base_url, api_key=api_key)

    conn = connect(args.db)
    try:
        result = ask(
            conn,
            document_root=args.document_root,
            question=args.question,
            client=client,
            model_id=model,
            revision=revision,
            quantization=quantization,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0 if result.get("status") in ("cache_hit", "cache_miss_then_written") else 1


def _cmd_providence(args: argparse.Namespace) -> int:
    """List providence_cache records for a document URI or source_root."""
    conn = connect(args.db)
    try:
        if args.document_uri:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "WHERE document_uri = ? ORDER BY created_at DESC",
                (args.document_uri,),
            ).fetchall()
        elif args.source_root:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "WHERE source_root = ? ORDER BY created_at DESC",
                (args.source_root,),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "ORDER BY created_at DESC LIMIT ?",
                (args.limit,),
            ).fetchall()
    finally:
        conn.close()
    out = [
        {
            "cache_key": r["cache_key"],
            "state": r["falsification_state"],
            "hit_count": r["hit_count"],
            "question": r["question_text"],
            "answer": r["answer_text"],
            "created_at": r["created_at"],
        }
        for r in rows
    ]
    print(json.dumps(out, indent=2))
    return 0


def _cmd_evict(args: argparse.Namespace) -> int:
    from aborist.evict import evict_to_cold

    conn = connect(args.db)
    try:
        result = evict_to_cold(
            conn,
            source_type=args.source_type,
            older_than_days=args.older_than_days,
            document_roots=args.document_root or None,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def _cmd_rehydrate(args: argparse.Namespace) -> int:
    from aborist.evict import rehydrate

    conn = connect(args.db)
    try:
        if args.all_cold:
            roots = [
                r["document_root"]
                for r in conn.execute(
                    "SELECT DISTINCT document_root FROM chunks WHERE tier = 'cold'"
                ).fetchall()
            ]
        else:
            roots = list(args.document_root or [])
        if not roots:
            print(
                "rehydrate needs --document-root R or --all-cold",
                file=sys.stderr,
            )
            return 2
        results = []
        for r in roots:
            res = rehydrate(conn, r)
            res["document_root"] = r
            results.append(res)
    finally:
        conn.close()
    print(json.dumps(results, indent=2))
    drift = sum(1 for r in results if r.get("status") == "drift_detected")
    return 1 if drift else 0


def _cmd_stats(args: argparse.Namespace) -> int:
    conn = connect(args.db)
    try:
        result = stats(conn)
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="aborist",
        description="An arborist for trees and forests of cross-linked information.",
    )
    p.add_argument("--version", action="version", version=f"aborist {__version__}")
    p.add_argument(
        "--db",
        type=Path,
        default=DEFAULT_DB_PATH,
        help=f"path to aborist SQLite db (default: {DEFAULT_DB_PATH})",
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    ingest = sub.add_parser("ingest", help="ingest documents from a source")
    ingest.add_argument(
        "--source",
        required=True,
        choices=["wikipedia_cur", "wikipedia_old", "html"],
        help="source type",
    )
    ingest.add_argument("--path", help="path to dump file (file-backed sources)")
    ingest.add_argument(
        "--url", action="append", help="URL to ingest (html source; repeatable)"
    )
    ingest.add_argument(
        "--urls-from",
        dest="urls_from",
        help="file with one URL per line (html source)",
    )
    ingest.add_argument(
        "--no-robots",
        dest="no_robots",
        action="store_true",
        help="do not consult robots.txt (use only for explicitly authorized sites)",
    )
    ingest.add_argument(
        "--chunker", default=None, help="chunker name (default: tok-512-v1)"
    )
    ingest.add_argument(
        "--limit", type=int, default=None, help="cap number of documents"
    )
    ingest.set_defaults(func=_cmd_ingest)

    search = sub.add_parser("search", help="keyword search (VISUAL audit mode)")
    search.add_argument("query", help="query string")
    search.add_argument("--limit", type=int, default=20)
    search.add_argument("--json", action="store_true", help="output JSON")
    search.set_defaults(func=_cmd_search)

    verify = sub.add_parser(
        "verify", help="round-trip Merkle proofs for N random documents"
    )
    verify.add_argument("-n", type=int, default=10)
    verify.set_defaults(func=_cmd_verify)

    distill = sub.add_parser(
        "distill",
        help="compress docs into Merkle-signed cores (surface->core or core->core)",
    )
    distill.add_argument(
        "--process", default="first-sentence-v1", help="distiller name"
    )
    distill.add_argument(
        "--kind",
        choices=["surface", "core"],
        default="surface",
        help="source kind to scan; 'core' runs recursive distillation",
    )
    distill.add_argument(
        "--source-type",
        dest="source_type",
        default=None,
        help="restrict to one source_type",
    )
    distill.add_argument(
        "--chunker", default=None, help="chunker for the core doc"
    )
    distill.add_argument(
        "--limit", type=int, default=None, help="cap number of docs scanned"
    )
    distill.set_defaults(func=_cmd_distill)

    ask_cmd = sub.add_parser(
        "ask",
        help="answer a question about a document (cache-first, STRICT)",
    )
    ask_cmd.add_argument(
        "--document-root",
        dest="document_root",
        required=True,
        help="document_root to ask about",
    )
    ask_cmd.add_argument(
        "--question", required=True, help="question text"
    )
    ask_cmd.add_argument(
        "--model",
        default=None,
        help="model_id (default $ABORIST_LLM_MODEL or hermes-3)",
    )
    ask_cmd.add_argument(
        "--endpoint",
        default=None,
        help="OpenAI-compatible base URL (default $ABORIST_LLM_ENDPOINT)",
    )
    ask_cmd.add_argument(
        "--dry-run",
        dest="dry_run",
        action="store_true",
        help="use StubClient — no network call",
    )
    ask_cmd.set_defaults(func=_cmd_ask)

    prov_cmd = sub.add_parser(
        "providence",
        help="list providence_cache records",
    )
    prov_cmd.add_argument("--document-uri", dest="document_uri", default=None)
    prov_cmd.add_argument("--source-root", dest="source_root", default=None)
    prov_cmd.add_argument("--limit", type=int, default=20)
    prov_cmd.set_defaults(func=_cmd_providence)

    evict_cmd = sub.add_parser(
        "evict",
        help="demote surface chunks hot→cold (NULL content, retain leaf_hash)",
    )
    evict_cmd.add_argument(
        "--source-type",
        dest="source_type",
        default=None,
        help="restrict to one source_type",
    )
    evict_cmd.add_argument(
        "--older-than-days",
        dest="older_than_days",
        type=int,
        default=None,
        help="only evict docs older than N days",
    )
    evict_cmd.add_argument(
        "--document-root",
        action="append",
        default=None,
        help="explicit document_root(s) to evict; repeatable",
    )
    evict_cmd.set_defaults(func=_cmd_evict)

    rehydrate_cmd = sub.add_parser(
        "rehydrate",
        help="refetch URI, verify leaves, restore cold content if root matches",
    )
    rehydrate_cmd.add_argument(
        "--document-root",
        action="append",
        default=None,
        help="explicit document_root(s) to rehydrate; repeatable",
    )
    rehydrate_cmd.add_argument(
        "--all-cold",
        dest="all_cold",
        action="store_true",
        help="rehydrate every document with cold chunks",
    )
    rehydrate_cmd.set_defaults(func=_cmd_rehydrate)

    stats_cmd = sub.add_parser("stats", help="counts: docs, chunks, edges, audit")
    stats_cmd.set_defaults(func=_cmd_stats)

    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return args.func(args)


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