app/validator.py

"""
Response validator: ensures every response is schema-compliant and grounded.

This is the MOST CRITICAL module — the automated evaluator will reject
any response that doesn't match the exact schema. This validator is the
last line of defense before returning a response.
"""

import json
import re
from app.models import ChatResponse, Recommendation
from app.catalog import Catalog


class ResponseValidator:
    """
    Validates and sanitizes LLM output to guarantee schema compliance.
    """
    
    def __init__(self, catalog: Catalog):
        self.catalog = catalog
    
    def parse_llm_output(self, raw_output: str) -> dict:
        """
        Parse LLM output into a dict, handling various formatting issues.
        The LLM sometimes wraps JSON in markdown code fences or adds extra text.
        """
        text = raw_output.strip()
        
        # Remove markdown code fences if present
        text = re.sub(r'^```(?:json)?\s*', '', text)
        text = re.sub(r'\s*```$', '', text)
        text = text.strip()
        
        # Try to find JSON object in the text
        # Look for the first { and last }
        first_brace = text.find('{')
        last_brace = text.rfind('}')
        
        if first_brace == -1 or last_brace == -1:
            raise ValueError("No JSON object found in LLM output")
        
        json_str = text[first_brace:last_brace + 1]
        
        try:
            return json.loads(json_str, strict=False)
        except json.JSONDecodeError as e:
            # Try to fix common issues
            # Fix trailing commas
            json_str = re.sub(r',\s*}', '}', json_str)
            json_str = re.sub(r',\s*]', ']', json_str)
            return json.loads(json_str, strict=False)
    
    def validate_and_fix(self, parsed: dict) -> ChatResponse:
        """
        Validate parsed LLM output and fix any issues.
        Returns a guaranteed-valid ChatResponse.
        """
        # Extract reply
        reply = parsed.get("reply", "")
        if not reply:
            reply = "I can help you find the right SHL assessment. Could you tell me more about the role?"
        
        # Extract and validate recommendations
        raw_recs = parsed.get("recommendations", [])
        if raw_recs is None:
            raw_recs = []
        
        valid_recs = []
        seen_urls = set()
        
        for rec in raw_recs:
            if not isinstance(rec, dict):
                continue
            
            name = rec.get("name", "")
            url = rec.get("url", "")
            test_type = rec.get("test_type", "K")
            
            # Skip if missing required fields
            if not name or not url:
                continue
            
            # Skip duplicates
            if url in seen_urls:
                continue
            
            # Validate URL is from catalog
            if not self.catalog.validate_url(url):
                # Try to find the assessment by name and use correct URL
                item = self.catalog.find_by_name(name)
                if item:
                    url = item.url
                    test_type = item.test_type
                else:
                    # Skip this recommendation entirely — not in catalog
                    continue
            else:
                # URL is valid — verify test_type from catalog
                item = self.catalog.find_by_url(url)
                if item:
                    test_type = item.test_type
            
            seen_urls.add(url)
            valid_recs.append(Recommendation(
                name=name,
                url=url,
                test_type=test_type,
            ))
        
        # Enforce 1-10 limit
        if len(valid_recs) > 10:
            valid_recs = valid_recs[:10]
        
        # Extract end_of_conversation
        eoc = parsed.get("end_of_conversation", False)
        if not isinstance(eoc, bool):
            eoc = str(eoc).lower() in ("true", "1", "yes")
        
        return ChatResponse(
            reply=reply,
            recommendations=valid_recs,
            end_of_conversation=eoc,
        )
    
    def create_safe_response(self, reply: str = "", eoc: bool = False) -> ChatResponse:
        """Create a safe response when LLM output is completely unusable."""
        if not reply:
            reply = (
                "I apologize, but I encountered an issue processing your request. "
                "Could you rephrase what you're looking for? I can help you find "
                "the right SHL assessment for your hiring needs."
            )
        
        return ChatResponse(
            reply=reply,
            recommendations=[],
            end_of_conversation=eoc,
        )
    
    def create_refusal_response(self, refusal_message: str) -> ChatResponse:
        """Create a response for safety refusals."""
        return ChatResponse(
            reply=refusal_message,
            recommendations=[],
            end_of_conversation=False,
        )