generation/validator.py

"""
generation/validator.py — Post-generation citation validator

KEY DIFFERENTIATOR: Most RAG systems only enforce citations at the prompt
level. This validator structurally checks every sentence post-generation.

For every sentence in the draft:
  - Does it contain [CHUNK-id]?
  - Is that chunk_id in the evidence set?
  - Sentences without citations are flagged

Output: ValidationReport with uncited_ratio and support_score.
"""

import re
from dataclasses import dataclass, field

from models import EvidenceChunk
from generation.prompts import EXPECTED_SECTIONS

CITATION_PATTERN = re.compile(r"\[CHUNK-([a-f0-9]{8})\]")
NOT_EVIDENCED_PATTERN = re.compile(r"\[NOT EVIDENCED\]")
CONTRADICTION_PATTERN = re.compile(r"\[CONTRADICTION:")


@dataclass
class ValidationReport:
    """Result of post-generation citation validation."""
    total_sentences: int = 0
    cited_sentences: int = 0
    uncited_sentences: list[str] = field(default_factory=list)
    invalid_citations: list[str] = field(default_factory=list)
    not_evidenced_count: int = 0
    contradiction_count: int = 0
    uncited_ratio: float = 0.0
    support_score: float = 1.0
    section_completeness: float = 0.0
    missing_sections: list[str] = field(default_factory=list)

    def is_clean(self) -> bool:
        """A clean report has < 10% uncited and no invalid citations."""
        return self.uncited_ratio < 0.10 and not self.invalid_citations

    def severity(self) -> str:
        """Return severity level for UI display."""
        if self.uncited_ratio < 0.10:
            return "green"
        elif self.uncited_ratio < 0.25:
            return "yellow"
        return "red"

    def summary(self) -> str:
        return (
            f"Sentences: {self.total_sentences} total, "
            f"{self.cited_sentences} cited, "
            f"{len(self.uncited_sentences)} uncited | "
            f"uncited_ratio: {self.uncited_ratio:.2%} | "
            f"support_score: {self.support_score:.2f} | "
            f"sections: {self.section_completeness:.0%}"
        )

    def to_dict(self) -> dict:
        return {
            "total_sentences": self.total_sentences,
            "cited_sentences": self.cited_sentences,
            "uncited_sentences": self.uncited_sentences,
            "invalid_citations": self.invalid_citations,
            "not_evidenced_count": self.not_evidenced_count,
            "contradiction_count": self.contradiction_count,
            "uncited_ratio": self.uncited_ratio,
            "support_score": self.support_score,
            "section_completeness": self.section_completeness,
            "missing_sections": self.missing_sections,
            "severity": self.severity(),
            "is_clean": self.is_clean(),
        }


class CitationValidator:
    """
    Post-generation validator. Checks every sentence for proper citations.
    """

    def validate(
        self,
        draft_text: str,
        evidence: list[EvidenceChunk],
    ) -> ValidationReport:
        """
        Validate all citations in a generated draft.

        Parameters
        ----------
        draft_text : str
            The generated draft text.
        evidence : list[EvidenceChunk]
            Evidence chunks used for generation.

        Returns
        -------
        ValidationReport
        """
        valid_ids = {c.chunk_id[:8] for c in evidence}
        sentences = self._split_sentences(draft_text)

        # Filter to content sentences (skip headers and short lines)
        content_sentences = [
            s for s in sentences
            if len(s.strip()) > 20
            and not s.strip().startswith("#")
        ]

        cited: list[str] = []
        uncited: list[str] = []
        invalid_citations: list[str] = []
        not_evidenced_count = 0
        contradiction_count = 0

        for sentence in content_sentences:
            # Check for explicit markers
            if NOT_EVIDENCED_PATTERN.search(sentence):
                not_evidenced_count += 1
                cited.append(sentence)
                continue

            if CONTRADICTION_PATTERN.search(sentence):
                contradiction_count += 1
                cited.append(sentence)
                continue

            # Check for citation references
            found_ids = CITATION_PATTERN.findall(sentence)

            if not found_ids:
                uncited.append(sentence.strip())
            else:
                cited.append(sentence)
                for cid in found_ids:
                    if cid not in valid_ids:
                        invalid_citations.append(
                            f"Unknown chunk [{cid}] in: "
                            f"{sentence.strip()[:80]}..."
                        )

        total = len(content_sentences)
        uncited_ratio = len(uncited) / total if total > 0 else 0.0

        # Section completeness check
        missing = self._check_sections(draft_text)

        return ValidationReport(
            total_sentences=total,
            cited_sentences=len(cited),
            uncited_sentences=uncited,
            invalid_citations=invalid_citations,
            not_evidenced_count=not_evidenced_count,
            contradiction_count=contradiction_count,
            uncited_ratio=round(uncited_ratio, 3),
            support_score=round(1.0 - uncited_ratio, 3),
            section_completeness=round(
                1.0 - len(missing) / len(EXPECTED_SECTIONS), 2
            ),
            missing_sections=missing,
        )

    def _split_sentences(self, text: str) -> list[str]:
        """
        Split text into sentences, keeping citations attached.

        Splits on newlines first, then on sentence-ending punctuation
        only when followed by an uppercase letter (not a bracket like [CHUNK-...]).
        """
        # First split on newlines
        lines = [line.strip() for line in text.split("\n") if line.strip()]

        # Then split lines that contain multiple sentences
        # Only split on ". " (or "! ", "? ") followed by an uppercase letter
        result: list[str] = []
        for line in lines:
            # Don't sub-split header lines — they contain periods in numbering
            if line.startswith("#"):
                result.append(line)
            else:
                parts = re.split(r"(?<=[.!?])\s+(?=[A-Z])", line)
                result.extend(p.strip() for p in parts if p.strip())

        return result

    def _check_sections(self, draft_text: str) -> list[str]:
        """Check which expected sections are missing from the draft."""
        missing: list[str] = []
        text_lower = draft_text.lower()
        for section in EXPECTED_SECTIONS:
            if section.lower() not in text_lower:
                missing.append(section)
        return missing