arborist/aborist/qa/repair.py

"""Mechanical answer repair — apply sidecar repair suggestions deterministically.

When the verifier flags an answer's quoted span as unverified, the
sidecar (`aborist/qa/inspect.py:_classify_span`) emits a `repair`
field naming a concrete fix: split a `[...]`-elided quote into two
verbatim quotes, trim a trailing `(Source: ...)` artifact, restore a
dropped parenthetical aside, drop a fully-invented claim. This module
applies those fixes by string substitution — no LLM call, no
non-determinism.

Gated by `policy["repair_enabled"]`. Off by default so existing
callers don't see answer text mutate under their feet.

Repair philosophy follows the toy-Hermes design pass (fox 2026-04-30):

- Mechanical first. Cheap, deterministic, idempotent.
- Re-prompt second (TODO; one extra LLM call to ask the model to
  rewrite around the failed claim). Adds latency; quality bonus.
- Repair stage is observational from the cache_key's perspective: the
  cache_key inputs are unchanged, only the persisted answer text differs.
  An audit event records the original→repaired transition.

The verifier-stays-binary discipline holds: this module never adds
soft-signal fields to the verifier output. It produces a CHANGED
answer text + a CHANGE LOG that the audit chain stores.
"""

from __future__ import annotations

from typing import Any

from .inspect import _classify_span, _normalize


def mechanical_repair(
    answer_text: str, unverified_quotes: list[str], context: str
) -> dict:
    """Apply repair suggestions to ``answer_text``.

    For each unverified quote, run the sidecar classifier to get a
    repair suggestion, then apply it by string substitution where the
    suggestion has a concrete replacement (split, trim, drop). Skips
    suggestions that need semantic judgment (paraphrase downgrade, aside
    restoration without a definitive verbatim form).

    Returns::

        {
          "repaired_text": str,    # the modified answer (== input if no repair fired)
          "changes":       [dict], # per-quote repair log
        }

    Each change entry::

        {
          "action": str,   # "split_into_two_quotes" | "trim_trailing_artifact" | ...
          "before": str,   # the original quoted span
          "after":  str | None,  # the replacement text, or None if removed
          "diagnosis": str,
        }

    Idempotent: applying repair twice produces the same result (the
    repaired quotes substring-match in source, so a second pass finds
    nothing more to fix).
    """
    norm_ctx = _normalize(context)
    repaired = answer_text
    changes: list[dict] = []

    for quote in unverified_quotes:
        diag = _classify_span(quote, norm_ctx, norm_ctx)
        repair = diag.get("repair") or {}
        action = repair.get("action")
        diagnosis = diag.get("diagnosis", "?")

        if action == "split_into_two_quotes":
            # `"prefix [...] suffix"` becomes `"prefix" ... "suffix"`. The
            # model's `[...]` ellipsis-marker becomes prose between two
            # verbatim-quoted spans the verifier can independently check.
            quotes = repair.get("quotes") or []
            if len(quotes) == 2 and all(quotes):
                new_segment = f'"{quotes[0]}" ... "{quotes[1]}"'
                old_segment = f'"{quote}"'
                if old_segment in repaired:
                    repaired = repaired.replace(old_segment, new_segment, 1)
                    changes.append({
                        "action": action,
                        "diagnosis": diagnosis,
                        "before": quote,
                        "after": new_segment,
                    })

        elif action == "trim_trailing_artifact":
            # `"prose. (Source: https://...)"` becomes `"prose."`. The
            # verbatim prefix kept; the model-appended tail dropped.
            kept = repair.get("kept_prefix")
            if kept:
                old_segment = f'"{quote}"'
                new_segment = f'"{kept}"'
                if old_segment in repaired:
                    repaired = repaired.replace(old_segment, new_segment, 1)
                    changes.append({
                        "action": action,
                        "diagnosis": diagnosis,
                        "before": quote,
                        "after": kept,
                    })

        elif action == "remove_claim":
            # Drop the line containing this quote entirely. Catches
            # full-invention spans (no_overlap diagnosis) where there's
            # no verbatim source content to substitute.
            new_lines = []
            removed = False
            for line in repaired.splitlines(keepends=True):
                if quote in line and not removed:
                    removed = True
                    continue
                new_lines.append(line)
            if removed:
                repaired = "".join(new_lines)
                changes.append({
                    "action": action,
                    "diagnosis": diagnosis,
                    "before": quote,
                    "after": None,
                })

        # Skipped: include_aside_for_verbatim (needs precise source-span
        # extraction we'd rather defer to a re-prompt), paraphrase
        # downgrade (needs prose rewriting), partial_paraphrase
        # (split_or_remove judgment). These cases pass through the
        # mechanical loop unchanged & remain UNSUPPORTED in the
        # post-repair verdict — falling back to a re-prompt loop is
        # future work (TODO: re-prompt feedback path).

    return {"repaired_text": repaired, "changes": changes}


def reprompt_repair(
    *,
    chat_client: Any,
    model_id: str,
    original_messages: list[dict],
    original_answer: str,
    failed_quotes: list[str],
    policy: dict,
) -> str | None:
    """Second tier of the hybrid repair loop: add a feedback user-turn
    naming the failed quotes & call the LLM once more.

    Mechanical repair handles deterministic substitutions (split, trim,
    drop). For the cases it declines (paraphrase, partial_paraphrase,
    interior_elision needing semantic judgment), this function asks the
    model to rewrite around the failed claims using only verbatim
    citations. Caller decides whether to accept the new answer based on
    a re-verify.

    Returns the new answer text or ``None`` if the LLM call fails or
    there's nothing to feed back. Caller is responsible for verifying
    & deciding whether to keep the new answer.

    Hard rule: only fires when ``policy["repair_max_reprompts"] > 0``.
    Caller enforces the cap (this function does not loop internally —
    one call per invocation).
    """
    if not failed_quotes:
        return None

    failed_listing = "\n".join(
        f"- \"{q[:200]}{'...' if len(q) > 200 else ''}\""
        for q in failed_quotes[:5]
    )
    feedback = (
        "Your previous answer contained quoted spans that did not "
        "verify verbatim against the provided sources:\n\n"
        f"{failed_listing}\n\n"
        "Rewrite the answer using ONLY verbatim quotes that appear "
        "word-for-word in the sources. Do not insert `[...]` ellipsis "
        "markers inside quotes — split into two separate quotes "
        "instead. Do not append `(Source: ...)` citation tails to "
        "quotes — keep the citation as prose outside the quoted "
        "span. If a claim cannot be verbatim-cited from a source, "
        "omit the claim or downgrade it to a non-quoted paraphrase."
    )

    new_messages = list(original_messages) + [
        {"role": "assistant", "content": original_answer},
        {"role": "user", "content": feedback},
    ]

    try:
        return chat_client.chat_completion(
            new_messages,
            model=model_id,
            temperature=policy.get("temperature", 0.1),
            max_tokens=policy.get("max_tokens", 768),
            top_p=policy.get("top_p", 1.0),
        )
    except Exception:
        return None