arborist/tests/test_query.py

"""Multi-source corpus Q&A: search → context → cache → Hermes.

Stub client only — no network. Validates:
- FTS5 finds the right docs across a small corpus
- context_root is deterministic (sorted source roots, then Merkle)
- cache hit on identical question returns same record without calling client
- different question -> different cache_key
"""

from __future__ import annotations

import json
from typing import Iterator

from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.qa import query
from aborist.qa.client import StubClient
from aborist.source import Source
from aborist.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.rsplit("/", 1)[-1])


# Three docs with distinguishable content so FTS5 can pick winners.
DOCS = [
    _doc(
        "test://anarchism",
        "Anarchism is a political philosophy that opposes the state. " * 12
        + "Mutual aid is central to anarchist theory. " * 8,
    ),
    _doc(
        "test://capitalism",
        "Capitalism is an economic system based on private ownership. " * 12
        + "Market exchange and capital accumulation drive growth. " * 8,
    ),
    _doc(
        "test://anarcho-capitalism",
        "Anarcho-capitalism combines anarchism's opposition to the state with capitalism's private property. " * 12
        + "Murray Rothbard developed many of its core ideas. " * 8,
    ),
]


def test_query_picks_relevant_docs_and_writes_record(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    # Verbatim quote from DOCS[2] → audit_mode=STRICT.
    client = StubClient(
        answer=(
            'Per the source: '
            '"Anarcho-capitalism combines anarchism\'s opposition to the state '
            'with capitalism\'s private property"'
        )
    )
    result = query(
        question="What is anarcho-capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="test-model",
        single_db=main_db,
        top_k=3,
    )

    assert result["status"] == "cache_miss_then_written"
    assert result["audit_mode"] == "STRICT"
    assert "Anarcho-capitalism" in result["answer_text"]
    assert len(result["sources"]) >= 1
    assert any("anarcho-capitalism" in s["document_uri"] for s in result["sources"])

    # context_root is deterministic.
    sorted_roots = sorted(s["document_root"] for s in result["sources"])
    if len(sorted_roots) == 1:
        assert result["context_root"] == sorted_roots[0]
    # Repeat call with same question → cache hit, no LLM call.
    n_calls_before = len(client.calls)
    result2 = query(
        question="What is anarcho-capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="test-model",
        single_db=main_db,
        top_k=3,
    )
    assert result2["status"] == "cache_hit"
    assert result2["cache_key"] == result["cache_key"]
    assert len(client.calls) == n_calls_before  # no new call


def test_query_different_question_different_cache(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    client = StubClient(answer="answer text")
    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    r2 = query(
        question="What is capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    assert r1["cache_key"] != r2["cache_key"]
    assert len(client.calls) == 2  # both missed cache, both called client


def test_query_question_equivalence_class_dedups_cache(tmp_path):
    """Fox 2026-04-29 catch: 'who is batman?', 'who is batman', and
    'who is the batman?' should all hit the same cache. Question_hash
    canonicalizes correctly, but conversation_hash used to hash the
    LITERAL question text in the user message — so each variant got
    its own chash and missed cache. Fix: canonical_question form is
    substituted into the messages list used for conversation_hash,
    while the LLM still receives the verbatim question."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    captured: list = []

    def _capture(messages, **kw):
        captured.append(messages)
        return "stub answer"

    # First variant — populates cache.
    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
    )
    # Second & third variants — must hit the same cache_key.
    r2 = query(
        question="What is anarchism",  # no trailing ?
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
    )
    r3 = query(
        question="what is the anarchism?",  # leading article
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
    )

    assert r1["cache_key"] == r2["cache_key"] == r3["cache_key"]
    assert r1["status"] == "cache_miss_then_written"
    assert r2["status"] == "cache_hit"
    assert r3["status"] == "cache_hit"
    # Only the first call reached the LLM.
    assert len(captured) == 1
    # And it received the verbatim question, not the canonical form.
    user_text = "\n".join(
        m["content"] for m in captured[0] if m["role"] == "user"
    )
    assert "What is anarchism?" in user_text


def test_query_strict_dedup_distinguishes_question_variants(tmp_path):
    """policy['question_dedup']='strict' makes every variant get its
    own cache_key. 'Who is X?' and 'who is the X?' write separate
    records under strict policy."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    strict_policy = dict(query.__globals__["DEFAULT_QUERY_POLICY"])
    strict_policy["question_dedup"] = "strict"

    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="a"),
        model_id="m",
        single_db=main_db,
        policy=strict_policy,
        fidelity="strict",
    )
    r2 = query(
        question="what is anarchism",
        qa_db=qa_db,
        chat_client=StubClient(answer="b"),
        model_id="m",
        single_db=main_db,
        policy=strict_policy,
        fidelity="strict",
    )
    # Different cache_keys; both populated independently.
    assert r1["cache_key"] != r2["cache_key"]
    assert r1["status"] == r2["status"] == "cache_miss_then_written"
    assert r1["lookup_path"] == "miss"


def test_query_equivalence_class_fidelity_falls_back_across_dedup_modes(tmp_path):
    """A record written under equivalence_class policy gets reused by a
    later strict-policy lookup that asks for fidelity='equivalence_class'.
    Verifies the cross-silo fallback: strict ckey misses, alternate
    equivalence_class ckey hits, lookup_path reports the fallback."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    DEFAULT = query.__globals__["DEFAULT_QUERY_POLICY"]
    eq_policy = dict(DEFAULT)
    eq_policy["question_dedup"] = "equivalence_class"
    strict_policy = dict(DEFAULT)
    strict_policy["question_dedup"] = "strict"

    # Agent A writes equivalence-class record for 'what is anarchism?'.
    captured: list = []

    def _capture(messages, **kw):
        captured.append(messages)
        return "stub"

    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
        policy=eq_policy,
    )
    assert r1["status"] == "cache_miss_then_written"

    # Agent B (strict policy, equivalence_class fidelity) asks the same
    # question. Strict ckey misses; the eq_class fallback hits A's record.
    r2 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
        policy=strict_policy,
        fidelity="equivalence_class",
    )
    assert r2["status"] == "cache_hit"
    assert r2["lookup_path"] == "equivalence_class_fallback"
    assert len(captured) == 1  # only A's call reached the LLM


def test_query_strict_fidelity_does_not_fall_back(tmp_path):
    """Audit-grade lookup: strict-fidelity refuses to read records from
    the other dedup mode's silo. Even if equivalence_class has a hit,
    strict fidelity reports cache miss."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    DEFAULT = query.__globals__["DEFAULT_QUERY_POLICY"]
    eq_policy = dict(DEFAULT)
    eq_policy["question_dedup"] = "equivalence_class"
    strict_policy = dict(DEFAULT)
    strict_policy["question_dedup"] = "strict"

    # Agent A writes under equivalence_class.
    query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="a"),
        model_id="m",
        single_db=main_db,
        policy=eq_policy,
    )

    # Agent B (strict policy + strict fidelity) asks same question.
    # Should NOT find A's record; runs LLM fresh.
    captured: list = []

    def _capture(messages, **kw):
        captured.append(messages)
        return "stub"

    r2 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
        policy=strict_policy,
        fidelity="strict",
    )
    assert r2["status"] == "cache_miss_then_written"
    assert r2["lookup_path"] == "miss"
    assert len(captured) == 1


def test_classify_source_role_separates_primary_from_noisy():
    """Direct unit test on the role classifier. JP film-score should be
    noisy_background; JP (film) should be primary; JP franchise should
    be secondary; The Lost World should be sequel; off-topic background.
    Catches the case where peripheral pages with strong title overlap
    used to share the primary slot."""
    from aborist.qa.query import _classify_source_role

    qstems = {"dinosaur", "jurassic", "park", "film"}
    assert _classify_source_role("Jurassic Park (film)", qstems) == "primary_answer_source"
    assert _classify_source_role("Jurassic Park (film score)", qstems) == "noisy_background_source"
    assert _classify_source_role("Jurassic Park video games", qstems) == "noisy_background_source"
    assert _classify_source_role("Jurassic Park (franchise)", qstems) == "secondary_context_source"
    assert _classify_source_role("List of Jurassic Park characters", qstems) == "secondary_context_source"
    assert _classify_source_role("The Lost World: Jurassic Park", qstems) == "sequel_background_source"
    # Off-topic title (no shared stems): falls through to background.
    assert _classify_source_role("Anarchism", qstems) == "background_source"


def test_query_role_weighted_budget_persists_role_on_sources(tmp_path):
    """Each source in the providence record's merkle_proof.sources gains
    a `source_role` field — verifies the role made it into the audit
    trail so an inspector can see which slot a source occupied."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()
    result = query(
        question="What is anarcho-capitalism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="x"),
        model_id="m",
        single_db=main_db,
        top_k=3,
    )
    assert all("source_role" in s for s in result["sources"])


def test_query_persists_run_dag_root_on_record_and_result(tmp_path):
    """Per-run Merkle-DAG fingerprint surfaces on both the result dict
    & the persisted providence_cache row. Recomputing the root from
    the persisted blob matches what was stored."""
    from aborist.qa.dag import verify_run_dag

    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    result = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer='The source: "a political philosophy that opposes the state"'),
        model_id="m",
        single_db=main_db,
    )
    assert "run_dag_root" in result
    assert isinstance(result["run_dag_root"], str)
    assert len(result["run_dag_root"]) == 64  # sha256 hex

    qa_conn = connect(qa_db)
    try:
        row = qa_conn.execute(
            "SELECT run_dag_root, run_dag_blob FROM providence_cache "
            "WHERE cache_key = ?",
            (result["cache_key"],),
        ).fetchone()
    finally:
        qa_conn.close()
    assert row["run_dag_root"] == result["run_dag_root"]
    assert verify_run_dag(row["run_dag_blob"]) is True


def test_query_no_sources_when_empty_corpus(tmp_path):
    main_db = tmp_path / "empty.db"
    qa_db = tmp_path / "qa.db"
    connect(main_db).close()  # creates schema, no docs
    result = query(
        question="What is anything?",
        qa_db=qa_db,
        chat_client=StubClient(),
        model_id="m",
        single_db=main_db,
    )
    assert result["status"] == "no_sources"


def test_query_filter_requires_breadth_for_multi_token_queries(tmp_path):
    """Fox 2026-04-29 catch: 'supermans girlfriend' returned 7-of-8
    unrelated `Girlfriends`-titled articles because title-overlap
    accepted ANY single-token match. The fix tightens both title and
    body filters to require ALL query tokens (≤2-token queries) so a
    doc whose title only matches ONE of the two qtokens doesn't pass.

    Synthetic corpus pins the new behavior:

      - "Lois Lane"        — neither qtoken in title; body has both → keep
      - "Girlfriends"      — only "girlfriend" in title/body; no "superman" → drop
      - "Superman album"   — "superman" in title; body lacks "girlfriend" → drop
    """
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    docs = [
        _doc(
            "test://lois-lane",
            "Lois Lane is a fictional character who works for the Daily Planet. "
            "She is Superman's girlfriend and frequently appears in his stories. "
            * 5,
        ),
        _doc(
            "test://girlfriends-tv",
            "Girlfriends is a sitcom about four women in Los Angeles. " * 10,
        ),
        _doc(
            "test://superman-music",
            "Superman is an album of rock music recorded in Tokyo. " * 10,
        ),
    ]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()

    result = query(
        question="who is supermans girlfriend?",
        qa_db=qa_db,
        chat_client=StubClient(answer="x"),
        model_id="m",
        single_db=main_db,
        top_k=8,
    )
    src_uris = [s["document_uri"] for s in result["sources"]]
    # Lois Lane MUST be in the result set — it's the only doc with both
    # query tokens in its body.
    assert any("lois-lane" in u for u in src_uris), (
        f"breadth filter regressed: lois-lane not in {src_uris}"
    )


def test_query_filter_one_token_query_still_synonym_expands(tmp_path):
    """1-token queries keep the loose synonym-expanded any-match — pin
    that we didn't over-tighten the single-token case."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    docs = [
        _doc("test://anarchism", "Anarchism is a political philosophy. " * 30),
    ]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()
    result = query(
        question="anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="x"),
        model_id="m",
        single_db=main_db,
        top_k=3,
    )
    assert len(result["sources"]) == 1


def test_query_burn_existing_forces_fresh_inference(tmp_path):
    """Fox 2026-04-29: `make query Q=... BURN=1` busts any matching live
    cache record before lookup so a fresh inference runs.

    Sequence:
      1. First query → cache_miss_then_written, populates cache
      2. Second query, NO burn → cache_hit (no new LLM call)
      3. Third query, BURN=True → cache_miss_then_written (cache busted),
         result reports burned_existing=1
    """
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    client = StubClient(answer="anarchism is a thing")

    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    assert r1["status"] == "cache_miss_then_written"
    assert r1["burned_existing"] == 0
    n_calls_after_first = len(client.calls)

    r2 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    assert r2["status"] == "cache_hit"
    assert r2["burned_existing"] == 0
    assert len(client.calls) == n_calls_after_first  # no new LLM call

    r3 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
        burn_existing=True,
    )
    assert r3["status"] == "cache_miss_then_written"
    assert r3["burned_existing"] == 1
    assert len(client.calls) == n_calls_after_first + 1   # new LLM call after burn


def test_query_burn_existing_writes_audit_event(tmp_path):
    """Each --burn writes a providence_burn audit event so the chain
    records the bust. Verifies one event lands per burn."""
    from aborist.store import connect as _connect

    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    client = StubClient(answer="x")
    query(
        question="What is anarchism?",
        qa_db=qa_db, chat_client=client, model_id="m", single_db=main_db,
    )
    qc = _connect(qa_db)
    try:
        burns_before = qc.execute(
            "SELECT COUNT(*) FROM audit_events WHERE event_type='providence_burn'"
        ).fetchone()[0]
    finally:
        qc.close()

    query(
        question="What is anarchism?",
        qa_db=qa_db, chat_client=client, model_id="m", single_db=main_db,
        burn_existing=True,
    )
    qc = _connect(qa_db)
    try:
        burns_after = qc.execute(
            "SELECT COUNT(*) FROM audit_events WHERE event_type='providence_burn'"
        ).fetchone()[0]
    finally:
        qc.close()
    assert burns_after == burns_before + 1


def test_query_burn_existing_with_no_prior_record_is_noop(tmp_path):
    """First-time query with --burn: nothing to burn → burned_existing=0,
    proceeds to fresh inference normally."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()
    r = query(
        question="What is anarchism?",
        qa_db=qa_db, chat_client=StubClient(answer="x"),
        model_id="m", single_db=main_db,
        burn_existing=True,
    )
    assert r["status"] == "cache_miss_then_written"
    assert r["burned_existing"] == 0


def test_query_per_source_cap_prevents_huge_doc_monopoly(tmp_path):
    """Fox 2026-04-29 catch: a top-ranked huge document (e.g.
    List_of_Batman_comics, 80 KB+ bibliography) used to consume the
    entire 60 KB budget at hit #1, dropping every subsequent doc with
    char_budget <= 0. Now each of the top_k hits gets at most
    `max_context_chars / top_k` chars; multiple sources land in
    context even when hit #1 is huge."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"

    # Hit #1 is intentionally huge: every query token AND the most copies
    # so it FTS5-ranks first. Hits #2 and #3 are smaller but still
    # contain the query token.
    bulk_token = "anarchism " * 5000   # ~50 KB after canonicalize
    docs = [
        _doc("test://huge-bibliography", bulk_token),
        _doc("test://anarchism-bio", "Anarchism is a political philosophy. " * 30),
        _doc("test://anarchism-history", "Anarchism's history begins with Proudhon. " * 30),
    ]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()

    captured = {"messages": None}

    def _capture(messages, **kw):
        captured["messages"] = messages
        return "stub"

    result = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=_capture),
        model_id="m",
        single_db=main_db,
        top_k=3,
        max_context_chars=60000,
    )
    # All three sources should be in the result, not just the huge one.
    assert len(result["sources"]) >= 2, (
        f"per-source cap broke: only {len(result['sources'])} sources reached "
        f"context (huge doc monopolized again)"
    )
    # Verify the user-turn context contains content from the smaller docs
    # too — not just a 60K slab of the bulk doc.
    user_text = "\n".join(m["content"] for m in captured["messages"] if m["role"] == "user")
    assert "anarchism-bio" in user_text or "anarchism-history" in user_text


def test_query_per_source_cap_respects_top_k(tmp_path):
    """top_k=1 → cap = max_context_chars (legacy behavior preserved when
    operator wants a single large source)."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    big = _doc("test://big", "Anarchism is a political philosophy. " * 2000)
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource([big]))
    finally:
        conn.close()
    result = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="x"),
        model_id="m",
        single_db=main_db,
        top_k=1,
        max_context_chars=60000,
    )
    # Single source allowed up to full budget.
    assert len(result["sources"]) == 1


def test_query_persists_audit_event(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()
    query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="X"),
        model_id="m",
        single_db=main_db,
    )
    # qa.db should now have one providence_query event in its audit chain.
    qc = connect(qa_db)
    try:
        events = qc.execute(
            "SELECT event_type FROM audit_events ORDER BY seq"
        ).fetchall()
    finally:
        qc.close()
    types = [e["event_type"] for e in events]
    assert "providence_query" in types