arborist/tests/test_reclassify.py

"""`arborist reclassify` — re-runs the verifier against live providence records.

Covers:
- Stale UNGROUNDED records become STRICT/HYBRID/etc when the new verifier
  finds grounding via span/entity paths.
- Already-correct records stay unchanged.
- Cold-source records (where any source doc has no hot chunks) are
  skipped — we cannot reclassify without context.
- --dry-run reports transitions without writing.
- Each changed record writes one 'providence_reclassify' audit event.
"""

from __future__ import annotations

import argparse
import json
from typing import Iterator

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


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)


def _build_corpus_and_query(tmp_path, answer_text):
    """Ingest one doc, run a query with `answer_text` from a stub client,
    return (qa_db, main_db, cache_key, conn-factory)."""
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"

    conn = connect(main_db)
    try:
        ingest_source(
            conn,
            _FakeSource([
                _doc(
                    "test://matrix",
                    "Keanu Reeves stars as Thomas A. Anderson. "
                    "Laurence Fishburne plays Morpheus. "
                    "Carrie-Anne Moss plays Trinity. "
                    "Hugo Weaving plays Agent Smith.",
                )
            ]),
        )
    finally:
        conn.close()

    client = StubClient(answer=answer_text)
    result = query(
        question="who stars in The Matrix?",
        qa_db=qa_db,
        chat_client=client,
        model_id="test-model",
        single_db=main_db,
        top_k=1,
    )
    return qa_db, main_db, result["cache_key"]


def _make_args(qa_db, dry_run=False, limit=0, entity_policy=None, compare=False):
    return argparse.Namespace(
        qa_db=str(qa_db),
        global_shards_dir=None,
        dry_run=dry_run,
        limit=limit,
        entity_policy=entity_policy,
        compare=compare,
    )


def test_reclassify_promotes_stale_visual_via_entity_path(tmp_path, capsys):
    """Stub answer with no quotes but verbatim entity names. The query
    persists UNGROUNDED/none if the verifier ran first; we then mutate the
    row to simulate an OLDER record (pre-layered-verifier), and reclassify."""
    answer = (
        "The cast includes Keanu Reeves, Laurence Fishburne, "
        "Carrie-Anne Moss, and Hugo Weaving."
    )
    qa_db, _, ckey = _build_corpus_and_query(tmp_path, answer)

    # Force the row into a stale state: UNGROUNDED/none/0 — what the old
    # quote-only verifier would have written.
    conn = connect(qa_db)
    try:
        conn.execute(
            "UPDATE providence_cache SET audit_mode='UNGROUNDED', "
            " verifier_method='none', n_quotes=0, n_verified=0, "
            " unverified_quotes=NULL "
            "WHERE cache_key = ?",
            (ckey,),
        )
    finally:
        conn.close()

    rc = _cmd_reclassify(_make_args(qa_db))
    assert rc == 0
    out = json.loads(capsys.readouterr().out)
    assert out["examined"] == 1
    assert out["changed"] == 1
    assert out["skipped_cold"] == 0

    # Row reflects the new verdict.
    conn = connect(qa_db)
    try:
        row = conn.execute(
            "SELECT audit_mode, verifier_method, n_verified FROM providence_cache "
            "WHERE cache_key = ?",
            (ckey,),
        ).fetchone()
        assert row["audit_mode"] in ("STRICT", "HYBRID")
        assert row["verifier_method"] == "entity"
        assert row["n_verified"] >= 3
        # Audit event recorded.
        last = conn.execute(
            "SELECT event_type FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()
        assert last["event_type"] == "providence_reclassify"
    finally:
        conn.close()


def test_reclassify_unchanged_when_already_correct(tmp_path, capsys):
    """Re-running reclassify on a correctly-labeled record is a no-op."""
    answer = (
        "The cast includes Keanu Reeves, Laurence Fishburne, "
        "Carrie-Anne Moss, and Hugo Weaving."
    )
    qa_db, _, _ = _build_corpus_and_query(tmp_path, answer)

    # First pass classifies correctly via the entity path. Reclassify
    # should report unchanged.
    rc = _cmd_reclassify(_make_args(qa_db))
    assert rc == 0
    out = json.loads(capsys.readouterr().out)
    assert out["examined"] == 1
    assert out["unchanged"] == 1
    assert out["changed"] == 0


def test_reclassify_dry_run_does_not_write(tmp_path, capsys):
    """--dry-run reports transitions but leaves the row untouched."""
    answer = (
        "The cast includes Keanu Reeves, Laurence Fishburne, "
        "Carrie-Anne Moss, and Hugo Weaving."
    )
    qa_db, _, ckey = _build_corpus_and_query(tmp_path, answer)

    # Force stale state.
    conn = connect(qa_db)
    try:
        conn.execute(
            "UPDATE providence_cache SET audit_mode='UNGROUNDED', "
            " verifier_method='none', n_quotes=0, n_verified=0 "
            "WHERE cache_key = ?",
            (ckey,),
        )
    finally:
        conn.close()

    rc = _cmd_reclassify(_make_args(qa_db, dry_run=True))
    assert rc == 0
    out = json.loads(capsys.readouterr().out)
    assert out["dry_run"] is True
    assert out["changed"] == 1

    # Row was NOT updated.
    conn = connect(qa_db)
    try:
        row = conn.execute(
            "SELECT audit_mode, verifier_method FROM providence_cache "
            "WHERE cache_key = ?",
            (ckey,),
        ).fetchone()
        assert row["audit_mode"] == "UNGROUNDED"
        assert row["verifier_method"] == "none"
    finally:
        conn.close()


def test_reclassify_skips_cold_source(tmp_path, capsys):
    """If the source doc's chunks are cold (content NULL), the loader
    returns None and we cannot reclassify — record is skipped."""
    answer = "Keanu Reeves stars in this film."
    qa_db, main_db, ckey = _build_corpus_and_query(tmp_path, answer)

    # Evict all chunks of the source doc to cold. Loader will return None.
    conn = connect(main_db)
    try:
        conn.execute("UPDATE chunks SET content = NULL, tier = 'cold'")
    finally:
        conn.close()

    rc = _cmd_reclassify(_make_args(qa_db))
    assert rc == 0
    out = json.loads(capsys.readouterr().out)
    assert out["examined"] == 1
    assert out["skipped_cold"] == 1
    assert out["changed"] == 0