evaluators/registry.py

"""
Evaluator Registry - Central registration for all evaluators.

This module provides a registry pattern for managing evaluators.
Each metric name maps to exactly one evaluator class with optional UI metadata.
"""
from typing import Dict, Type, Optional
from dataclasses import dataclass
from evaluators.base import Evaluator


@dataclass
class MetricInfo:
    """Metadata for a metric (used for UI display)."""
    key: str
    label: str
    description: str = ""
    category: str = ""  # Optional: group metrics by category


class EvaluatorRegistry:
    """Registry for managing evaluator classes and their metadata."""
    
    def __init__(self):
        self._registry: Dict[str, Type[Evaluator]] = {}
        self._metadata: Dict[str, MetricInfo] = {}
    
    def register(
        self, 
        metric_name: str, 
        evaluator_class: Type[Evaluator],
        label: Optional[str] = None,
        description: str = "",
        category: str = ""
    ):
        """
        Register an evaluator class for a metric with optional UI metadata.
        
        Args:
            metric_name: The metric name (e.g., "talk_type", "empathy_er")
            evaluator_class: The evaluator class
            label: Human-readable label for UI (defaults to formatted metric_name)
            description: Description of what this metric measures
            category: Optional category for grouping metrics in UI
        """
        if metric_name in self._registry:
            raise ValueError(f"Metric '{metric_name}' is already registered")
        
        self._registry[metric_name] = evaluator_class
        self._metadata[metric_name] = MetricInfo(
            key=metric_name,
            label=label or metric_name.replace('_', ' ').title(),
            description=description,
            category=category
        )
    
    def get(self, metric_name: str) -> Optional[Type[Evaluator]]:
        """
        Get the evaluator class for a metric.
        
        Args:
            metric_name: The metric name
            
        Returns:
            Evaluator class or None if not found
        """
        return self._registry.get(metric_name)
    
    def get_metadata(self, metric_name: str) -> Optional[MetricInfo]:
        """Get metadata for a metric."""
        return self._metadata.get(metric_name)
    
    def list_metrics(self) -> list[str]:
        """Get list of all registered metric names."""
        return list(self._registry.keys())
    
    def get_ui_labels(self) -> Dict[str, str]:
        """Get metric key -> label mapping for UI display."""
        return {k: v.label for k, v in self._metadata.items()}
    
    def get_metrics_by_category(self) -> Dict[str, list[str]]:
        """Get metrics grouped by category."""
        categories: Dict[str, list[str]] = {}
        for key, info in self._metadata.items():
            cat = info.category or "Other"
            if cat not in categories:
                categories[cat] = []
            categories[cat].append(key)
        return categories
    
    def create_evaluator(self, metric_name: str, **kwargs) -> Optional[Evaluator]:
        """
        Create an evaluator instance for a metric.
        
        Args:
            metric_name: The metric name
            **kwargs: Arguments to pass to evaluator constructor
            
        Returns:
            Evaluator instance or None if metric not found
        """
        evaluator_class = self.get(metric_name)
        if evaluator_class:
            return evaluator_class(**kwargs)
        return None


# Global registry instance
_global_registry = EvaluatorRegistry()


def register_evaluator(
    metric_name: str,
    label: Optional[str] = None,
    description: str = "",
    category: str = ""
):
    """
    Decorator to register an evaluator class with optional UI metadata.
    
    Args:
        metric_name: Unique metric identifier
        label: Human-readable label for UI (optional)
        description: What this metric measures
        category: Category for grouping in UI (e.g., "Empathy", "Communication")
    
    Usage:
        @register_evaluator("talk_type", label="Talk Type", category="Communication")
        class TalkTypeEvaluator(Evaluator):
            METRIC_NAME = "talk_type"
            ...
    """
    def decorator(evaluator_class: Type[Evaluator]):
        _global_registry.register(metric_name, evaluator_class, label, description, category)
        return evaluator_class
    return decorator


def get_evaluator_class(metric_name: str) -> Optional[Type[Evaluator]]:
    """Get evaluator class for a metric."""
    return _global_registry.get(metric_name)


def create_evaluator(metric_name: str, **kwargs) -> Optional[Evaluator]:
    """Create evaluator instance for a metric."""
    return _global_registry.create_evaluator(metric_name, **kwargs)


def get_metric_metadata(metric_name: str) -> Optional[MetricInfo]:
    """Get metadata for a specific metric."""
    return _global_registry.get_metadata(metric_name)


def list_available_metrics() -> list[str]:
    """Get list of all available metrics."""
    return _global_registry.list_metrics()


def get_ui_labels() -> Dict[str, str]:
    """Get metric key -> label mapping for UI display."""
    return _global_registry.get_ui_labels()


def get_metrics_by_category() -> Dict[str, list[str]]:
    """Get metrics grouped by category."""
    return _global_registry.get_metrics_by_category()


def get_registry() -> EvaluatorRegistry:
    """Get the global registry instance."""
    return _global_registry