core/postprocessor.py

import re
import logging
from typing import List

logger = logging.getLogger("axiom.postprocessor")


class PostProcessor:
    """
    Cleans generated text and injects source citations.
    """

    def __init__(self, groq_gen=None):
        self.groq_gen = groq_gen
        logger.info("Ready.")

    def _remove_repetition(self, text: str) -> str:
        """Removes HTML artifacts and duplicate sentences."""

        # Strip any HTML tags that leaked through
        text = re.sub(r'<[^>]+>', ' ', text)
        text = re.sub(r'\s+', ' ', text).strip()

        sentences = text.split(". ")
        seen      = set()
        cleaned   = []

        for sentence in sentences:
            normalized = sentence.strip().lower()
            if normalized and normalized not in seen:
                seen.add(normalized)
                cleaned.append(sentence.strip())

        return ". ".join(cleaned)

    def _inject_citations(self, text: str, sources: List[str]) -> str:
        """Appends source list at the bottom of the generated text."""
        if not sources:
            return text

        citation_block = "\n\nSources:\n"
        for i, source in enumerate(sources):
            citation_block += f"  [{i+1}] {source}\n"

        return text + citation_block

    def _check_faithfulness(self, generated: str, context: str) -> dict:
        """
        LLM-based faithfulness check.
        """
        if not self.groq_gen:
            return {
                "faithfulness_score": 0.5,
                "supported_sentences": 0,
                "unsupported_sentences": 0
            }

        # Truncate context to avoid token limits causing failures
        truncated_context = context[:3000] if len(context) > 3000 else context
        truncated_generated = generated[:1000] if len(generated) > 1000 else generated

        prompt = f"""You are a strict fact-checker. Evaluate if the Generated Text is faithful to the Context.

Context:
{truncated_context}

Generated Text:
{truncated_generated}

Rules:
- faithfulness_score: 1.0 means every claim is supported by context
- faithfulness_score: 0.0 means major claims are not in context
- Count each sentence as supported or unsupported

Respond ONLY with valid JSON, no markdown, no explanation:
{{"faithfulness_score": 0.95, "supported_sentences": 4, "unsupported_sentences": 0}}"""

        try:
            res = self.groq_gen.generate(query=prompt, context="")
            # More aggressive cleaning
            clean_res = res.strip()
            clean_res = re.sub(r'```json\s*', '', clean_res)
            clean_res = re.sub(r'```\s*', '', clean_res)
            clean_res = clean_res.strip()

            # Find JSON object in response even if there's extra text
            json_match = re.search(r'\{[^}]+\}', clean_res)
            if json_match:
                clean_res = json_match.group()

            import json
            data = json.loads(clean_res)

            score = float(data.get("faithfulness_score", 0.5))
            # Clamp between 0 and 1
            score = max(0.0, min(1.0, score))

            return {
                "faithfulness_score": score,
                "supported_sentences": int(data.get("supported_sentences", 0)),
                "unsupported_sentences": int(data.get("unsupported_sentences", 0))
            }
        except Exception as e:
            logger.error(f"Faithfulness check failed: {e}")
            # Return 0.5 instead of 0.0 so badge still shows
            return {
                "faithfulness_score": 0.5,
                "supported_sentences": 0,
                "unsupported_sentences": 0
            }

    def process(self, generated: str, sources: List[str], context: str) -> dict:
        """
        Full post-processing pipeline.
        Returns cleaned text + citations + faithfulness check.
        """
        cleaned        = self._remove_repetition(generated)
        with_citations = self._inject_citations(cleaned, sources)
        faithfulness   = self._check_faithfulness(cleaned, context)

        return {
            "final_output": with_citations,
            "clean_output": cleaned,
            "sources":      sources,
            "faithfulness": faithfulness
        }