arborist/aborist/distill/runner.py

"""Distill existing surface (or core) docs into deeper cores.

Per source row, the runner:

1. Loads the source's chunks (must be hot tier — content present).
2. Runs the Distiller (pure function: source -> core artifact).
3. Computes the core's own Merkle tree.
4. Generates a Merkle proof for every contributing source chunk against the
   source's document_root and packs them into proof_blob (JSON).

These are read-only / pure steps. Writes are accumulated into batches and
flushed in a single transaction per batch — same fsync-amortization approach
as ingest. With WAL+synchronous=NORMAL this gives ~10x throughput.

Idempotent: re-running with the same distiller skips existing core_root.
"""

from __future__ import annotations

import json
import sqlite3
import time
from dataclasses import dataclass

from aborist import (
    CANONICALIZATION_VERSION,
    SCHEMA_VERSION,
)
from aborist.compress import unpack_chunk
from aborist.distill.base import Distiller
from aborist.document import Document, canonicalize, get_chunker
from aborist.merkle import MerkleTree, hash_leaf, proof_to_dict
from aborist.store import append_audit, transaction


DEFAULT_BATCH_SIZE = 200


@dataclass
class _Pending:
    """Pre-computed distillation artifacts ready for a batched DB write."""

    src_root: str
    src_uri: str
    src_compression_depth: int
    core_uri: str
    core_source_type: str
    core_title: str | None
    core_root: str
    core_chunk_strs: list[str]
    core_leaves: list[bytes]
    core_tree: MerkleTree
    proof_blob: str
    src_chars: int
    core_chars: int
    contributing_count: int


def distill_existing(
    conn: sqlite3.Connection,
    distiller: Distiller,
    *,
    kind: str = "surface",
    source_type: str | None = None,
    limit: int | None = None,
    chunker_name: str | None = None,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict:
    """Distill documents of `kind` already in the store. Returns counts.

    `kind='surface'` (default) compresses ingested docs into depth=1 cores.
    `kind='core'` runs recursive distillation: a depth=N core compresses
    into a depth=N+1 core.
    """
    if kind not in ("surface", "core"):
        raise ValueError("kind must be 'surface' or 'core'")
    chunker = get_chunker(chunker_name)

    where_clauses = ["kind = ?"]
    params: list = [kind]
    if source_type:
        where_clauses.append("source_type = ?")
        params.append(source_type)
    sql = (
        "SELECT document_root, document_uri, title, source_type, compression_depth "
        "FROM documents WHERE " + " AND ".join(where_clauses) + " ORDER BY ingest_ts ASC"
    )
    if limit:
        sql += f" LIMIT {int(limit)}"
    src_rows = conn.execute(sql, params).fetchall()

    counters = {
        "scanned": len(src_rows),
        "distilled": 0,
        "skipped_existing": 0,
        "skipped_cold": 0,
        "skipped_empty": 0,
    }
    batch: list[_Pending] = []

    def flush() -> None:
        if not batch:
            return
        with transaction(conn):
            for entry in batch:
                outcome = _persist_no_tx(conn, entry, distiller.name, chunker.name, kind)
                counters[outcome] += 1
        batch.clear()

    for s in src_rows:
        src_root = s["document_root"]
        chunk_rows = conn.execute(
            "SELECT idx, leaf_hash, content FROM chunks "
            "WHERE document_root = ? ORDER BY idx ASC",
            (src_root,),
        ).fetchall()
        if any(r["content"] is None for r in chunk_rows):
            counters["skipped_cold"] += 1
            continue

        # Surface chunks may be zstd-compressed (post April 2026); cores
        # are stored plaintext so SQL LOWER+LIKE keeps working over them.
        # Unpack here so the distiller sees clean strings either way.
        chunk_strs = [unpack_chunk(r["content"]) or "" for r in chunk_rows]
        src_doc = Document(
            uri=s["document_uri"],
            content="",
            source_type=s["source_type"],
            title=s["title"],
        )
        result = distiller.distill(src_doc, chunk_strs)

        core_text = canonicalize(result.core.content)
        core_chunk_strs = chunker.split(core_text)
        if not core_chunk_strs:
            counters["skipped_empty"] += 1
            continue

        core_leaves = [hash_leaf(c.encode("utf-8")) for c in core_chunk_strs]
        core_tree = MerkleTree.build(core_leaves)
        core_root = core_tree.root.hex()

        src_leaves = [bytes.fromhex(r["leaf_hash"]) for r in chunk_rows]
        src_tree = MerkleTree.build(src_leaves)
        contributing: list[dict] = []
        for idx in result.contributing_chunk_indices:
            if 0 <= idx < len(src_leaves):
                p = src_tree.proof(idx)
                contributing.append(
                    {"src_chunk_idx": idx, "proof": proof_to_dict(p)}
                )
        proof_blob = json.dumps(
            {
                "process_id": distiller.name,
                "core_root": core_root,
                "src_root": src_root,
                "contributing": contributing,
            },
            separators=(",", ":"),
            sort_keys=True,
        )

        batch.append(
            _Pending(
                src_root=src_root,
                src_uri=s["document_uri"],
                src_compression_depth=s["compression_depth"] or 0,
                core_uri=result.core.uri,
                core_source_type=result.core.source_type,
                core_title=result.core.title,
                core_root=core_root,
                core_chunk_strs=core_chunk_strs,
                core_leaves=core_leaves,
                core_tree=core_tree,
                proof_blob=proof_blob,
                src_chars=sum(len(c) for c in chunk_strs),
                core_chars=sum(len(c) for c in core_chunk_strs),
                contributing_count=len(contributing),
            )
        )
        if len(batch) >= batch_size:
            flush()
    flush()
    return counters


def _persist_no_tx(
    conn: sqlite3.Connection,
    p: _Pending,
    process_id: str,
    chunker_name: str,
    src_kind: str,
) -> str:
    """Write one pending distillation. Caller wraps the batch in a transaction.

    Returns the counter key to bump: 'distilled' or 'skipped_existing'.
    """
    existing = conn.execute(
        "SELECT 1 FROM documents WHERE document_root = ?", (p.core_root,)
    ).fetchone()
    if existing:
        # Same core_root from another source — record the (core, src) edge
        # and the alternative derivation row, but don't re-insert the doc.
        now = int(time.time())
        conn.execute(
            "INSERT OR IGNORE INTO derivations "
            "(core_root, src_root, proof_blob, process_id, distilled_at) "
            "VALUES (?, ?, ?, ?, ?)",
            (p.core_root, p.src_root, p.proof_blob, process_id, now),
        )
        conn.execute(
            "INSERT OR IGNORE INTO edges "
            "(src_root, dst_root, dst_uri, edge_type, anchor) "
            "VALUES (?, ?, ?, 'derived_from', '')",
            (p.core_root, p.src_root, p.src_uri),
        )
        return "skipped_existing"

    now = int(time.time())
    new_depth = p.src_compression_depth + 1
    conn.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, compression_depth, "
        " title, chunking_version, canonicalization_version, schema_version, "
        " ingest_ts) "
        "VALUES (?, ?, ?, 'core', ?, ?, ?, ?, ?, ?)",
        (
            p.core_root,
            p.core_uri,
            p.core_source_type,
            new_depth,
            p.core_title,
            chunker_name,
            CANONICALIZATION_VERSION,
            SCHEMA_VERSION,
            now,
        ),
    )
    # Reserve a contiguous chunk_id range for this core's chunks; the
    # FTS5 rowid must equal chunks.chunk_id so the search-time JOIN works.
    next_id = (
        conn.execute("SELECT COALESCE(MAX(chunk_id), 0) FROM chunks")
        .fetchone()[0]
        + 1
    )
    chunk_ids = list(range(next_id, next_id + len(p.core_chunk_strs)))
    # Cores stay plaintext (small comma-separated keyword strings); SQL
    # LOWER+LIKE in qa.query._docs_with_core_keyword_match relies on this.
    conn.executemany(
        "INSERT INTO chunks (chunk_id, document_root, idx, leaf_hash, content) VALUES (?, ?, ?, ?, ?)",
        [
            (chunk_ids[i], p.core_root, i, p.core_leaves[i].hex(), p.core_chunk_strs[i])
            for i in range(len(p.core_chunk_strs))
        ],
    )
    conn.executemany(
        "INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)",
        [(chunk_ids[i], p.core_chunk_strs[i]) for i in range(len(p.core_chunk_strs))],
    )
    interior: list[tuple] = []
    for layer_idx in range(1, len(p.core_tree.layers)):
        for node_idx, h in enumerate(p.core_tree.layers[layer_idx]):
            interior.append((p.core_root, layer_idx, node_idx, h.hex()))
    if interior:
        conn.executemany(
            "INSERT INTO merkle_nodes (document_root, layer, idx, hash) VALUES (?, ?, ?, ?)",
            interior,
        )
    conn.execute(
        "INSERT INTO derivations "
        "(core_root, src_root, proof_blob, process_id, distilled_at) "
        "VALUES (?, ?, ?, ?, ?)",
        (p.core_root, p.src_root, p.proof_blob, process_id, now),
    )
    conn.execute(
        "INSERT OR IGNORE INTO edges "
        "(src_root, dst_root, dst_uri, edge_type, anchor) "
        "VALUES (?, ?, ?, 'derived_from', '')",
        (p.core_root, p.src_root, p.src_uri),
    )
    append_audit(
        conn,
        event_type="derive",
        subject_root=p.core_root,
        body={
            "src_root": p.src_root,
            "process_id": process_id,
            "src_kind": src_kind,
            "compression_depth": new_depth,
            "core_chunks": len(p.core_chunk_strs),
            "src_chunks_used": p.contributing_count,
            "src_chars": p.src_chars,
            "core_chars": p.core_chars,
            "compression_ratio": (
                round(p.core_chars / p.src_chars, 4) if p.src_chars else 0.0
            ),
        },
        ts=now,
    )
    return "distilled"