arborist/tests/test_distill_recursive.py

"""Recursive distillation: cores can be distilled into depth+1 cores."""

from __future__ import annotations

import json
from typing import Iterator

from arborist.distill import FirstSentenceDistiller
from arborist.distill.runner import distill_existing
from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.merkle import proof_from_dict, verify_proof
from arborist.source import Source
from arborist.store import connect


class FakeSource(Source):
    source_type = "test"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str) -> Document:
    return Document(uri=uri, content=content, source_type="test", title=uri)


# Multi-paragraph source so distill produces multi-chunk cores.
LONG = (
    "First paragraph one. " * 80
    + "\n\n"
    + "Second paragraph here. " * 80
    + "\n\n"
    + "Third paragraph follows. " * 80
    + "\n\n"
    + "Fourth paragraph closes. " * 80
)


def test_recursive_distill_increments_compression_depth(tmp_path):
    db = tmp_path / "rec.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("test://x", LONG)]))
        d = FirstSentenceDistiller()

        # Round 1: surface -> depth=1 cores
        r1 = distill_existing(conn, d, kind="surface")
        assert r1["distilled"] == 1

        depths = conn.execute(
            "SELECT compression_depth FROM documents WHERE kind='core'"
        ).fetchall()
        assert all(r["compression_depth"] == 1 for r in depths)

        # Round 2: core -> depth=2 cores
        r2 = distill_existing(conn, d, kind="core")
        # The depth=1 core may itself produce a depth=2 core if it's long
        # enough for the chunker. With LONG above, the core has 4 sentences
        # -> at least 1 chunk -> recursive distill yields 1 deeper core.
        assert r2["distilled"] >= 1 or r2["skipped_existing"] >= 1

        if r2["distilled"] >= 1:
            d2 = conn.execute(
                "SELECT compression_depth FROM documents WHERE compression_depth >= 2"
            ).fetchall()
            assert len(d2) >= 1
    finally:
        conn.close()


def test_recursive_distill_proofs_still_verify(tmp_path):
    """A depth=2 core's derivation_proof must reconstruct its depth=1 source."""
    db = tmp_path / "rec_proof.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("test://verify", LONG)]))
        d = FirstSentenceDistiller()
        distill_existing(conn, d, kind="surface")
        distill_existing(conn, d, kind="core")

        # Find any depth=2 derivation if present.
        rows = conn.execute(
            "SELECT der.proof_blob, der.src_root "
            "FROM derivations der "
            "JOIN documents core ON core.document_root = der.core_root "
            "WHERE core.compression_depth = 2"
        ).fetchall()
        if not rows:
            # Acceptable: if the depth=1 core was too short to produce a
            # meaningful depth=2 core, the test is a no-op.
            return

        for row in rows:
            blob = json.loads(row["proof_blob"])
            for entry in blob["contributing"]:
                proof = proof_from_dict(entry["proof"])
                assert verify_proof(proof)
                assert proof.root.hex() == row["src_root"]
    finally:
        conn.close()