arborist/tests/test_qa_corpus_functional.py

"""Functional tests — full retrieval → context → LLM → verify, both adapters.

Builds the smallest plausible "ask-shaped" pipeline that the future
Corpus-aware query() refactor will collapse into:

    hits   = corpus.fts_body(question, limit=k)
    chunks = []
    for h in hits:
        chunks.extend(corpus.chunks_for_doc(h.document_root, limit=1))
    evidence_map = build_evidence_map([...chunks-as-evidence-dicts...])
    prompt = system + EVIDENCE + question + reminder
    answer = chat.chat_completion(...)
    verdict = verify_claim_lattice(answer, evidence_map)

This module drives that pipeline through SqliteShardCorpus AND
SidecarBucketCorpus with a StubClient (deterministic answer) and
asserts both reach the same audit_mode + cite the same primary
source_root.

It's the SMALLEST canary for the "DRY the cloud + local diff" work:
when the refactor moves query() onto the protocol, this test gates
both adapters in one go.
"""
from __future__ import annotations

import os
import shutil
import sqlite3
import threading
from http.server import ThreadingHTTPServer

from arborist.wallet._range_http_server import RangeHandler as _RangeHandler
from pathlib import Path
from typing import Iterator

import pytest

apsw = pytest.importorskip("apsw")

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa.client import StubClient
from arborist.qa.corpus import SidecarBucketCorpus, SqliteShardCorpus
from arborist.qa.evidence import (
    build_evidence_map,
    render_evidence_map,
)
from arborist.qa.prompts import (
    CLAIM_LATTICE_GROUNDING_REMINDER,
    CLAIM_LATTICE_SYSTEM_PROMPT,
)
from arborist.qa.verify import verify_claim_lattice
from arborist.source import Source
from arborist.store import connect
from arborist.wallet.bucket import BucketManifest, MultiShardSidecarCorpus
from arborist.wallet.fts_sidecar_build import build as build_fts_sidecar


# ---------------------------------------------------------------------------
# Fixture: same corpus exposed via both adapters.
# ---------------------------------------------------------------------------


CORPUS_DOCS = [
    Document(
        uri="test://anarchism", source_type="test", title="Anarchism",
        content=(
            "Anarchism is a political philosophy that promotes a stateless "
            "society. " * 6
            + 'The phrase "abolition of authority" describes the core principle '
              'of anarchism. ' * 4
        ),
    ),
    Document(
        uri="test://capital", source_type="test", title="Capital",
        content=(
            "The eight forms of capital include living, social, intellectual, "
            "material, financial, experiential, cultural, and spiritual. " * 6
        ),
    ),
    Document(
        uri="test://noise-1", source_type="test", title="Random Noise",
        content=("Lorem ipsum dolor sit amet. " * 12),
    ),
]


class _FakeSource(Source):
    source_type = "test"
    def __init__(self, ds): self.ds = ds
    def iter_documents(self) -> Iterator[Document]:
        yield from self.ds


@pytest.fixture
def both_adapters(tmp_path: Path):
    """Yield (sqlite_corpus, sidecar_corpus) over the same content."""
    db = tmp_path / "shard.db"
    c0 = connect(db)
    ingest_source(c0, _FakeSource(CORPUS_DOCS))
    c0.execute("PRAGMA wal_checkpoint(FULL)")
    c0.close()

    # bucket layout + slim FTS5 sidecar
    bucket = tmp_path / "bucket"
    clones = bucket / "clones" / "snap-1"
    clones.mkdir(parents=True)
    bucket_shard = clones / "000.db"
    shutil.copy(db, bucket_shard)
    sidecar_path = clones / "000.idx.db"
    build_fts_sidecar(str(bucket_shard), str(sidecar_path), verbose=False)

    # in-process Range-aware server
    cwd_before = os.getcwd()
    os.chdir(bucket)
    httpd = ThreadingHTTPServer(("127.0.0.1", 0), _RangeHandler)
    port = httpd.server_address[1]
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    base_url = f"http://127.0.0.1:{port}"

    manifest = BucketManifest(
        shards=[{
            "url": f"{base_url}/clones/snap-1/000.db",
            "shard_idx": 0,
            "label": "test-shard",
            "fts_sidecar_url": f"{base_url}/clones/snap-1/000.idx.db",
        }],
        blob_base=f"{base_url}/blobs",
        snapshot_root=None,
        snapshot_ts=None,
    )
    multi = MultiShardSidecarCorpus(
        manifest, fts_sidecar_cache_dir=str(tmp_path / "cache"),
    )
    conn = connect(db)

    try:
        yield SqliteShardCorpus(conn), SidecarBucketCorpus(multi)
    finally:
        multi.close()
        conn.close()
        httpd.shutdown()
        os.chdir(cwd_before)


# ---------------------------------------------------------------------------
# The shared end-to-end harness — the prototype of what query() will do
# once it takes a Corpus.
# ---------------------------------------------------------------------------


def _retrieve_and_verify(
    corpus, question: str, *, stub_answer: str, top_k: int = 3,
):
    """Mini retrieval pipeline through the Corpus protocol. Returns a
    dict with audit_mode + cited document_roots + verdict shape."""
    hits = corpus.fts_body(question, limit=top_k)
    if not hits:
        return {"audit_mode": "UNGROUNDED", "n_quotes": 0, "n_verified": 0,
                "sources": [], "cited_roots": set()}

    chunks_payload: list[dict] = []
    for rank, h in enumerate(hits, 1):
        rows = corpus.chunks_for_doc(h.document_root, limit=1)
        if not rows:
            continue
        r = rows[0]
        chunks_payload.append({
            "source_root": h.document_root,
            "document_uri": h.document_uri,
            "title": h.title,
            "chunk_idx": r.idx,
            "chunk_root": r.leaf_hash or "0" * 64,
            "span": r.content[:1200],
            "source_role": (
                "primary_answer_source" if rank == 1 else "background_source"
            ),
        })

    evidence_map = build_evidence_map(chunks_payload)
    evidence_text = render_evidence_map(evidence_map)
    messages = [
        {"role": "system", "content": CLAIM_LATTICE_SYSTEM_PROMPT},
        {"role": "user", "content": (
            f"EVIDENCE:\n\n{evidence_text}\n\n"
            f"QUESTION: {question}\n\n"
            f"{CLAIM_LATTICE_GROUNDING_REMINDER}"
        )},
    ]
    answer = StubClient(answer=stub_answer).chat_completion(
        messages, model="stub", temperature=0.0,
    )
    verdict = verify_claim_lattice(
        answer, evidence_map, question=question, warrant_check_enabled=False,
    )

    cited_roots = {h.document_root for h in hits}
    return {
        "audit_mode": verdict["audit_mode"],
        "n_quotes": verdict["n_quotes"],
        "n_verified": verdict["n_verified"],
        "cited_roots": cited_roots,
        "primary_root": hits[0].document_root if hits else "",
        "primary_title": hits[0].title if hits else "",
    }


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_pipeline_grounds_anarchism_through_both_adapters(both_adapters):
    """Verbatim-quote answer should land at the same audit_mode on both
    adapters and primary-cite the Anarchism doc on both."""
    sq, sc = both_adapters
    question = "what is the core principle of anarchism?"
    # Verbatim from CORPUS_DOCS[0]: 'abolition of authority'
    stub = (
        'The core principle of anarchism is the "abolition of authority". [E1]'
    )
    sq_result = _retrieve_and_verify(sq, question, stub_answer=stub)
    sc_result = _retrieve_and_verify(sc, question, stub_answer=stub)
    assert sq_result["audit_mode"] == sc_result["audit_mode"], (
        f"audit_mode differs: sqlite={sq_result['audit_mode']} "
        f"sidecar={sc_result['audit_mode']}"
    )
    # Both must surface the Anarchism doc as primary.
    assert "Anarchism" in sq_result["primary_title"]
    assert "Anarchism" in sc_result["primary_title"]


def test_pipeline_handles_noise_question_consistently(both_adapters):
    """An off-topic question (no matching doc) must produce UNGROUNDED
    / empty-evidence on both adapters."""
    sq, sc = both_adapters
    question = "what is the population of mars?"
    stub = "Mars has no population because it is uninhabited. [E1]"
    sq_result = _retrieve_and_verify(sq, question, stub_answer=stub)
    sc_result = _retrieve_and_verify(sc, question, stub_answer=stub)
    # Off-topic should land at UNGROUNDED or HYBRID on BOTH paths —
    # neither corpus has Mars content. STRICT would be a false positive.
    for r in (sq_result, sc_result):
        assert r["audit_mode"] in ("UNGROUNDED", "HYBRID"), (
            f"unexpected STRICT on off-topic: {r}"
        )


def test_pipeline_chunks_for_primary_are_same_bytes(both_adapters):
    """Both adapters must surface byte-identical chunk content for the
    same document_root — wallet-side Merkle verification depends on it."""
    sq, sc = both_adapters
    question = "anarchism"
    sq_hits = sq.fts_body(question, limit=1)
    assert sq_hits
    droot = sq_hits[0].document_root

    sq_chunks = sq.chunks_for_doc(droot)
    sc_chunks = sc.chunks_for_doc(droot)
    assert sq_chunks and sc_chunks
    for a, b in zip(sq_chunks, sc_chunks):
        assert a.content == b.content, (
            f"chunk[{a.idx}] content differs by adapter"
        )
        assert a.leaf_hash == b.leaf_hash