Merge pull request #11 from dk3yyyy/fix/rag-answer-diagnostics-repair

Improve RAG answer reliability and diagnostics

Author: dk3yyyy

README.md

@@ -140,7 +140,30 @@ uv run local-ai-agent evaluate --report-dir evaluation/results/my-run
 The versioned case manifest is tied to the dataset SHA-256 and uses immutable
 content-derived source IDs for gold relevance. The report records model tags
 and immutable Ollama digests, dataset and case-set hashes, retrieval limit,
-runtime versions, per-case RAG and BM25 rankings, and aggregate metrics.
+runtime versions, per-case RAG and BM25 rankings, and aggregate metrics. Each
+RAG observation also retains the first raw model response, any one-shot repair
+response, the initial and final structured validation reasons, and whether a
+repair was attempted. These diagnostics stay in the evaluation artifact; the
+CLI and dashboard continue to expose only validated answers or safe fallback
+messages.
+
+Compare answer models without changing the embedding model or case set:
+
+```bash
+uv run local-ai-agent evaluate \
+  --chat-model llama3.2 \
+  --report-dir /tmp/local-ai-agent-llama3.2
+
+uv run local-ai-agent evaluate \
+  --chat-model <installed-7b-or-8b-instruct-model> \
+  --report-dir /tmp/local-ai-agent-stronger-model
+```
+
+Keep the reports outside the repository so benchmark artifacts do not make the
+worktree dirty. Compare answer success, citation validity, abstention recall,
+latency, and the per-case failure reasons rather than optimizing one aggregate
+score.
+
 Retrieval quality is reported as recall@k, hit rate@k, and MRR@k for both
 semantic search and the BM25 baseline. Relevance judgments are known-positive,
 not exhaustive. The generated report is evidence for this fixed benchmark

agent.py

@@ -17,13 +17,19 @@
 )
 CITATION_PATTERN = re.compile(r"\[([A-Za-z0-9][A-Za-z0-9_-]*)\]")
 INSUFFICIENT_EVIDENCE_TOKEN = "INSUFFICIENT_EVIDENCE"
+REPAIRABLE_FAILURE_REASONS = frozenset(
+    {"missing_citations", "out_of_range_citation", "unknown_citation"}
+)
 
 ANSWER_PROMPT = """You are a review analyst.
 Answer the question using only the supplied reviews. Do not add facts that are not present.
-When evidence is mixed or limited, say so clearly. Every factual claim must cite one or more
-retrieved evidence numbers exactly as shown, for example [1]. Never cite an evidence number
-that is not supplied. Source IDs are validation metadata; do not copy them into the answer.
-If the supplied reviews do not answer the question, reply exactly INSUFFICIENT_EVIDENCE.
+Always answer every part that at least one supplied review supports. Mixed or incomplete evidence is
+not a reason to abstain; describe the limitation and answer only the supported portion. Every
+factual claim must cite one or more retrieved evidence numbers exactly as shown, for example
+[1]. Never cite an evidence number that is not supplied. Source IDs are validation metadata;
+do not copy them into the answer. Reply exactly INSUFFICIENT_EVIDENCE only when no supplied
+review answers any part of the question. Never use INSUFFICIENT_EVIDENCE as prose or place it
+beside an answer.
 
 Question:
 {question}
@@ -34,6 +40,27 @@
 Answer:
 """
 
+REPAIR_PROMPT = """You are repairing one rejected review-analysis response.
+Using only the supplied reviews, rewrite it once so every factual claim has one or more valid
+evidence citations such as [1]. Use only evidence numbers that appear below. Preserve supported
+meaning, remove unsupported claims, and answer every supported part even when evidence is mixed
+or incomplete. If no supplied review answers any part, reply exactly INSUFFICIENT_EVIDENCE.
+Never use INSUFFICIENT_EVIDENCE as prose or place it beside an answer.
+
+Question:
+{question}
+
+Supplied review records:
+{context}
+
+Rejected response:
+<rejected_response>
+{rejected_response}
+</rejected_response>
+
+Rewritten answer:
+"""
+
 
 @dataclass(frozen=True)
 class CitedReview:
@@ -48,6 +75,11 @@ class AnswerResult:
     sources: tuple[CitedReview, ...]
     retrieved_source_ids: tuple[str, ...] = ()
     abstained: bool = False
+    raw_response: str = ""
+    repair_response: str | None = None
+    initial_failure_reason: str | None = None
+    failure_reason: str | None = None
+    repair_attempted: bool = False
 
 
 def create_chat_model(
@@ -99,27 +131,28 @@ def _remove_standalone_control_token(answer: str) -> str | None:
 def _validate_and_number_citations(
     answer: str,
     matches: list[ReviewMatch],
-) -> tuple[str, tuple[CitedReview, ...]] | None:
+) -> tuple[tuple[str, tuple[CitedReview, ...]] | None, str | None]:
     retrieved: dict[str, ReviewMatch] = {}
     evidence_aliases: dict[str, str] = {}
     for evidence_number, match in enumerate(matches, start=1):
         source_id = str(
             match.document.metadata.get("source_id") or match.document.id or ""
         )
         if not source_id:
-            return None
+            return None, "retrieved_source_missing_id"
         retrieved[source_id] = match
         evidence_aliases[str(evidence_number)] = source_id
 
     cited_tokens = CITATION_PATTERN.findall(answer)
     if not cited_tokens:
-        return None
+        return None, "missing_citations"
 
     resolved_ids: list[str] = []
     for token in cited_tokens:
         source_id = evidence_aliases.get(token, token)
         if source_id not in retrieved:
-            return None
+            reason = "out_of_range_citation" if token.isdigit() else "unknown_citation"
+            return None, reason
         resolved_ids.append(source_id)
 
     ordered_ids = list(dict.fromkeys(resolved_ids))
@@ -140,7 +173,26 @@ def _validate_and_number_citations(
         )
         for source_id in ordered_ids
     )
-    return numbered_answer, sources
+    return (numbered_answer, sources), None
+
+
+def _evaluate_model_response(
+    response: str,
+    matches: list[ReviewMatch],
+) -> tuple[tuple[str, tuple[CitedReview, ...]] | None, str | None]:
+    normalized = response.strip()
+    if normalized == INSUFFICIENT_EVIDENCE_TOKEN:
+        return None, "clean_abstention"
+    normalized_without_control = _remove_standalone_control_token(normalized)
+    if normalized_without_control is None:
+        return None, "embedded_control_token"
+    if not normalized_without_control:
+        return None, "invalid_remainder"
+    return _validate_and_number_citations(normalized_without_control, matches)
+
+
+def _response_text(response: Any) -> str:
+    return str(response.content if hasattr(response, "content") else response)
 
 
 def answer_question(
@@ -177,47 +229,95 @@ def answer_question(
         countries=countries,
     )
     if not matches:
-        return AnswerResult(answer=NO_MATCH_MESSAGE, sources=())
+        return AnswerResult(
+            answer=NO_MATCH_MESSAGE,
+            sources=(),
+            failure_reason="empty_retrieval",
+        )
 
     retrieved_source_ids = tuple(
         str(match.document.metadata.get("source_id") or match.document.id or "")
         for match in matches
     )
     if any(not source_id for source_id in retrieved_source_ids):
-        return AnswerResult(answer=CITATION_VALIDATION_MESSAGE, sources=())
+        return AnswerResult(
+            answer=CITATION_VALIDATION_MESSAGE,
+            sources=(),
+            failure_reason="retrieved_source_missing_id",
+        )
 
     answer_model = model or create_chat_model(model=chat_model, base_url=ollama_host)
+    context = _format_context(matches)
     prompt = ANSWER_PROMPT.format(
         question=normalized_question,
-        context=_format_context(matches),
+        context=context,
     )
-    response = answer_model.invoke(prompt)
-    answer = response.content if hasattr(response, "content") else str(response)
-    normalized_answer = answer.strip()
-    if normalized_answer == INSUFFICIENT_EVIDENCE_TOKEN:
+    raw_response = _response_text(answer_model.invoke(prompt))
+    validated, failure_reason = _evaluate_model_response(raw_response, matches)
+    if validated is not None:
+        validated_answer, sources = validated
+        return AnswerResult(
+            answer=validated_answer,
+            sources=sources,
+            retrieved_source_ids=retrieved_source_ids,
+            raw_response=raw_response,
+        )
+    if failure_reason == "clean_abstention":
         return AnswerResult(
             answer=NO_MATCH_MESSAGE,
             sources=(),
             retrieved_source_ids=retrieved_source_ids,
             abstained=True,
+            raw_response=raw_response,
+            failure_reason=failure_reason,
         )
-    normalized_answer = _remove_standalone_control_token(normalized_answer)
-    if normalized_answer is None:
+    if failure_reason not in REPAIRABLE_FAILURE_REASONS:
         return AnswerResult(
             answer=CITATION_VALIDATION_MESSAGE,
             sources=(),
             retrieved_source_ids=retrieved_source_ids,
+            raw_response=raw_response,
+            failure_reason=failure_reason,
         )
-    validated = _validate_and_number_citations(normalized_answer, matches)
-    if validated is None:
+
+    initial_failure_reason = failure_reason
+    repair_prompt = REPAIR_PROMPT.format(
+        question=normalized_question,
+        context=context,
+        rejected_response=raw_response,
+    )
+    repair_response = _response_text(answer_model.invoke(repair_prompt))
+    repaired, repair_failure_reason = _evaluate_model_response(repair_response, matches)
+    if repaired is not None:
+        repaired_answer, sources = repaired
         return AnswerResult(
-            answer=CITATION_VALIDATION_MESSAGE,
+            answer=repaired_answer,
+            sources=sources,
+            retrieved_source_ids=retrieved_source_ids,
+            raw_response=raw_response,
+            repair_response=repair_response,
+            initial_failure_reason=initial_failure_reason,
+            repair_attempted=True,
+        )
+    if repair_failure_reason == "clean_abstention":
+        return AnswerResult(
+            answer=NO_MATCH_MESSAGE,
             sources=(),
             retrieved_source_ids=retrieved_source_ids,
+            abstained=True,
+            raw_response=raw_response,
+            repair_response=repair_response,
+            initial_failure_reason=initial_failure_reason,
+            failure_reason=repair_failure_reason,
+            repair_attempted=True,
         )
-    validated_answer, sources = validated
     return AnswerResult(
-        answer=validated_answer,
-        sources=sources,
+        answer=CITATION_VALIDATION_MESSAGE,
+        sources=(),
         retrieved_source_ids=retrieved_source_ids,
+        raw_response=raw_response,
+        repair_response=repair_response,
+        initial_failure_reason=initial_failure_reason,
+        failure_reason=repair_failure_reason,
+        repair_attempted=True,
     )