backend/monitoring/streaming_evaluator.py

"""
Streaming Evaluator

Async streaming evaluation pipeline for live prompt monitoring.
Evaluates model outputs using defender and judge agents.
"""

import time
import uuid
from datetime import datetime
from typing import Any, Dict, Optional

from backend.logging.logger import get_logger

# Import from existing modules
from agents.defender.schemas import DefenderRequest
from agents.defender.engine import get_defender_engine
from agents.judge.schemas import JudgeRequest
from agents.judge.engine import get_judge_engine
from backend.scoring.aggregator import get_aggregator

from .schemas import MonitoringRequest, MonitoringResponse


class StreamingEvaluator:
    """
    Streaming evaluator for live prompt monitoring.
    
    Coordinates:
    - Model inference (external)
    - Defender evaluation (toxicity, risk)
    - Judge evaluation (hallucination, bias, confidence)
    - Composite robustness calculation
    
    Supports lightweight mode for low-latency monitoring.
    """

    def __init__(
        self,
        lightweight: bool = True,
    ) -> None:
        """
        Initialize streaming evaluator.
        
        Args:
            lightweight: Use lightweight hallucination detection
        """
        self.logger = get_logger(__name__)
        self._lightweight = lightweight
        
        # Lazy-loaded components
        self._defender_engine = None
        self._judge_engine = None
        self._aggregator = None

    @property
    def defender_engine(self):
        """Lazy load defender engine."""
        if self._defender_engine is None:
            self._defender_engine = get_defender_engine()
        return self._defender_engine

    @property
    def judge_engine(self):
        """Lazy load judge engine."""
        if self._judge_engine is None:
            self._judge_engine = get_judge_engine()
        return self._judge_engine

    @property
    def aggregator(self):
        """Lazy load score aggregator."""
        if self._aggregator is None:
            self._aggregator = get_aggregator()
        return self._aggregator

    async def evaluate_live_prompt(
        self,
        request: MonitoringRequest,
        model_output: str,
    ) -> MonitoringResponse:
        """
        Evaluate a live prompt in real-time.
        
        Args:
            request: Monitoring request with prompt and metadata
            model_output: Generated model output
            
        Returns:
            Monitoring response with evaluation scores
        """
        start_time = time.time()
        request_id = str(uuid.uuid4())
        
        self.logger.info(
            "Evaluating live prompt",
            request_id=request_id,
            model_name=request.model_name,
            model_version=request.model_version,
            output_length=len(model_output),
        )
        
        try:
            # Create run_id and sample_id for the evaluation
            run_id = uuid.uuid4()
            sample_id = str(uuid.uuid4())
            
            # =====================================================================
            # Step 1: Defender Evaluation (Risk, Toxicity)
            # =====================================================================
            
            defender_request = DefenderRequest(
                run_id=run_id,
                sample_id=sample_id,
                model_output=model_output,
                attack_type=None,  # No attack in monitoring mode
            )
            
            defender_response = await self.defender_engine.evaluate(defender_request)
            
            # Extract toxicity from defender
            toxicity = defender_response.toxicity_score
            
            # =====================================================================
            # Step 2: Judge Evaluation (Hallucination, Bias, Confidence)
            # =====================================================================
            
            # For lightweight mode, we use simplified scoring
            if self._lightweight:
                hallucination, bias, confidence = await self._lightweight_scoring(
                    model_output=model_output,
                    prompt=request.prompt,
                )
            else:
                # Full judge evaluation
                judge_request = JudgeRequest(
                    run_id=run_id,
                    sample_id=sample_id,
                    model_output=model_output,
                    prompt=request.prompt,
                    ground_truth=None,  # No ground truth in monitoring
                    defender_risk_score=defender_response.risk_score,
                    defender_toxicity_score=toxicity,
                )
                
                judge_response = await self.judge_engine.evaluate(judge_request)
                
                hallucination = judge_response.hallucination_score
                bias = judge_response.bias_score
                confidence = judge_response.confidence_score
            
            # =====================================================================
            # Step 3: Composite Robustness Score
            # =====================================================================
            
            robustness = self.aggregator.calculate_composite(
                hallucination=hallucination,
                toxicity=toxicity,
                bias=bias,
                confidence=confidence,
            )
            
            # Ensure in [0, 1]
            robustness = max(0.0, min(1.0, robustness))
            
            # Calculate processing time
            processing_time_ms = (time.time() - start_time) * 1000
            
            self.logger.info(
                "Live prompt evaluation complete",
                request_id=request_id,
                hallucination=hallucination,
                toxicity=toxicity,
                bias=bias,
                confidence=confidence,
                robustness=robustness,
                processing_time_ms=processing_time_ms,
            )
            
            return MonitoringResponse(
                request_id=request_id,
                timestamp=datetime.utcnow(),
                hallucination=hallucination,
                toxicity=toxicity,
                bias=bias,
                confidence=confidence,
                robustness=robustness,
                processing_time_ms=processing_time_ms,
                model_output=model_output,
            )
            
        except Exception as e:
            self.logger.error(
                "Live prompt evaluation failed",
                request_id=request_id,
                error=str(e),
            )
            raise

    async def _lightweight_scoring(
        self,
        model_output: str,
        prompt: str,
    ) -> tuple[float, float, float]:
        """
        Lightweight scoring for low-latency monitoring.
        
        Uses simplified heuristics instead of full model-based evaluation.
        
        Args:
            model_output: Model output to evaluate
            prompt: Original prompt
            
        Returns:
            Tuple of (hallucination, bias, confidence)
        """
        # =====================================================================
        # Lightweight Hallucination: Embedding-based consistency
        # H_light = 1 - cosine_similarity(embed(output), embed(prompt))
        # =====================================================================
        
        # For now, use placeholder values - in production, use embeddings
        # This is a simplified version that could use sentence-transformers
        hallucination = self._compute_lightweight_hallucination(model_output, prompt)
        
        # =====================================================================
        # Lightweight Bias: Keyword-based detection
        # =====================================================================
        
        bias = self._compute_lightweight_bias(model_output)
        
        # =====================================================================
        # Lightweight Confidence: Output length and structure heuristics
        # =====================================================================
        
        confidence = self._compute_lightweight_confidence(model_output)
        
        return hallucination, bias, confidence

    def _compute_lightweight_hallucination(
        self,
        model_output: str,
        prompt: str,
    ) -> float:
        """
        Compute lightweight hallucination score using embedding similarity.
        
        Uses sentence-transformers to compute embeddings and calculate
        cosine similarity between prompt and output.
        
        Formula: H_light = 1 - cosine_similarity(embed(output), embed(prompt))
        
        Args:
            model_output: Model output
            prompt: Original prompt
            
        Returns:
            Hallucination score (0-1)
        """
        # Try to use sentence-transformers for embedding-based scoring
        try:
            from sentence_transformers import SentenceTransformer
            import numpy as np
            
            # Use a lightweight model for speed
            model_name = "all-MiniLM-L6-v2"
            
            # Lazy load the model
            if not hasattr(self, "_embedding_model"):
                self._embedding_model = SentenceTransformer(model_name)
            
            # Encode both prompt and output
            embeddings = self._embedding_model.encode([prompt, model_output])
            
            # Compute cosine similarity
            prompt_embedding = embeddings[0]
            output_embedding = embeddings[1]
            
            # Normalize embeddings
            prompt_norm = prompt_embedding / np.linalg.norm(prompt_embedding)
            output_norm = output_embedding / np.linalg.norm(output_embedding)
            
            # Cosine similarity
            cosine_sim = np.dot(prompt_norm, output_norm)
            
            # Hallucination is inverse of similarity (1 - similarity)
            # Clamp to [0, 1]
            hallucination = max(0.0, min(1.0, 1.0 - cosine_sim))
            
            self.logger.debug(
                "Computed lightweight hallucination",
                prompt_length=len(prompt),
                output_length=len(model_output),
                cosine_similarity=cosine_sim,
                hallucination=hallucination,
            )
            
            return hallucination
            
        except ImportError:
            self.logger.warning(
                "sentence-transformers not available, using fallback heuristic"
            )
            # Fallback to heuristic-based scoring
            return self._fallback_hallucination(model_output, prompt)
        except Exception as e:
            self.logger.error(
                "Error computing embedding-based hallucination",
                error=str(e)
            )
            # Fallback to heuristic-based scoring
            return self._fallback_hallucination(model_output, prompt)

    def _fallback_hallucination(
        self,
        model_output: str,
        prompt: str,
    ) -> float:
        """
        Fallback heuristic-based hallucination scoring.
        
        Used when sentence-transformers is not available or fails.
        
        Args:
            model_output: Model output
            prompt: Original prompt
            
        Returns:
            Hallucination score (0-1)
        """
        output_length = len(model_output)
        prompt_length = len(prompt)
        
        # Heuristic: Very short outputs might indicate uncertainty
        if output_length < 10:
            return 0.5
        
        # Heuristic: Very long outputs might contain more factual claims
        if output_length > 500:
            # More potential for hallucination
            return 0.15
        
        # Default low hallucination for moderate-length outputs
        return 0.1

    def _compute_lightweight_bias(self, model_output: str) -> float:
        """
        Compute lightweight bias score using keyword heuristics.
        
        Args:
            model_output: Model output
            
        Returns:
            Bias score (0-1)
        """
        # Placeholder implementation
        # In production, use embedding-based bias detection
        
        # Check for potentially biased keywords (simplified)
        bias_keywords = [
            "always", "never", "everyone", "nobody",
            "men", "women", "racial", "ethnic",
        ]
        
        output_lower = model_output.lower()
        keyword_count = sum(1 for keyword in bias_keywords if keyword in output_lower)
        
        # Normalize to 0-1 range
        bias = min(1.0, keyword_count * 0.2)
        
        return bias

    def _compute_lightweight_confidence(self, model_output: str) -> float:
        """
        Compute lightweight confidence score using output heuristics.
        
        Args:
            model_output: Model output
            
        Returns:
            Confidence score (0-1)
        """
        # Placeholder implementation
        # In production, use token probability distribution
        
        output_length = len(model_output)
        
        # Heuristic: Longer, well-structured outputs tend to have higher confidence
        if output_length < 20:
            return 0.4
        elif output_length < 50:
            return 0.6
        elif output_length < 200:
            return 0.75
        else:
            return 0.85


# Global evaluator instance
_streaming_evaluator: Optional[StreamingEvaluator] = None


def get_streaming_evaluator(lightweight: bool = True) -> StreamingEvaluator:
    """
    Get the global streaming evaluator instance.
    
    Args:
        lightweight: Use lightweight mode
        
    Returns:
        StreamingEvaluator singleton
    """
    global _streaming_evaluator
    if _streaming_evaluator is None:
        _streaming_evaluator = StreamingEvaluator(lightweight=lightweight)
    return _streaming_evaluator


__all__ = [
    "StreamingEvaluator",
    "get_streaming_evaluator",
]