arborist/tests/test_repair.py

"""Mechanical answer-repair loop.

`mechanical_repair` applies sidecar repair suggestions (synthetic_elision
split, trailing_artifact trim, no_overlap remove) to an answer text
deterministically. The query/ask runners gate this behind
`policy["repair_enabled"]` and re-verify the repaired text; if the
post-repair verdict isn't worse, the repaired answer is persisted with
a `providence_repair` audit event recording the pre→post transition.

These tests cover:
- Each repair action produces the right substitution.
- Idempotence: running repair on already-clean text is a no-op.
- query() integration: repair_enabled=True can promote HYBRID/quote →
  STRICT/quote on synthetic_elision cases without an extra LLM call.
- query() default repair_enabled=False leaves answer text untouched.
"""

from __future__ import annotations

from typing import Iterator

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa import query
from arborist.qa.client import StubClient
from arborist.qa.repair import mechanical_repair
from arborist.source import Source
from arborist.store import connect


# ---------------------------------------------------------------- mechanical_repair


def test_mechanical_repair_splits_synthetic_elision():
    """`"prefix [...] suffix"` (both halves verbatim) becomes
    `"prefix" ... "suffix"`."""
    context = (
        "The film centers on the fictional Isla Nublar, in Costa Rica. "
        "Universal Studios acquired the rights to the novel."
    )
    answer = (
        'The plot states "The film centers on the fictional Isla Nublar [...] '
        'Universal Studios acquired the rights to the novel".'
    )
    bad_quote = (
        "The film centers on the fictional Isla Nublar [...] "
        "Universal Studios acquired the rights to the novel"
    )
    out = mechanical_repair(answer, [bad_quote], context)
    assert len(out["changes"]) == 1
    assert out["changes"][0]["action"] == "split_into_two_quotes"
    # Two separate quoted spans now appear:
    assert '"The film centers on the fictional Isla Nublar"' in out["repaired_text"]
    assert '"Universal Studios acquired the rights to the novel"' in out["repaired_text"]
    # `[...]` no longer appears inside any single quoted span.
    assert "[...]" not in out["repaired_text"]


def test_mechanical_repair_trims_trailing_citation():
    """`"prose. (Source: ...)"` becomes `"prose."`."""
    context = (
        "Pikachu can store electricity in its cheeks and release it in "
        "lightning-based attacks. Pikachu evolves from Pichu."
    )
    bad_quote = (
        "Pikachu can store electricity in its cheeks and release it in "
        "lightning-based attacks. (Source: https://en.wikipedia.org/wiki/Pikachu)"
    )
    answer = f'According to source: "{bad_quote}"'
    out = mechanical_repair(answer, [bad_quote], context)
    assert len(out["changes"]) == 1
    assert out["changes"][0]["action"] == "trim_trailing_artifact"
    assert "(Source:" not in out["repaired_text"]


def test_mechanical_repair_removes_no_overlap_line():
    """Full-invention spans get the line stripped from the answer."""
    context = "Pikachu is a Pokémon species."
    bad_quote = "The Roman Senate convened in 49 BC to debate Caesar's rebellion"
    answer = (
        "- Pikachu lives in the wild\n"
        f'- "{bad_quote}"\n'
        "- Pichu evolves into Pikachu\n"
    )
    out = mechanical_repair(answer, [bad_quote], context)
    assert len(out["changes"]) == 1
    assert out["changes"][0]["action"] == "remove_claim"
    assert bad_quote not in out["repaired_text"]
    # Other bullets preserved.
    assert "Pikachu lives in the wild" in out["repaired_text"]
    assert "Pichu evolves into Pikachu" in out["repaired_text"]


def test_mechanical_repair_idempotent_on_clean_text():
    """No unverified quotes → no change."""
    context = "Cloud is the protagonist."
    answer = 'The source: "Cloud is the protagonist".'
    out = mechanical_repair(answer, [], context)
    assert out["changes"] == []
    assert out["repaired_text"] == answer


# ---------------------------------------------------------------- query() integration


class FakeSource(Source):
    source_type = "test"

    def __init__(self, docs):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri, content):
    return Document(uri=uri, content=content, source_type="test", title=uri.rsplit("/", 1)[-1])


def test_query_repair_disabled_default_leaves_answer_unchanged(tmp_path):
    """`policy["repair_enabled"]` defaults to False — answer text in the
    persisted record matches what the LLM produced."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    long_text = (
        "Capitalism is an economic system based on private ownership "
        "of the means of production. " * 20
    )
    docs = [_doc("test://capitalism", long_text)]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()
    # Answer with a synthetic_elision-style bad quote.
    bad_answer = (
        '"Capitalism is an economic system based on private ownership '
        '[...] of the means of production."'
    )
    result = query(
        question="What is capitalism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=bad_answer),
        model_id="m",
        single_db=main_db,
    )
    # Repair off → answer text unchanged from LLM output.
    qa_conn = connect(qa_db)
    try:
        row = qa_conn.execute(
            "SELECT answer_text FROM providence_cache WHERE cache_key=?",
            (result["cache_key"],),
        ).fetchone()
    finally:
        qa_conn.close()
    assert row["answer_text"] == bad_answer
    assert result["repair_changes"] == []


def test_query_reprompt_rewrites_on_paraphrase_failure(tmp_path):
    """Re-prompt tier handles cases mechanical declines (paraphrase,
    interior_elision needing semantic judgment). The stub returns a
    failing answer on first call & a clean verbatim quote on the
    second; with `repair_max_reprompts=1`, the system promotes the
    verdict to STRICT and the persisted answer is the re-written one."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    src_text = (
        "Capitalism is an economic system based on private ownership "
        "of the means of production. " * 10
    )
    docs = [_doc("test://capitalism", src_text)]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()

    # First call: paraphrase inside a quote that won't substring-match
    # AND won't trigger mechanical repair (no [...], no Source-tail, not
    # full invention — clearly a paraphrase).
    bad_answer = (
        '"Capitalism is an economic philosophy based on personal control "'
        '"over production"'
    )
    # Second call (re-prompt): clean verbatim quote.
    good_answer = (
        '"Capitalism is an economic system based on private ownership '
        'of the means of production"'
    )

    class _SeqClient:
        def __init__(self):
            self.calls = []
            self.answers = [bad_answer, good_answer]

        def chat_completion(self, messages, **kw):
            self.calls.append(messages)
            return self.answers[len(self.calls) - 1] if self.calls else self.answers[0]

    client = _SeqClient()

    from arborist.qa.query import DEFAULT_QUERY_POLICY
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["repair_enabled"] = True
    policy["repair_max_reprompts"] = 1

    result = query(
        question="What is capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
        policy=policy,
    )

    assert len(client.calls) == 2  # one initial + one re-prompt
    assert result["audit_mode"] == "STRICT"
    assert any(
        c["action"] == "reprompt_rewrite" for c in result["repair_changes"]
    )


def test_query_repair_enabled_promotes_synthetic_elision_to_strict(tmp_path):
    """With `repair_enabled=True`, the mechanical loop splits a
    `[...]`-elided quote into two verbatim spans, re-verifies, & lands
    STRICT instead of HYBRID. Persisted answer is the repaired text."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    # Source phrase the model will fail to quote verbatim:
    src_text = (
        "Capitalism is an economic system based on private ownership "
        "of the means of production. " * 20
    )
    docs = [_doc("test://capitalism", src_text)]
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(docs))
    finally:
        conn.close()

    bad_answer = (
        'Per the source: "Capitalism is an economic system based on private ownership '
        '[...] of the means of production".'
    )

    # Build a policy variant with repair_enabled=True.
    from arborist.qa.query import DEFAULT_QUERY_POLICY
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["repair_enabled"] = True

    result = query(
        question="What is capitalism?",
        qa_db=qa_db,
        chat_client=StubClient(answer=bad_answer),
        model_id="m",
        single_db=main_db,
        policy=policy,
    )

    # Repair fired & promoted the verdict.
    assert result["pre_repair_audit_mode"] in ("HYBRID", "UNGROUNDED")
    assert result["audit_mode"] == "STRICT"
    assert len(result["repair_changes"]) >= 1
    assert result["repair_changes"][0]["action"] == "split_into_two_quotes"

    # Persisted answer is the REPAIRED text (no `[...]` inside any quote).
    qa_conn = connect(qa_db)
    try:
        row = qa_conn.execute(
            "SELECT answer_text FROM providence_cache WHERE cache_key=?",
            (result["cache_key"],),
        ).fetchone()
        # Audit chain has the providence_repair event.
        evt = qa_conn.execute(
            "SELECT event_type, body FROM audit_events "
            "WHERE event_type='providence_repair' ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        qa_conn.close()
    assert "[...]" not in row["answer_text"]
    assert evt is not None
    assert evt["event_type"] == "providence_repair"