diff --git a/aborist/qa/dag.py b/aborist/qa/dag.py index ba4c20c..0c8ced2 100644 --- a/aborist/qa/dag.py +++ b/aborist/qa/dag.py @@ -52,6 +52,47 @@ def _canonical_json(obj) -> str: return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False) +def localize_failure( + *, + audit_mode: str, + n_sources: int, + n_quotes: int, + n_verified: int, +) -> str | None: + """Map a non-STRICT verdict to the pipeline stage that introduced + the failure. Returns ``None`` for STRICT outcomes. + + Stage labels (in pipeline order): + + - ``retrieval`` — no admitted sources. Title/body gates rejected + everything, or the corpus genuinely lacks the topic. Repair path: + ingest more sources or relax the breadth threshold. + - ``context`` — sources admitted but no quotes extracted. Could be + a context-truncation issue (per-source cap dropped the relevant + paragraph) or a model that declined to cite anything. Repair path: + raise per-source cap; tighten prompt. + - ``answer`` — sources retrieved & quotes extracted but they don't + verify. The model either fabricated content, paraphrased inside + quotes, or appended citation tails. Repair path: the + ``mechanical_repair`` pass + (when wired) the re-prompt feedback + loop. + + The toy-Hermes design pass calls this "chain-segment failure + localization" — debugging becomes typed instead of vague. An + operator reading ``failure_stage='answer'`` knows retrieval & + context were fine; the model is what to fix. ``failure_stage='retrieval'`` + means stop tuning the verifier & go ingest a relevant source. + """ + if audit_mode == "STRICT": + return None + if n_sources == 0: + return "retrieval" + if n_quotes == 0: + return "context" + # Quotes were extracted but didn't all verify (or none did). + return "answer" + + def build_run_dag( *, question_hash: str, @@ -84,12 +125,19 @@ def build_run_dag( ] retrieval_hash = _sha256_hex(_canonical_json(sources_summary)) answer_hash = _sha256_hex(answer_text) + failure_stage = localize_failure( + audit_mode=audit_mode, + n_sources=len(sources), + n_quotes=n_quotes, + n_verified=n_verified, + ) verify_payload = { "audit_mode": audit_mode, "verifier_method": verifier_method, "n_quotes": n_quotes, "n_verified": n_verified, "claim_statuses": claim_statuses or [], + "failure_stage": failure_stage, } verify_hash = _sha256_hex(_canonical_json(verify_payload)) final_label_payload = { diff --git a/aborist/qa/query.py b/aborist/qa/query.py index f709d0c..33eef5f 100644 --- a/aborist/qa/query.py +++ b/aborist/qa/query.py @@ -65,7 +65,7 @@ from aborist.qa.keys import ( question_hash, ) from aborist.qa.dag import build_run_dag -from aborist.qa.repair import mechanical_repair +from aborist.qa.repair import mechanical_repair, reprompt_repair from aborist.qa.verify import verify_quotes try: @@ -270,6 +270,10 @@ DEFAULT_QUERY_POLICY = { # `providence_repair` records the pre→post transition. Bumps # governance_policy_hash so on/off agents share no cache silos. "repair_enabled": False, + # Maximum re-prompt iterations after mechanical repair. 0 = no + # re-prompt (mechanical only). 1 = at most one extra LLM call + # asking the model to rewrite around the failed quotes. + "repair_max_reprompts": 0, } @@ -1066,37 +1070,72 @@ def query( proximity_window=policy.get("entity_proximity_window", 300), ) - # 5c. Optional mechanical repair pass. When `policy["repair_enabled"]` - # is on AND the first verdict isn't STRICT, walk the unverified - # quotes through the sidecar classifier; for each repairable - # diagnosis (synthetic_elision split, trailing_artifact trim, - # no_overlap remove) apply the suggestion deterministically & - # re-verify. Persist the post-repair answer + the change log; - # cache_key inputs are unchanged so the cache_key itself stays. + # 5c. Optional hybrid repair loop. Mechanical first (deterministic + # string substitutions for synthetic_elision / trailing_artifact / + # no_overlap), then up to `repair_max_reprompts` LLM re-prompts + # for the cases mechanical declined (paraphrase, interior_elision + # needing semantic judgment). Both gated by `policy["repair_enabled"]`. + # Persist the post-repair answer + the change log; cache_key + # inputs are unchanged. repair_changes: list[dict] = [] pre_repair_verdict: dict | None = None + + def _verify(text: str) -> dict: + return verify_quotes( + text, + context, + entity_policy=policy.get("entity_policy", "hybrid"), + proximity_n=policy.get("entity_proximity_n", 3), + proximity_window=policy.get("entity_proximity_window", 300), + ) + if ( policy.get("repair_enabled") and verdict["audit_mode"] != "STRICT" and verdict.get("unverified_quotes") ): + # Tier 1: mechanical (deterministic, no extra LLM call). repair_result = mechanical_repair( answer_text, verdict["unverified_quotes"], context ) if repair_result["changes"]: - new_verdict = verify_quotes( - repair_result["repaired_text"], - context, - entity_policy=policy.get("entity_policy", "hybrid"), - proximity_n=policy.get("entity_proximity_n", 3), - proximity_window=policy.get("entity_proximity_window", 300), - ) - # Accept the repair only if it didn't make things worse. + new_verdict = _verify(repair_result["repaired_text"]) if new_verdict["n_verified"] >= verdict["n_verified"]: pre_repair_verdict = verdict answer_text = repair_result["repaired_text"] verdict = new_verdict - repair_changes = repair_result["changes"] + repair_changes = list(repair_result["changes"]) + + # Tier 2: re-prompt feedback (one extra LLM call max). + max_reprompts = int(policy.get("repair_max_reprompts", 0)) + for _ in range(max_reprompts): + if ( + verdict["audit_mode"] == "STRICT" + or not verdict.get("unverified_quotes") + ): + break + new_text = reprompt_repair( + chat_client=chat_client, + model_id=model_id, + original_messages=messages, + original_answer=answer_text, + failed_quotes=verdict["unverified_quotes"], + policy=policy, + ) + if not new_text: + break + new_verdict = _verify(new_text) + if new_verdict["n_verified"] > verdict["n_verified"]: + if pre_repair_verdict is None: + pre_repair_verdict = verdict + answer_text = new_text + verdict = new_verdict + repair_changes.append({ + "action": "reprompt_rewrite", + "diagnosis": "model_feedback_loop", + }) + else: + break unverified_blob = ( json.dumps(verdict["unverified_quotes"], separators=(",", ":")) @@ -1225,12 +1264,27 @@ def query( qa_conn.close() persist_ms = _ms_since(t_phase) + # Pull the verify node's failure_stage out of the run_dag for the result. + failure_stage = next( + (n.get("hash") for n in run_dag["nodes"] if n["stage"] == "verify"), + None, + ) + # The actual label is in the verify_payload, which we computed in + # localize_failure earlier — recompute for the result dict. + from aborist.qa.dag import localize_failure as _localize + failure_stage = _localize( + audit_mode=verdict["audit_mode"], + n_sources=len(chosen), + n_quotes=verdict["n_quotes"], + n_verified=verdict["n_verified"], + ) return { "status": "cache_miss_then_written", "audit_mode": verdict["audit_mode"], "cache_key": ckey, "run_dag_root": run_dag["root"], "lookup_path": "miss", + "failure_stage": failure_stage, "repair_changes": repair_changes, "pre_repair_audit_mode": ( pre_repair_verdict["audit_mode"] if pre_repair_verdict else None diff --git a/aborist/qa/repair.py b/aborist/qa/repair.py index 425efd4..8f978fb 100644 --- a/aborist/qa/repair.py +++ b/aborist/qa/repair.py @@ -27,6 +27,8 @@ answer text + a CHANGE LOG that the audit chain stores. from __future__ import annotations +from typing import Any + from .inspect import _classify_span, _normalize @@ -133,3 +135,67 @@ def mechanical_repair( # 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 diff --git a/aborist/qa/runner.py b/aborist/qa/runner.py index b0778e6..adac39c 100644 --- a/aborist/qa/runner.py +++ b/aborist/qa/runner.py @@ -35,7 +35,7 @@ from aborist.qa.keys import ( question_hash, ) from aborist.qa.dag import build_run_dag -from aborist.qa.repair import mechanical_repair +from aborist.qa.repair import mechanical_repair, reprompt_repair from aborist.qa.verify import verify_quotes from aborist.store import append_audit, transaction @@ -76,6 +76,7 @@ DEFAULT_POLICY = { # Mechanical answer repair after first verify. Off by default; see # aborist/qa/query.py for semantics. "repair_enabled": False, + "repair_max_reprompts": 0, # Strip wikitext markup before the LLM ever sees the context. Lets # Hermes quote prose verbatim and shrinks token bills (~43% on # Wikipedia chunks). Bumps governance_policy_hash so prior cached @@ -294,9 +295,19 @@ def ask( proximity_window=policy.get("entity_proximity_window", 300), ) - # Optional mechanical repair pass — see aborist/qa/query.py for shape. + # Optional hybrid repair loop — see aborist/qa/query.py for shape. repair_changes: list[dict] = [] pre_repair_verdict: dict | None = None + + def _verify(text: str) -> dict: + return verify_quotes( + text, + document_text, + entity_policy=policy.get("entity_policy", "hybrid"), + proximity_n=policy.get("entity_proximity_n", 3), + proximity_window=policy.get("entity_proximity_window", 300), + ) + if ( policy.get("repair_enabled") and verdict["audit_mode"] != "STRICT" @@ -306,18 +317,42 @@ def ask( answer_text, verdict["unverified_quotes"], document_text ) if repair_result["changes"]: - new_verdict = verify_quotes( - repair_result["repaired_text"], - document_text, - entity_policy=policy.get("entity_policy", "hybrid"), - proximity_n=policy.get("entity_proximity_n", 3), - proximity_window=policy.get("entity_proximity_window", 300), - ) + new_verdict = _verify(repair_result["repaired_text"]) if new_verdict["n_verified"] >= verdict["n_verified"]: pre_repair_verdict = verdict answer_text = repair_result["repaired_text"] verdict = new_verdict - repair_changes = repair_result["changes"] + repair_changes = list(repair_result["changes"]) + + max_reprompts = int(policy.get("repair_max_reprompts", 0)) + for _ in range(max_reprompts): + if ( + verdict["audit_mode"] == "STRICT" + or not verdict.get("unverified_quotes") + ): + break + new_text = reprompt_repair( + chat_client=client, + model_id=model_id, + original_messages=messages, + original_answer=answer_text, + failed_quotes=verdict["unverified_quotes"], + policy=policy, + ) + if not new_text: + break + new_verdict = _verify(new_text) + if new_verdict["n_verified"] > verdict["n_verified"]: + if pre_repair_verdict is None: + pre_repair_verdict = verdict + answer_text = new_text + verdict = new_verdict + repair_changes.append({ + "action": "reprompt_rewrite", + "diagnosis": "model_feedback_loop", + }) + else: + break unverified_blob = ( json.dumps(verdict["unverified_quotes"], separators=(",", ":")) @@ -429,12 +464,20 @@ def ask( ), ) + from aborist.qa.dag import localize_failure as _localize + failure_stage = _localize( + audit_mode=verdict["audit_mode"], + n_sources=1, # ask() runs against one document + n_quotes=verdict["n_quotes"], + n_verified=verdict["n_verified"], + ) return { "status": "cache_miss_then_written", "audit_mode": verdict["audit_mode"], "cache_key": ckey, "run_dag_root": run_dag["root"], "lookup_path": "miss", + "failure_stage": failure_stage, "repair_changes": repair_changes, "pre_repair_audit_mode": ( pre_repair_verdict["audit_mode"] if pre_repair_verdict else None diff --git a/tests/test_dag.py b/tests/test_dag.py index 76ce4ef..011cab4 100644 --- a/tests/test_dag.py +++ b/tests/test_dag.py @@ -13,7 +13,7 @@ exactly as recorded. from __future__ import annotations -from aborist.qa.dag import build_run_dag, verify_run_dag +from aborist.qa.dag import build_run_dag, localize_failure, verify_run_dag def _kw(**overrides): @@ -109,6 +109,60 @@ def test_verify_dag_detects_node_tampering(): assert verify_run_dag(out) is False +def test_localize_failure_strict_returns_none(): + """STRICT verdict has no failure to localize.""" + assert localize_failure(audit_mode="STRICT", n_sources=3, n_quotes=4, n_verified=4) is None + + +def test_localize_failure_no_sources_is_retrieval(): + """No admitted sources → retrieval-stage failure (gate over-rejected + or corpus genuinely lacks the topic).""" + assert localize_failure( + audit_mode="UNGROUNDED", n_sources=0, n_quotes=0, n_verified=0 + ) == "retrieval" + + +def test_localize_failure_no_quotes_is_context(): + """Sources retrieved but no quotes extracted → context-stage failure + (per-source cap dropped relevant content, or model declined to + cite). Distinct from answer-stage failures where quotes existed but + didn't verify.""" + assert localize_failure( + audit_mode="UNGROUNDED", n_sources=5, n_quotes=0, n_verified=0 + ) == "context" + + +def test_localize_failure_unverified_quotes_is_answer(): + """Quotes extracted but not all verify → answer-stage failure + (model fabricated, paraphrased inside quotes, or appended + citations). The case mechanical_repair targets.""" + assert localize_failure( + audit_mode="HYBRID", n_sources=3, n_quotes=4, n_verified=3 + ) == "answer" + assert localize_failure( + audit_mode="UNGROUNDED", n_sources=3, n_quotes=2, n_verified=0 + ) == "answer" + + +def test_localize_failure_lands_on_run_dag_verify_node(): + """failure_stage is folded into the verify node's payload so + auditors reading the persisted DAG can see which stage produced + a non-STRICT verdict.""" + out = build_run_dag(**_kw( + audit_mode="HYBRID", + n_quotes=4, + n_verified=3, + )) + # The verify_node hash should differ between failure_stage='answer' + # and a STRICT verdict — failure_stage enters the hashed payload. + out_strict = build_run_dag(**_kw( + audit_mode="STRICT", + n_quotes=4, + n_verified=4, + )) + out["root"] != out_strict["root"] # at minimum, audit_mode differs + + def test_verify_dag_accepts_json_string(): """Sidecar/audit tools persist the DAG as JSON; verify accepts both.""" import json diff --git a/tests/test_repair.py b/tests/test_repair.py index 21dd423..9612a39 100644 --- a/tests/test_repair.py +++ b/tests/test_repair.py @@ -157,6 +157,70 @@ def test_query_repair_disabled_default_leaves_answer_unchanged(tmp_path): 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 aborist.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