streaming_intent/config.py

"""
Configuration loader for streaming intent router.
"""

import os
from pathlib import Path
from typing import Dict, List, Optional
from dataclasses import dataclass, field

import yaml
import numpy as np


@dataclass
class StreamingConfig:
    """Configuration for streaming intent router."""
    
    intents: List[str] = field(default_factory=list)
    prior: Dict[str, float] = field(default_factory=dict)
    transition_matrix: np.ndarray = field(default_factory=lambda: np.array([]))
    intent_to_idx: Dict[str, int] = field(default_factory=dict)
    idx_to_intent: Dict[int, str] = field(default_factory=dict)
    
    # Emission parameters
    alpha: float = 1.5
    epsilon: float = 1e-8
    
    # Decision thresholds
    theta_hi: float = 0.85
    theta_med: float = 0.60
    theta_lock: float = 0.70
    K: int = 3
    
    # Window parameters
    max_tokens: int = 64
    max_tokens_limit: int = 128
    
    # Debounce parameters
    debounce_ms: int = 150
    min_change_chars: int = 3
    
    # Model parameters
    model_max_length: int = 48
    
    # Emission temperature (for temperature scaling transform)
    temperature: float = 1.0
    
    # SPRT parameters
    sprt_alpha: float = 0.05
    sprt_beta: float = 0.10
    sprt_p0: float = 0.20
    sprt_p1: float = 0.60
    
    @classmethod
    def from_yaml(cls, path: Optional[str] = None) -> "StreamingConfig":
        """Load configuration from YAML file."""
        if path is None:
            path = Path(__file__).parent.parent / "config" / "streaming_intent.yaml"
        
        with open(path, "r") as f:
            data = yaml.safe_load(f)
        
        config = cls()
        config.intents = data.get("intents", [])
        config.prior = data.get("prior", {})
        
        # Build intent index mappings
        config.intent_to_idx = {intent: i for i, intent in enumerate(config.intents)}
        config.idx_to_intent = {i: intent for i, intent in enumerate(config.intents)}
        
        # Build transition matrix as numpy array
        n = len(config.intents)
        config.transition_matrix = np.zeros((n, n))
        trans_dict = data.get("transition_matrix", {})
        for from_intent, to_probs in trans_dict.items():
            if from_intent in config.intent_to_idx:
                i = config.intent_to_idx[from_intent]
                for to_intent, prob in to_probs.items():
                    if to_intent in config.intent_to_idx:
                        j = config.intent_to_idx[to_intent]
                        config.transition_matrix[i, j] = prob
        
        # Normalize rows (ensure they sum to 1)
        row_sums = config.transition_matrix.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1  # Avoid division by zero
        config.transition_matrix = config.transition_matrix / row_sums
        
        # Emission parameters
        emission = data.get("emission", {})
        config.alpha = emission.get("alpha", 1.5)
        config.epsilon = emission.get("epsilon", 1e-8)
        
        # Decision thresholds
        thresholds = data.get("thresholds", {})
        config.theta_hi = thresholds.get("theta_hi", 0.85)
        config.theta_med = thresholds.get("theta_med", 0.60)
        config.theta_lock = thresholds.get("theta_lock", 0.70)
        config.K = thresholds.get("K", 3)
        
        # Window parameters
        window = data.get("window", {})
        config.max_tokens = window.get("max_tokens", 64)
        config.max_tokens_limit = window.get("max_tokens_limit", 128)
        
        # Debounce parameters
        debounce = data.get("debounce", {})
        config.debounce_ms = debounce.get("debounce_ms", 150)
        config.min_change_chars = debounce.get("min_change_chars", 3)
        
        # Model parameters
        model = data.get("model", {})
        config.model_max_length = model.get("max_length", 64)
        
        # SPRT parameters
        sprt = data.get("sprt", {})
        config.sprt_alpha = sprt.get("alpha", 0.05)
        config.sprt_beta = sprt.get("beta", 0.10)
        config.sprt_p0 = sprt.get("p0", 0.20)
        config.sprt_p1 = sprt.get("p1", 0.60)
        
        # Emission temperature
        config.temperature = emission.get("temperature", 1.0)
        
        # Validate configuration
        if not config.intents:
            raise ValueError(
                f"No intents found in config file: {path}. "
                "The 'intents' list must contain at least one intent."
            )
        
        return config
    
    def get_prior_vector(self) -> np.ndarray:
        """Get prior distribution as numpy array."""
        if not self.intents:
            raise ValueError(
                "Cannot compute prior vector: intents list is empty. "
                "Ensure streaming_intent.yaml contains valid intent definitions."
            )
        
        prior = np.zeros(len(self.intents))
        for intent, prob in self.prior.items():
            if intent in self.intent_to_idx:
                prior[self.intent_to_idx[intent]] = prob
        
        # Normalize
        if prior.sum() > 0:
            prior = prior / prior.sum()
        else:
            # Fallback to uniform distribution
            prior = np.ones(len(self.intents)) / len(self.intents)
        return prior
    
    def save_transition_matrix(self, path: Optional[str] = None) -> None:
        """Save current transition matrix back to config file."""
        if path is None:
            path = Path(__file__).parent.parent / "config" / "streaming_intent.yaml"
        
        with open(path, "r") as f:
            data = yaml.safe_load(f)
        
        # Update transition matrix in data
        trans_dict = {}
        for i, from_intent in enumerate(self.intents):
            trans_dict[from_intent] = {}
            for j, to_intent in enumerate(self.intents):
                trans_dict[from_intent][to_intent] = float(self.transition_matrix[i, j])
        
        data["transition_matrix"] = trans_dict
        
        with open(path, "w") as f:
            yaml.dump(data, f, default_flow_style=False, sort_keys=False)