arborist/tests/test_distill.py

"""Distillation: cores must Merkle-bind back to their source surface roots."""

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)


# A long, multi-paragraph source so the chunker produces multiple chunks.
LONG_TEXT = (
    "The quick brown fox jumps over the lazy dog. " * 30
    + "\n\n"
    + "Merkle providence proves answer derives from a specific source. " * 30
    + "\n\n"
    + "The eight forms of capital include living, social, and intellectual. " * 30
)


def test_distillation_produces_core_with_verifiable_proofs(tmp_path):
    db = tmp_path / "distill.db"
    src_doc = _doc("test://long", LONG_TEXT)
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([src_doc]))
        result = distill_existing(conn, FirstSentenceDistiller())

        assert result["distilled"] == 1
        assert result["skipped_existing"] == 0

        # Core document exists and is marked correctly.
        cores = conn.execute(
            "SELECT * FROM documents WHERE kind = 'core'"
        ).fetchall()
        assert len(cores) == 1
        core = cores[0]
        assert core["compression_depth"] == 1
        assert core["source_type"] == "core:first-sentence-v1"

        # Derivation row binds core to source.
        derivs = conn.execute("SELECT * FROM derivations").fetchall()
        assert len(derivs) == 1
        d = derivs[0]
        assert d["core_root"] == core["document_root"]
        assert d["process_id"] == "first-sentence-v1"

        # CRITICAL: every contributing-chunk Merkle proof in proof_blob must
        # reconstruct the source's document_root. This is the cryptographic
        # binding of compressed core back to surface.
        proof_data = json.loads(d["proof_blob"])
        assert proof_data["core_root"] == core["document_root"]
        src_root = proof_data["src_root"]
        assert len(proof_data["contributing"]) >= 1
        for entry in proof_data["contributing"]:
            proof = proof_from_dict(entry["proof"])
            assert verify_proof(proof), f"proof invalid for chunk {entry['src_chunk_idx']}"
            assert proof.root.hex() == src_root, (
                "contributing-chunk proof must reconstruct source root"
            )

        # derived_from edge in the forest.
        edges = conn.execute(
            "SELECT * FROM edges WHERE edge_type = 'derived_from'"
        ).fetchall()
        assert len(edges) == 1
        assert edges[0]["src_root"] == core["document_root"]
        assert edges[0]["dst_root"] == src_root

        # Audit chain has both ingest and derive events.
        events = conn.execute(
            "SELECT event_type FROM audit_events ORDER BY seq ASC"
        ).fetchall()
        types = [r["event_type"] for r in events]
        assert "ingest" in types and "derive" in types
    finally:
        conn.close()


def test_distill_idempotent(tmp_path):
    db = tmp_path / "idempotent.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("test://x", LONG_TEXT)]))
        d = FirstSentenceDistiller()
        first = distill_existing(conn, d)
        second = distill_existing(conn, d)
        assert first["distilled"] == 1
        assert second["distilled"] == 0
        assert second["skipped_existing"] == 1
        # Still only one core document.
        n_cores = conn.execute(
            "SELECT COUNT(*) FROM documents WHERE kind='core'"
        ).fetchone()[0]
        assert n_cores == 1
    finally:
        conn.close()


def test_distill_skips_cold_chunks(tmp_path):
    """If any source chunk has been evicted (content NULL), we must skip."""
    db = tmp_path / "cold.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("test://cold", LONG_TEXT)]))
        # Manually evict one chunk to cold tier.
        conn.execute("UPDATE chunks SET content=NULL, tier='cold' WHERE idx=0")
        result = distill_existing(conn, FirstSentenceDistiller())
        assert result["distilled"] == 0
        assert result["skipped_cold"] == 1
    finally:
        conn.close()