arborist/tests/test_llm_context_base.py

"""LLM context is wikitext-stripped before send.

Before this lands, runner.ask() and query.query() handed Hermes raw
[[wikitext]] markup. The model paraphrased it into clean prose then
labeled the result as a verbatim quote — verifier correctly flagged
those as UNGROUNDED. With base_version in policy, both paths run
to_base() on the assembled context before message construction; Hermes
sees prose so it CAN quote verbatim, and the verifier compares prose-
to-prose end-to-end (idempotent — verify_quotes also runs to_base).

These tests use StubClient to capture exactly what messages would be
sent to the LLM, asserting the wikitext markup is gone before the
model would see it.
"""

from __future__ import annotations

from pathlib import Path

import pytest

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa.client import StubClient
from arborist.qa.runner import DEFAULT_POLICY, ask
from arborist.store import connect
from arborist.wikitext import BASE_VERSION


WIKITEXT_DOC = (
    "{{Refimprove|date=October 2009}}\n"
    "[[File:foo.jpg|thumb|caption]]\n"
    "''[[Final Fantasy VII]]'' follows [[protagonist]] [[Cloud Strife]], "
    "a troubled mercenary.<ref>Smith 2010</ref> "
    "He joins [[AVALANCHE]] to stop [[Shinra Electric Power Company|Shinra]]."
)


class _OneDocSource:
    source_type = "llm_ctx_test"

    def __init__(self, content: str):
        self._doc = Document(
            uri="https://example.com/ff7",
            content=content,
            source_type=self.source_type,
            title="Test",
        )

    def iter_documents(self):
        yield self._doc


def _ingest_one(db_path: Path, content: str) -> str:
    conn = connect(db_path)
    try:
        ingest_source(conn, _OneDocSource(content))
        row = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri=?",
            ("https://example.com/ff7",),
        ).fetchone()
    finally:
        conn.close()
    return row["document_root"]


# ---------------------------------------------------------------------------
# DEFAULT_POLICY carries base_version
# ---------------------------------------------------------------------------


def test_default_policy_includes_base_version():
    """If base_version disappears from default policy, the LLM goes back
    to seeing raw wikitext silently. Tripwire."""
    assert DEFAULT_POLICY["base_version"] == BASE_VERSION


# ---------------------------------------------------------------------------
# Stripped context flows to the LLM (ask path)
# ---------------------------------------------------------------------------


def _user_messages_text(client: StubClient) -> str:
    """Concatenate every user-role message from the LAST captured call."""
    last = client.calls[-1]
    return "\n".join(m["content"] for m in last["messages"] if m["role"] == "user")


def test_ask_passes_base_form_to_llm(tmp_path):
    db = tmp_path / "qa.db"
    root = _ingest_one(db, WIKITEXT_DOC)

    client = StubClient(answer="I don't know based on the provided document.")
    conn = connect(db)
    try:
        ask(
            conn,
            document_root=root,
            question="who is the protagonist?",
            client=client,
            model_id="stub-model",
        )
    finally:
        conn.close()

    sent = _user_messages_text(client)
    # Wikitext markers MUST be gone.
    assert "[[" not in sent
    assert "]]" not in sent
    assert "{{Refimprove" not in sent
    assert "<ref>" not in sent
    assert "File:" not in sent
    # Topical content survives.
    assert "Cloud Strife" in sent
    assert "Final Fantasy VII" in sent
    assert "AVALANCHE" in sent


def test_ask_with_base_version_disabled_passes_raw_wikitext(tmp_path):
    """Operator escape hatch: setting base_version=None reverts to raw
    wikitext context. Confirms the gate is effective and reversible."""
    db = tmp_path / "qa.db"
    root = _ingest_one(db, WIKITEXT_DOC)

    custom_policy = dict(DEFAULT_POLICY)
    custom_policy["base_version"] = None

    client = StubClient(answer="stub")
    conn = connect(db)
    try:
        ask(
            conn,
            document_root=root,
            question="who is the protagonist?",
            client=client,
            model_id="stub-model",
            policy=custom_policy,
        )
    finally:
        conn.close()

    sent = _user_messages_text(client)
    # Raw wikitext markers should now BE present in what the LLM saw.
    assert "[[" in sent
    assert "]]" in sent


# ---------------------------------------------------------------------------
# governance_policy_hash flows base_version (cache_key separation)
# ---------------------------------------------------------------------------


def test_base_version_changes_governance_hash(tmp_path):
    """Two policies that differ only in base_version produce different
    governance_policy_hash values. This is what keeps prior cached
    answers (under raw-wikitext policy) from satisfying lookups under
    the new (stripped) policy on the same question/source — they remain
    distinct cache_keys."""
    from arborist.qa.keys import governance_policy_hash

    p_with = dict(DEFAULT_POLICY)
    p_with["base_version"] = "wikitext-base-v1"
    p_without = dict(DEFAULT_POLICY)
    p_without["base_version"] = None

    h_with = governance_policy_hash(p_with)
    h_without = governance_policy_hash(p_without)
    assert h_with != h_without