task_contract.py

# task_contract.py
# Central contract validation for TaskSpec and naturalized task text.
#
# This module provides:
# - normalize_token(): Consistent token normalization used everywhere
# - validate_task_spec_contract(): Single source of truth for TaskSpec validation
# - text_contains_token(): Safe token presence check (word-boundary aware)

import re
import logging

from config import STRICT_TASK_CONTRACT, DEBUG_SESSION

# Configure logger
logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('[TaskContract] %(levelname)s: %(message)s'))
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG if DEBUG_SESSION else logging.WARNING)


# =============================================================================
# TOKEN NORMALIZATION
# =============================================================================

def normalize_token(token: str) -> str:
    """
    Normalize a token for consistent comparison across the pipeline.
    
    Normalization rules:
    - Lowercase
    - Strip leading/trailing whitespace
    - Strip surrounding quotes (single or double)
    - Strip trailing punctuation (. , ! ?)
    
    Args:
        token: Raw token string
    
    Returns:
        Normalized token (may be empty string if input was just punctuation)
    
    Example:
        normalize_token("'Fire!'") -> "fire"
        normalize_token('"Blue"') -> "blue"
        normalize_token("  YES  ") -> "yes"
    """
    if not token:
        return ""
    
    # Lowercase and strip whitespace
    result = token.lower().strip()
    
    # Strip surrounding quotes
    if len(result) >= 2:
        if (result.startswith("'") and result.endswith("'")) or \
           (result.startswith('"') and result.endswith('"')):
            result = result[1:-1]
    
    # Strip trailing punctuation
    result = result.rstrip(".,!?;:")
    
    # Strip leading punctuation that might remain
    result = result.lstrip(".,!?;:")
    
    return result.strip()


def normalize_targets(targets: list) -> list:
    """
    Normalize a list of target tokens.
    
    Args:
        targets: List of raw target strings
    
    Returns:
        List of normalized, non-empty tokens (duplicates removed, order preserved)
    """
    if not targets:
        return []
    
    seen = set()
    result = []
    for t in targets:
        normalized = normalize_token(str(t) if t else "")
        if normalized and normalized not in seen:
            seen.add(normalized)
            result.append(normalized)
    
    return result


# =============================================================================
# TOKEN PRESENCE CHECK (WORD-BOUNDARY AWARE)
# =============================================================================

def text_contains_token(text: str, token: str) -> bool:
    """
    Check if text contains the token with word-boundary safety.
    
    This prevents matching 'red' inside 'credited' while still allowing
    matches within common sentence structures.
    
    Args:
        text: Text to search (naturalized task text, etc.)
        token: Token to find (already normalized)
    
    Returns:
        True if token appears as a distinct word in text
    """
    if not text or not token:
        return False
    
    text_lower = text.lower()
    token_lower = token.lower()
    
    # Quick check: if token not even present as substring, return False
    if token_lower not in text_lower:
        return False
    
    # Tokenize text into words (simple: alphanumeric sequences)
    text_tokens = set(re.findall(r"[a-z0-9]+", text_lower))
    
    # Check if exact token is present
    if token_lower in text_tokens:
        return True
    
    # Also check if token appears with common suffixes (e.g., "dog" in "dogs")
    # But be conservative - only check if token is prefix of a word
    for text_word in text_tokens:
        if text_word.startswith(token_lower) and len(text_word) <= len(token_lower) + 3:
            # Allow up to 3 extra chars (covers plurals, -ing, -ed, etc.)
            return True
    
    return False


# =============================================================================
# CONTRACT VALIDATION
# =============================================================================

class TaskContractViolation(ValueError):
    """Exception raised when TaskSpec contract is violated in strict mode."""
    pass


def validate_task_spec_contract(
    task_spec: dict,
    naturalized_task_text: str = None,
    strict: bool = None,
    episode_id: str = "",
    session_id: str = "",
    age_band: int = 0,
) -> list:
    """
    Validate TaskSpec contract and naturalized text invariants.
    
    This is the single source of truth for TaskSpec validation.
    Called once per episode after task generation and naturalization.
    
    Args:
        task_spec: The TaskSpec dict (must have targets, may have choices)
        naturalized_task_text: The naturalized text that will be shown to Basil
        strict: If True, raise TaskContractViolation on any violation
                If False/None, log warnings but continue
        episode_id: For logging context
        session_id: For logging context
        age_band: Basil's current age band (affects validation rules)
    
    Returns:
        List of warning messages (empty if no violations)
    
    Raises:
        TaskContractViolation: In strict mode, if any contract is violated
    """
    if strict is None:
        strict = STRICT_TASK_CONTRACT
    
    warnings = []
    context = f"[Session:{session_id or '?'} Episode:{episode_id or '?'}]"
    
    # =========================================================================
    # TARGETS VALIDATION
    # =========================================================================
    
    targets = task_spec.get("targets")
    
    # Check targets exists and is a list
    if targets is None:
        warnings.append(f"{context} Missing 'targets' field in TaskSpec")
        normalized_targets = []
    elif not isinstance(targets, list):
        warnings.append(f"{context} 'targets' must be list, got {type(targets).__name__}")
        normalized_targets = []
    else:
        # Normalize targets
        normalized_targets = normalize_targets(targets)
        
        # Check at least one non-empty target
        if not normalized_targets:
            warnings.append(f"{context} 'targets' is empty after normalization (raw: {targets})")
        
        # Check target count cap
        if len(normalized_targets) > 5:
            warnings.append(f"{context} Too many targets ({len(normalized_targets)} > 5)")
        
        # Check individual target constraints
        for target in normalized_targets:
            # Length cap
            if len(target) > 24:
                warnings.append(f"{context} Target too long: '{target[:30]}...' ({len(target)} > 24)")
            
            # For age_band=0 vocab/control: reject multiword targets
            task_category = task_spec.get("task_category", "vocab")
            if age_band == 0 and task_category in ("vocab", "control"):
                if " " in target:
                    warnings.append(f"{context} Multiword target '{target}' not allowed for age_band=0 {task_category}")
    
    # =========================================================================
    # CHOICES VALIDATION
    # =========================================================================
    
    choices = task_spec.get("choices")
    normalized_choices = []
    
    if choices is not None:
        if not isinstance(choices, list):
            warnings.append(f"{context} 'choices' must be list, got {type(choices).__name__}")
        else:
            normalized_choices = normalize_targets(choices)
            
            # Check choice count cap
            if len(normalized_choices) > 6:
                warnings.append(f"{context} Too many choices ({len(normalized_choices)} > 6)")
    
    # =========================================================================
    # YES/NO TASK VALIDATION
    # =========================================================================
    
    task_text_lower = (task_spec.get("task_text", "") or "").lower()
    is_yes_no_task = "yes or no" in task_text_lower or \
                     (set(normalized_targets) == {"yes", "no"}) or \
                     (set(normalized_choices) == {"yes", "no"})
    
    if is_yes_no_task:
        # Enforce targets contains yes and/or no
        if normalized_targets:
            has_yes_no = "yes" in normalized_targets or "no" in normalized_targets
            if not has_yes_no:
                warnings.append(f"{context} Yes/no task but targets={normalized_targets} missing 'yes' or 'no'")
    
    # =========================================================================
    # NATURALIZED TEXT INVARIANTS
    # =========================================================================
    
    if naturalized_task_text:
        # For each target, require it appears in naturalized text
        for target in normalized_targets:
            if not text_contains_token(naturalized_task_text, target):
                warnings.append(
                    f"{context} Target '{target}' not found in naturalized text: "
                    f"'{naturalized_task_text[:60]}...'"
                )
        
        # For MCQ tasks (choices non-empty), require choices appear in prompt
        if normalized_choices:
            for choice in normalized_choices:
                if not text_contains_token(naturalized_task_text, choice):
                    warnings.append(
                        f"{context} Choice '{choice}' not found in naturalized text: "
                        f"'{naturalized_task_text[:60]}...'"
                    )
    
    # =========================================================================
    # LOG WARNINGS OR RAISE
    # =========================================================================
    
    if warnings:
        for w in warnings:
            logger.warning(w)
        
        if strict:
            raise TaskContractViolation(
                f"TaskSpec contract violated ({len(warnings)} issue(s)): {warnings[0]}"
            )
    
    return warnings


def validate_naturalized_contains_targets(
    naturalized_text: str,
    targets: list,
    choices: list = None,
) -> tuple:
    """
    Validate that naturalized text contains all required target/choice tokens.
    
    Used by the naturalizer to validate LLM output.
    
    Args:
        naturalized_text: The naturalized task text
        targets: List of target tokens (may be raw, will be normalized)
        choices: Optional list of choice tokens for MCQ tasks
    
    Returns:
        Tuple of (is_valid: bool, missing_tokens: list)
    """
    if not naturalized_text:
        return (False, ["empty text"])
    
    missing = []
    
    # Check targets
    normalized_targets = normalize_targets(targets or [])
    for target in normalized_targets:
        if not text_contains_token(naturalized_text, target):
            missing.append(target)
    
    # Check choices (for MCQ)
    if choices:
        normalized_choices = normalize_targets(choices)
        for choice in normalized_choices:
            if not text_contains_token(naturalized_text, choice):
                if choice not in missing:  # Avoid duplicates
                    missing.append(choice)
    
    return (len(missing) == 0, missing)