arborist/tests/test_qa_corpus_query.py

"""Unit + functional tests for arborist.qa.corpus_query.run_query.

The single orchestrator both `arborist cloud query` and (eventually)
`arborist query` will share. Tests use SqliteShardCorpus + StubClient
to keep the suite offline + deterministic.

Coverage:
  * Happy path — verbatim quote in StubClient answer reaches the verifier
    on the assembled evidence, lands STRICT.
  * Off-topic — no FTS hits → UNGROUNDED, empty sources, fast return.
  * Timings dict populated with the expected phase keys.
  * Sources annotated with used / used_pointer_ids from [E#] tags in answer.
  * Capacity dict reports sys + evidence + question + answer char counts.
"""
from __future__ import annotations

from typing import Iterator

import pytest

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa.client import StubClient
from arborist.qa.corpus import SqliteShardCorpus
from arborist.qa.corpus_query import run_query
from arborist.source import Source
from arborist.store import connect


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


@pytest.fixture
def corpus(tmp_path):
    db = tmp_path / "shard.db"
    c = connect(db)
    ingest_source(c, _S([
        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 and social. " * 6),
        ),
    ]))
    yield SqliteShardCorpus(c)
    c.close()


# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------


def test_run_query_grounds_verbatim_quote(corpus):
    """Verbatim quote in the stub answer must verify on the assembled
    evidence — STRICT (or higher rung)."""
    stub = (
        'The core principle of anarchism is the "abolition of authority". [E1]'
    )
    result = run_query(
        corpus, "what is the core principle of anarchism?",
        StubClient(answer=stub),
        model_id="stub", top_k=2,
    )
    assert result["audit_mode"] == "STRICT"
    assert result["n_verified"] >= 1
    assert result["sources"]
    assert "Anarchism" in result["sources"][0]["title"]
    # Cited source must be marked used (E1).
    assert result["sources"][0]["used"] is True
    assert "E1" in result["sources"][0]["used_pointer_ids"]


def test_run_query_emits_capacity_and_timings(corpus):
    """Result must carry capacity (prompt char breakdown) + timings
    (per-phase + total) so callers don't have to reverse-engineer."""
    result = run_query(
        corpus, "anarchism",
        StubClient(answer="Anarchism is a philosophy. [E1]"),
        model_id="stub", top_k=1,
    )
    cap = result["capacity"]
    assert cap["sys_prompt_chars"] > 0
    assert cap["evidence_chars"] > 0
    assert cap["prompt_chars"] == (
        cap["sys_prompt_chars"]
        + cap["evidence_chars"]
        + len("EVIDENCE:\n\n\n\nQUESTION: anarchism\n\n")
        # The exact prompt scaffolding char count varies with the
        # grounding-reminder; assert it's positive + plausible.
    ) or cap["prompt_chars"] >= cap["sys_prompt_chars"] + cap["evidence_chars"]
    t = result["timings"]
    assert t["total"] > 0
    for k in ("search", "context", "llm", "verify"):
        assert k in t, f"missing timing phase: {k}"


# ---------------------------------------------------------------------------
# Off-topic / no hits
# ---------------------------------------------------------------------------


def test_run_query_returns_ungrounded_on_no_hits(corpus):
    """Query that no corpus doc matches must return UNGROUNDED quickly
    with empty sources — no LLM call, no verifier-method confusion."""
    stub = StubClient(answer="this should never be returned")
    result = run_query(
        corpus, "what is the price of tea in xylophone-land?",
        stub, model_id="stub", top_k=4,
    )
    assert result["audit_mode"] == "UNGROUNDED"
    assert result["sources"] == []
    # Critical: LLM must NOT have been called when retrieval is empty.
    assert stub.calls == [], "LLM was called despite empty retrieval"


# ---------------------------------------------------------------------------
# Bucket-adapter parity (not run by default; opt-in fixture-heavy)
# ---------------------------------------------------------------------------


def test_run_query_corpus_name_surfaces(corpus):
    """corpus.name leaks into result so bench harnesses can attribute
    which backend produced a row."""
    result = run_query(
        corpus, "anarchism",
        StubClient(answer="Anarchism is a philosophy. [E1]"),
        model_id="stub", top_k=1,
    )
    assert result.get("corpus_name") == "sqlite-shard"