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_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_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