rmi-backend/app/hallucination_guard.py

712 lines
27 KiB
Python

"""
Hallucination Guard Module — NLI-based hallucination detection for RAG outputs.
Primary method: Cross-encoder NLI model (DeBERTa-v3) scores (context, answer) pairs.
Secondary method: LLM self-check via OpenRouter API when DeBERTa is unavailable.
Used AFTER LLM generation to verify answers are grounded in retrieved context.
"""
from __future__ import annotations
import asyncio
import logging
import os
import re
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any
import httpx
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
NLI_MODEL_PRIMARY = "cross-encoder/nli-deberta-v3-small"
NLI_MODEL_FALLBACK = "moritz23/nli-deberta-v3-xsmall"
OPENROUTER_API_KEY_ENV = "OPENROUTER_API_KEY"
LLM_API_KEY_ENV = "LLM_API_KEY"
# Decode base64 LLM key if present, otherwise use plain LLM_API_KEY
# (safety net: ensures key is decoded even when imported without main.py)
if os.getenv("LLM_API_KEY_B64"):
import base64 as _b64
os.environ["LLM_API_KEY"] = _b64.b64decode(os.getenv("LLM_API_KEY_B64")).decode()
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api.deepseek.com/v1/chat/completions")
LLM_MODEL = os.getenv("LLM_MODEL", "deepseek-v4-flash")
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_MODEL = "openai/gpt-4.1-mini"
MAX_CHECK_TIMEOUT_S = 5.0
# NLI label indices (standard for cross-encoder NLI models)
LABEL_ENTAILMENT = 0
LABEL_NEUTRAL = 1
LABEL_CONTRADICTION = 2
# Confidence thresholds
CONFIDENCE_HIGH = 0.85
CONFIDENCE_MEDIUM = 0.60
class ConfidenceLevel(StrEnum):
HIGH = "high"
MEDIUM = "medium"
LOW = "low"
# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------
@dataclass
class CheckResult:
"""Result of a single hallucination check."""
is_faithful: bool
confidence: float
method: str
contradiction_details: list[str] = field(default_factory=list)
flagged_claims: list[str] = field(default_factory=list)
@property
def confidence_level(self) -> ConfidenceLevel:
if self.confidence > CONFIDENCE_HIGH:
return ConfidenceLevel.HIGH
elif self.confidence >= CONFIDENCE_MEDIUM:
return ConfidenceLevel.MEDIUM
return ConfidenceLevel.LOW
@dataclass
class CitationVerificationResult:
"""Result of verifying inline citations against sources."""
total_citations: int
verified_citations: int
unsupported_citations: list[dict[str, Any]] = field(default_factory=list)
citation_details: list[dict[str, Any]] = field(default_factory=list)
all_supported: bool = True
# ---------------------------------------------------------------------------
# Claim splitter — simple sentence-level extraction
# ---------------------------------------------------------------------------
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")
def _split_claims(text: str) -> list[str]:
"""Split answer into individual claim sentences for granular checking."""
sentences = _SENTENCE_RE.split(text.strip())
return [s.strip() for s in sentences if s.strip()]
# ---------------------------------------------------------------------------
# Citation parser
# ---------------------------------------------------------------------------
_CITATION_RE = re.compile(r"\[(\d+)\]")
def _extract_citations(text: str) -> list[tuple[int, str]]:
"""
Parse inline citations like [1], [2] from the answer.
Returns list of (citation_number, sentence_containing_citation).
"""
results: list[tuple[int, str]] = []
sentences = _split_claims(text)
for sentence in sentences:
nums = _CITATION_RE.findall(sentence)
for n in nums:
results.append((int(n), sentence))
return results
# ---------------------------------------------------------------------------
# HallucinationGuard
# ---------------------------------------------------------------------------
class HallucinationGuard:
"""NLI-based hallucination detector with LLM self-check fallback."""
_instance: HallucinationGuard | None = None
def __init__(
self,
nli_model_name: str = NLI_MODEL_PRIMARY,
fallback_model_name: str = NLI_MODEL_FALLBACK,
timeout_s: float = MAX_CHECK_TIMEOUT_S,
) -> None:
self._nli_model_name = nli_model_name
self._fallback_model_name = fallback_model_name
self._timeout_s = timeout_s
# Lazy-loaded model artefacts
self._tokenizer = None
self._model = None
self._model_loaded = False
self._model_load_failed = False
self._active_model_name: str | None = None
# -----------------------------------------------------------------------
# Singleton
# -----------------------------------------------------------------------
@classmethod
async def get_guard(cls) -> HallucinationGuard:
"""Return the singleton HallucinationGuard, creating if needed."""
if cls._instance is None:
cls._instance = cls()
if not cls._instance._model_loaded and not cls._instance._model_load_failed:
await cls._instance._load_model()
return cls._instance
# -----------------------------------------------------------------------
# Model loading
# -----------------------------------------------------------------------
async def _load_model(self) -> None:
"""Attempt to load the primary NLI model, then the fallback."""
for model_name in (self._nli_model_name, self._fallback_model_name):
try:
logger.info("Loading NLI model: %s", model_name)
from transformers import AutoModelForSequenceClassification, AutoTokenizer
loop = asyncio.get_event_loop()
tokenizer, model = await loop.run_in_executor(
None,
lambda mn=model_name: (
AutoTokenizer.from_pretrained(mn),
AutoModelForSequenceClassification.from_pretrained(mn),
),
)
model.eval()
self._tokenizer = tokenizer
self._model = model
self._model_loaded = True
self._active_model_name = model_name
logger.info("Successfully loaded NLI model: %s", model_name)
return
except Exception as exc:
logger.warning("Failed to load NLI model %s: %s", model_name, exc)
continue
self._model_load_failed = True
logger.warning("All NLI models failed to load. Falling back to LLM self-check.")
# -----------------------------------------------------------------------
# Warm-up & health
# -----------------------------------------------------------------------
async def warm_up(self) -> None:
"""Pre-load the NLI model so first real check is fast."""
if not self._model_loaded and not self._model_load_failed:
await self._load_model()
logger.info(
"HallucinationGuard warm-up complete. model_loaded=%s active=%s",
self._model_loaded,
self._active_model_name,
)
def health_check(self) -> dict[str, Any]:
"""Return current model status for health endpoints."""
return {
"nli_model_loaded": self._model_loaded,
"nli_model_name": self._active_model_name,
"fallback_active": self._model_load_failed,
"primary_model": self._nli_model_name,
"fallback_model": self._fallback_model_name,
"method": "nli" if self._model_loaded else "llm_self_check",
}
# -----------------------------------------------------------------------
# Core NLI scoring
# -----------------------------------------------------------------------
async def _nli_score(self, premise: str, hypothesis: str) -> dict[str, float]:
"""
Run NLI scoring on a (premise, hypothesis) pair.
Returns dict with keys: entailment, neutral, contradiction.
"""
import torch
def _sync_score() -> dict[str, float]:
inputs = self._tokenizer(
premise,
hypothesis,
truncation=True,
max_length=512,
return_tensors="pt",
)
with torch.no_grad():
logits = self._model(**inputs).logits
probs = torch.softmax(logits, dim=-1).squeeze().tolist()
# Ensure probs is a list even for single-entry batches
if isinstance(probs, float):
probs = [probs]
return {
"entailment": float(probs[LABEL_ENTAILMENT]),
"neutral": float(probs[LABEL_NEUTRAL]),
"contradiction": float(probs[LABEL_CONTRADICTION]),
}
loop = asyncio.get_event_loop()
result = await asyncio.wait_for(
loop.run_in_executor(None, _sync_score),
timeout=self._timeout_s,
)
return result
# -----------------------------------------------------------------------
# LLM self-check fallback
# -----------------------------------------------------------------------
@staticmethod
def _build_llm_prompt(context: str, answer: str) -> str:
return (
"Given ONLY the following retrieved context, determine whether the answer "
"is faithful to the context. An answer is faithful if every claim it makes "
"is supported by the context. If the answer contradicts the context or adds "
"information not present in the context, it is not faithful.\n\n"
f"Context:\n{context}\n\n"
f"Answer:\n{answer}\n\n"
"Is this answer faithful to the context? Answer Yes or No, then provide "
"brief reasoning.\n\nFormat:\nFaithful: Yes/No\nReasoning: ..."
)
async def _llm_self_check(self, answer: str, context: str) -> CheckResult:
"""Use an LLM (DeepSeek primary, OpenRouter fallback) to check faithfulness when NLI is unavailable."""
api_key = os.environ.get(LLM_API_KEY_ENV, "") or os.environ.get(OPENROUTER_API_KEY_ENV, "")
if not api_key:
logger.error("LLM_API_KEY / OPENROUTER_API_KEY not set; LLM self-check unavailable.")
return CheckResult(
is_faithful=False,
confidence=0.0,
method="llm_self_check_unavailable",
contradiction_details=["LLM API key not configured"],
flagged_claims=[],
)
# Use DeepSeek URL/model if LLM_API_KEY is set, else fall back to OpenRouter
use_llm = bool(os.environ.get(LLM_API_KEY_ENV, ""))
api_url = LLM_BASE_URL if use_llm else OPENROUTER_API_URL
model = LLM_MODEL if use_llm else OPENROUTER_MODEL
prompt = self._build_llm_prompt(context, answer)
payload = {
"model": model,
"messages": [{"role": "user", "content": prompt}],
"temperature": 0.0,
"max_tokens": 256,
}
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
try:
async with httpx.AsyncClient(timeout=self._timeout_s) as client:
resp = await client.post(api_url, json=payload, headers=headers)
resp.raise_for_status()
data = resp.json()
content = data["choices"][0]["message"]["content"].strip()
is_faithful = content.lower().startswith("faithful: yes") or ("faithful: yes" in content.lower()[:80])
# Extract reasoning
reasoning_parts = content.split("Reasoning:")
reasoning = reasoning_parts[1].strip() if len(reasoning_parts) > 1 else content
confidence = 0.7 if is_faithful else 0.3
contradiction_details: list[str] = []
flagged_claims: list[str] = []
if not is_faithful:
contradiction_details.append(reasoning)
flagged_claims.append(answer[:200])
return CheckResult(
is_faithful=is_faithful,
confidence=confidence,
method="llm_self_check",
contradiction_details=contradiction_details,
flagged_claims=flagged_claims,
)
except TimeoutError:
logger.warning("LLM self-check timed out after %ss", self._timeout_s)
return CheckResult(
is_faithful=False,
confidence=0.0,
method="llm_self_check_timeout",
contradiction_details=["LLM self-check timed out"],
flagged_claims=[],
)
except Exception as exc:
logger.error("LLM self-check failed: %s", exc)
return CheckResult(
is_faithful=False,
confidence=0.0,
method="llm_self_check_error",
contradiction_details=[str(exc)],
flagged_claims=[],
)
# -----------------------------------------------------------------------
# Public API — single check
# -----------------------------------------------------------------------
async def check_answer(self, answer: str, context: str) -> CheckResult:
"""
Check a single (answer, context) pair for hallucinations.
Uses NLI model if loaded, otherwise falls back to LLM self-check.
Checks each claim sentence individually and aggregates.
"""
if self._model_loaded:
return await self._check_answer_nli(answer, context)
else:
return await self._llm_self_check(answer, context)
async def _check_answer_nli(self, answer: str, context: str) -> CheckResult:
"""NLI-based check: score each claim sentence against the context."""
claims = _split_claims(answer)
if not claims:
return CheckResult(
is_faithful=True,
confidence=1.0,
method="nli",
contradiction_details=[],
flagged_claims=[],
)
flagged_claims: list[str] = []
contradiction_details: list[str] = []
total_entailment = 0.0
total_neutral = 0.0
total_contradiction = 0.0
for claim in claims:
try:
scores = await asyncio.wait_for(
self._nli_score(context, claim),
timeout=self._timeout_s,
)
except TimeoutError:
logger.warning("NLI scoring timed out for claim: %.80s", claim)
flagged_claims.append(claim)
contradiction_details.append(f"Timeout checking: {claim[:80]}")
total_contradiction += 1.0
continue
except Exception as exc:
logger.error("NLI scoring error: %s", exc)
flagged_claims.append(claim)
contradiction_details.append(f"Error checking: {claim[:80]}")
total_contradiction += 1.0
continue
entailment = scores["entailment"]
neutral = scores["neutral"]
contradiction = scores["contradiction"]
total_entailment += entailment
total_neutral += neutral
total_contradiction += contradiction
if contradiction > entailment and contradiction > neutral:
flagged_claims.append(claim)
contradiction_details.append(f"Contradiction (p={contradiction:.2f}): {claim[:120]}")
n = len(claims)
avg_entailment = total_entailment / n
avg_contradiction = total_contradiction / n
# Faithful = no flagged contradictions. A claim is only flagged when
# contradiction > entailment AND contradiction > neutral (line 411).
# High neutral scores just mean "not enough info to confirm" — not a hallucination.
is_faithful = len(flagged_claims) == 0
confidence = (avg_entailment + (1.0 - avg_contradiction)) / 2.0 if is_faithful else 1.0 - avg_contradiction
return CheckResult(
is_faithful=is_faithful,
confidence=round(confidence, 4),
method="nli",
contradiction_details=contradiction_details,
flagged_claims=flagged_claims,
)
# -----------------------------------------------------------------------
# Public API — multi-source check
# -----------------------------------------------------------------------
async def check_answer_with_sources(self, answer: str, sources: list[dict[str, Any]]) -> CheckResult:
"""
Check an answer against each source individually and aggregate.
Each source dict should have a 'text' or 'content' key with the
source content. The answer is evaluated against each source; if ANY
source entails a claim, it is considered supported.
"""
if not sources:
return CheckResult(
is_faithful=False,
confidence=0.0,
method="nli" if self._model_loaded else "llm_self_check",
contradiction_details=["No sources provided for verification"],
flagged_claims=[],
)
claims = _split_claims(answer)
if not claims:
return CheckResult(
is_faithful=True,
confidence=1.0,
method="nli" if self._model_loaded else "llm_self_check",
)
# Extract text from each source
source_texts: list[str] = []
for src in sources:
text = src.get("text", src.get("content", src.get("page_content", "")))
if text:
source_texts.append(str(text))
if not source_texts:
return CheckResult(
is_faithful=False,
confidence=0.0,
method="nli" if self._model_loaded else "llm_self_check",
contradiction_details=["Sources contain no usable text"],
flagged_claims=[],
)
if self._model_loaded:
return await self._check_answer_with_sources_nli(claims, source_texts)
else:
# Combine sources for LLM self-check
combined_context = "\n\n---\n\n".join(source_texts)
return await self._llm_self_check(answer, combined_context)
async def _check_answer_with_sources_nli(self, claims: list[str], source_texts: list[str]) -> CheckResult:
"""NLI-based multi-source check: a claim is supported if ANY source entails it."""
flagged_claims: list[str] = []
contradiction_details: list[str] = []
claim_verdicts: list[float] = [] # best entailment score per claim
for claim in claims:
best_entailment = 0.0
best_contradiction = 0.0
for source in source_texts:
try:
scores = await asyncio.wait_for(
self._nli_score(source, claim),
timeout=self._timeout_s,
)
except TimeoutError:
logger.warning("NLI timeout for claim against source: %.60s", claim[:60])
continue
except Exception as exc:
logger.error("NLI error in multi-source check: %s", exc)
continue
best_entailment = max(best_entailment, scores["entailment"])
best_contradiction = max(best_contradiction, scores["contradiction"])
claim_verdicts.append(best_entailment)
if best_contradiction > best_entailment and best_contradiction > 0.5:
flagged_claims.append(claim)
contradiction_details.append(f"Contradiction (best p={best_contradiction:.2f}): {claim[:120]}")
n = len(claims)
avg_entailment = sum(claim_verdicts) / n if n else 0.0
has_contradiction = len(flagged_claims) > 0
is_faithful = not has_contradiction
confidence = avg_entailment if is_faithful else 1.0 - (len(flagged_claims) / n)
return CheckResult(
is_faithful=is_faithful,
confidence=round(confidence, 4),
method="nli",
contradiction_details=contradiction_details,
flagged_claims=flagged_claims,
)
# -----------------------------------------------------------------------
# Citation verification
# -----------------------------------------------------------------------
async def verify_citations(self, answer: str, sources: list[dict[str, Any]]) -> CitationVerificationResult:
"""
Parse inline citations [1], [2] from the answer and verify each cited
source actually supports the claim it is attached to.
Each source dict is expected to have a numeric index or be ordered
such that sources[0] corresponds to [1], sources[1] to [2], etc.
"""
citations = _extract_citations(answer)
if not citations:
return CitationVerificationResult(
total_citations=0,
verified_citations=0,
unsupported_citations=[],
citation_details=[],
all_supported=True,
)
citation_details: list[dict[str, Any]] = []
unsupported: list[dict[str, Any]] = []
verified = 0
for cit_num, sentence in citations:
src_idx = cit_num - 1 # [1] -> index 0
source_text = ""
if 0 <= src_idx < len(sources):
src = sources[src_idx]
source_text = src.get("text", src.get("content", src.get("page_content", "")))
source_text = str(source_text) if source_text else ""
detail: dict[str, Any] = {
"citation_number": cit_num,
"sentence": sentence,
"source_found": bool(source_text),
}
if not source_text:
detail["supported"] = False
detail["reason"] = "Source not found or empty"
unsupported.append(detail)
citation_details.append(detail)
continue
# Verify citation supports the claim
if self._model_loaded:
try:
scores = await asyncio.wait_for(
self._nli_score(source_text, sentence),
timeout=self._timeout_s,
)
except TimeoutError:
detail["supported"] = False
detail["reason"] = "NLI timeout"
unsupported.append(detail)
citation_details.append(detail)
continue
except Exception as exc:
detail["supported"] = False
detail["reason"] = f"NLI error: {exc}"
unsupported.append(detail)
citation_details.append(detail)
continue
entailment = scores["entailment"]
contradiction = scores["contradiction"]
neutral = scores["neutral"]
if entailment >= contradiction and entailment >= neutral:
detail["supported"] = True
detail["score"] = scores
verified += 1
elif contradiction > entailment:
detail["supported"] = False
detail["score"] = scores
detail["reason"] = f"Contradiction detected (p={contradiction:.2f})"
unsupported.append(detail)
else:
# Neutral — the source doesn't contradict but doesn't entail either
detail["supported"] = False
detail["score"] = scores
detail["reason"] = "Citation does not support claim (neutral)"
unsupported.append(detail)
else:
# Fallback: LLM self-check for citation
try:
llm_result = await asyncio.wait_for(
self._llm_self_check(sentence, source_text),
timeout=self._timeout_s,
)
detail["supported"] = llm_result.is_faithful
detail["method"] = "llm_self_check"
if not llm_result.is_faithful:
detail["reason"] = "; ".join(llm_result.contradiction_details)
unsupported.append(detail)
else:
verified += 1
except TimeoutError:
detail["supported"] = False
detail["reason"] = "LLM check timed out"
unsupported.append(detail)
citation_details.append(detail)
all_supported = len(unsupported) == 0
return CitationVerificationResult(
total_citations=len(citations),
verified_citations=verified,
unsupported_citations=unsupported,
citation_details=citation_details,
all_supported=all_supported,
)
# -----------------------------------------------------------------------
# Batch checking
# -----------------------------------------------------------------------
async def check_multiple(self, answers_and_contexts: list[tuple[str, str]]) -> list[CheckResult]:
"""
Check multiple (answer, context) pairs concurrently.
Each call is bounded by the timeout; results are returned in order.
"""
if not answers_and_contexts:
return []
tasks = [self.check_answer(answer, context) for answer, context in answers_and_contexts]
results = await asyncio.gather(*tasks, return_exceptions=True)
output: list[CheckResult] = []
for i, result in enumerate(results):
if isinstance(result, Exception):
logger.error("Batch check %d failed: %s", i, result)
output.append(
CheckResult(
is_faithful=False,
confidence=0.0,
method="error",
contradiction_details=[f"Batch check error: {result}"],
flagged_claims=[],
)
)
else:
output.append(result)
return output
# -----------------------------------------------------------------------
# Reset (mainly for testing)
# -----------------------------------------------------------------------
@classmethod
def _reset_singleton(cls) -> None:
"""Reset the singleton instance (for testing purposes only)."""
cls._instance = None
# ---------------------------------------------------------------------------
# Module-level convenience
# ---------------------------------------------------------------------------
async def get_guard() -> HallucinationGuard:
"""Async singleton accessor — the recommended entry point."""
return await HallucinationGuard.get_guard()