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"


# ---------------------------------------------------------------------------
# policy= kwarg (Phase 1 step 6a of #53)
# ---------------------------------------------------------------------------


def test_run_query_no_policy_default_behavior(corpus):
    """policy=None (default) must produce the SAME verdict as before
    step 6a — this is the byte-identity gate at the orchestrator level."""
    answer = 'Anarchism is a "stateless society". [E1]'
    r_no_policy = run_query(
        corpus, "anarchism", StubClient(answer=answer),
        model_id="stub", top_k=1,
    )
    # Same call again to confirm determinism + capture the legacy shape.
    r_again = run_query(
        corpus, "anarchism", StubClient(answer=answer),
        model_id="stub", top_k=1,
    )
    assert r_no_policy["audit_mode"] == r_again["audit_mode"]
    assert r_no_policy["n_verified"] == r_again["n_verified"]
    assert r_no_policy["raw_answer"] == r_again["raw_answer"]


def test_run_query_policy_forwards_verifier_kwarg(corpus):
    """policy={"max_claims_per_answer": 0} should trip TOO_MANY_CLAIMS
    on any non-empty answer (default cap is 12). Proves policy actually
    reaches the verifier."""
    answer = 'Anarchism is a "stateless society". [E1]'
    r_strict = run_query(
        corpus, "anarchism", StubClient(answer=answer),
        model_id="stub", top_k=1,
    )
    r_capped = run_query(
        corpus, "anarchism", StubClient(answer=answer),
        model_id="stub", top_k=1,
        policy={"max_claims_per_answer": 0},
    )
    # Default-policy run should STRICT or higher; capped run should
    # have at least one TOO_MANY_CLAIMS violation.
    assert r_strict["audit_mode"] in ("STRICT", "HYBRID")
    violations = r_capped.get("violations") or []
    assert any(
        v.get("kind") == "TOO_MANY_CLAIMS" or v.get("violation_type") == "TOO_MANY_CLAIMS"
        for v in violations
    ), f"expected TOO_MANY_CLAIMS violation in capped run; got {violations!r}"


def test_run_query_policy_classifies_source_role(corpus):
    """With policy provided, hit roles come from the title-based
    classifier (arborist.qa.source_roles.classify_source_role), not
    a flat 'rank 1 = primary, else = background' default. For the
    'anarchism' query and the Anarchism fixture title, the top hit
    must classify as primary_answer_source either way (matches
    the legacy behavior for single-token clean-title queries)."""
    result = run_query(
        corpus, "anarchism",
        StubClient(answer="Anarchism is a philosophy. [E1]"),
        model_id="stub", top_k=2, policy={},
    )
    src = result["sources"][0]
    # Classifier identifies clean Anarchism title as primary.
    assert src["source_role"] == "primary_answer_source"


def test_run_query_multi_route_merges_body_and_title(corpus):
    """policy={"multi_route": True} fans body + title + phrase +
    core_keyword. core_keyword raises NotSupportedError on
    SqliteShardCorpus when the shard has no derivations — fail-open
    means the merge still produces hits from body and title.

    For the Anarchism fixture (no derivations), this test just
    confirms multi-route doesn't blow up and the merged hit set is
    non-empty."""
    result = run_query(
        corpus, "anarchism",
        StubClient(answer="Anarchism is a philosophy. [E1]"),
        model_id="stub", top_k=2,
        policy={"multi_route": True},
    )
    assert result["sources"]
    assert any("Anarchism" in s["title"] for s in result["sources"])


def test_run_query_policy_ignores_unknown_keys(corpus):
    """policy with verifier-irrelevant keys (e.g. base_version that
    Phase 1 step 6a doesn't honor yet) must not blow up — unknown keys
    are silently ignored so a legacy DEFAULT_QUERY_POLICY can be
    passed without filtering."""
    answer = 'Anarchism is a "stateless society". [E1]'
    result = run_query(
        corpus, "anarchism", StubClient(answer=answer),
        model_id="stub", top_k=1,
        policy={
            "base_version": "wikitext-base-v1",   # ignored (step 6d work)
            "some_future_field": 42,              # ignored
            "max_claims_per_answer": 12,          # honored — default
        },
    )
    assert result["audit_mode"] in ("STRICT", "HYBRID", "UNGROUNDED")