router.py

"""Smart model router for selecting optimal models based on task complexity."""

import re
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from dataclasses import dataclass, field
from abc import ABC, abstractmethod

from .utils import count_tokens, estimate_cost


class TaskComplexity(Enum):
    """Task complexity levels."""
    SIMPLE = "simple"      # Basic Q&A, simple formatting
    MEDIUM = "medium"      # Moderate reasoning, summarization
    COMPLEX = "complex"    # Multi-step reasoning, analysis
    CRITICAL = "critical"  # High-stakes, requires best model


class Provider(Enum):
    """Supported API providers."""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    OPENROUTER = "openrouter"


@dataclass
class ModelConfig:
    """Configuration for a model."""
    name: str
    provider: Provider
    max_tokens: int
    input_cost_per_1m: float  # Cost per 1M input tokens
    output_cost_per_1m: float  # Cost per 1M output tokens
    complexity_levels: List[TaskComplexity]  # Suitable complexity levels
    priority: int = 0  # Higher = preferred when multiple options
    rate_limit: int = 0  # Requests per minute (0 = unlimited)
    context_window: int = 128000  # Context window size

    def estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost for given token counts."""
        input_cost = (input_tokens / 1_000_000) * self.input_cost_per_1m
        output_cost = (output_tokens / 1_000_000) * self.output_cost_per_1m
        return input_cost + output_cost


@dataclass
class RoutingRule:
    """Custom routing rule."""
    name: str
    condition: Callable[[str, Dict[str, Any]], bool]
    model: str
    priority: int = 100  # Higher = checked first


@dataclass
class RoutingDecision:
    """Result of routing decision."""
    model: str
    provider: Provider
    complexity: TaskComplexity
    estimated_cost: float
    reason: str
    alternatives: List[Tuple[str, float]] = field(default_factory=list)


# Default model configurations
# Using OpenRouter as the unified API provider
# Model names follow OpenRouter format: provider/model-name
DEFAULT_MODELS: Dict[str, ModelConfig] = {
    # OpenAI models via OpenRouter
    "openai/gpt-4o": ModelConfig(
        name="openai/gpt-4o",
        provider=Provider.OPENROUTER,
        max_tokens=16384,
        input_cost_per_1m=2.50,
        output_cost_per_1m=10.00,
        complexity_levels=[TaskComplexity.MEDIUM, TaskComplexity.COMPLEX],
        priority=70,  # Lower priority - use Opus for complex
        context_window=128000
    ),
    "openai/gpt-4o-mini": ModelConfig(
        name="openai/gpt-4o-mini",
        provider=Provider.OPENROUTER,
        max_tokens=16384,
        input_cost_per_1m=0.15,
        output_cost_per_1m=0.60,
        complexity_levels=[TaskComplexity.SIMPLE, TaskComplexity.MEDIUM],
        priority=100,  # HIGHEST priority for simple tasks
        context_window=128000
    ),
    "openai/gpt-4-turbo": ModelConfig(
        name="openai/gpt-4-turbo",
        provider=Provider.OPENROUTER,
        max_tokens=4096,
        input_cost_per_1m=10.00,
        output_cost_per_1m=30.00,
        complexity_levels=[TaskComplexity.COMPLEX],
        priority=60,
        context_window=128000
    ),
    "openai/gpt-3.5-turbo": ModelConfig(
        name="openai/gpt-3.5-turbo",
        provider=Provider.OPENROUTER,
        max_tokens=4096,
        input_cost_per_1m=0.50,
        output_cost_per_1m=1.50,
        complexity_levels=[TaskComplexity.SIMPLE],
        priority=50,  # Lower than 4o-mini
        context_window=16385
    ),
    # Anthropic models - Claude Opus 4 for complex/critical tasks
    "anthropic/claude-opus-4.5": ModelConfig(
        name="anthropic/claude-opus-4.5",
        provider=Provider.OPENROUTER,
        max_tokens=32000,
        input_cost_per_1m=15.00,
        output_cost_per_1m=75.00,
        complexity_levels=[TaskComplexity.COMPLEX, TaskComplexity.CRITICAL],
        priority=100,  # HIGHEST priority for complex/critical tasks
        context_window=200000
    ),
    "anthropic/claude-sonnet-4-20250514": ModelConfig(
        name="anthropic/claude-sonnet-4-20250514",
        provider=Provider.OPENROUTER,
        max_tokens=16000,
        input_cost_per_1m=3.00,
        output_cost_per_1m=15.00,
        complexity_levels=[TaskComplexity.MEDIUM, TaskComplexity.COMPLEX],
        priority=80,
        context_window=200000
    ),
    "anthropic/claude-3.5-sonnet": ModelConfig(
        name="anthropic/claude-3.5-sonnet",
        provider=Provider.OPENROUTER,
        max_tokens=8192,
        input_cost_per_1m=3.00,
        output_cost_per_1m=15.00,
        complexity_levels=[TaskComplexity.MEDIUM, TaskComplexity.COMPLEX],
        priority=75,
        context_window=200000
    ),
    "anthropic/claude-3.5-haiku": ModelConfig(
        name="anthropic/claude-3.5-haiku",
        provider=Provider.OPENROUTER,
        max_tokens=8192,
        input_cost_per_1m=0.80,
        output_cost_per_1m=4.00,
        complexity_levels=[TaskComplexity.SIMPLE, TaskComplexity.MEDIUM],
        priority=85,
        context_window=200000
    ),
    "anthropic/claude-3-opus": ModelConfig(
        name="anthropic/claude-3-opus",
        provider=Provider.OPENROUTER,
        max_tokens=4096,
        input_cost_per_1m=15.00,
        output_cost_per_1m=75.00,
        complexity_levels=[TaskComplexity.CRITICAL],
        priority=90,
        context_window=200000
    ),
    "anthropic/claude-3-haiku": ModelConfig(
        name="anthropic/claude-3-haiku",
        provider=Provider.OPENROUTER,
        max_tokens=4096,
        input_cost_per_1m=0.25,
        output_cost_per_1m=1.25,
        complexity_levels=[TaskComplexity.SIMPLE],
        priority=60,
        context_window=200000
    ),



class ComplexityAnalyzer:
    """Analyzes prompt complexity."""

    # Patterns indicating higher complexity
    COMPLEX_PATTERNS = [
        r"analyze",
        r"compare",
        r"explain.*why",
        r"step[- ]by[- ]step",
        r"reason(ing)?",
        r"prove",
        r"derive",
        r"evaluate",
        r"critique",
        r"synthesize",
    ]

    CRITICAL_PATTERNS = [
        r"critical",
        r"important",
        r"accurate",
        r"precise",
        r"legal",
        r"medical",
        r"financial",
        r"security",
        r"compliance",
    ]

    SIMPLE_PATTERNS = [
        r"translate",
        r"summarize",
        r"list",
        r"format",
        r"convert",
        r"extract",
    ]

    def __init__(self):
        self.complex_re = [re.compile(p, re.IGNORECASE) for p in self.COMPLEX_PATTERNS]
        self.critical_re = [re.compile(p, re.IGNORECASE) for p in self.CRITICAL_PATTERNS]
        self.simple_re = [re.compile(p, re.IGNORECASE) for p in self.SIMPLE_PATTERNS]

    def analyze(self, prompt: str, metadata: Optional[Dict[str, Any]] = None) -> TaskComplexity:
        """
        Analyze prompt complexity.

        Args:
            prompt: The prompt text
            metadata: Optional metadata with hints

        Returns:
            Detected complexity level
        """
        metadata = metadata or {}

        # Check for explicit complexity hint
        if "complexity" in metadata:
            try:
                return TaskComplexity(metadata["complexity"])
            except ValueError:
                pass

        # Check for critical patterns first
        critical_score = sum(1 for r in self.critical_re if r.search(prompt))
        if critical_score >= 2 or metadata.get("critical", False):
            return TaskComplexity.CRITICAL

        # Check for complex patterns
        complex_score = sum(1 for r in self.complex_re if r.search(prompt))
        simple_score = sum(1 for r in self.simple_re if r.search(prompt))

        # Token count affects complexity
        token_count = count_tokens(prompt)

        # Scoring logic
        if complex_score >= 3 or token_count > 2000:
            return TaskComplexity.COMPLEX
        elif complex_score >= 1 or token_count > 500:
            return TaskComplexity.MEDIUM
        elif simple_score >= 1 or token_count < 200:
            return TaskComplexity.SIMPLE
        else:
            return TaskComplexity.MEDIUM


class ModelRouter:
    """Smart router for selecting optimal models."""

    def __init__(
        self,
        models: Optional[Dict[str, ModelConfig]] = None,
        default_model: str = "openai/gpt-4o-mini",
        allowed_providers: Optional[List[Provider]] = None,
        cost_optimization: bool = True,
        max_cost_per_request: Optional[float] = None
    ):
        """
        Initialize the router.

        Args:
            models: Model configurations (uses defaults if None)
            default_model: Fallback model
            allowed_providers: List of allowed providers (all if None)
            cost_optimization: Whether to prefer cheaper models
            max_cost_per_request: Maximum allowed cost per request
        """
        self.models = models or DEFAULT_MODELS.copy()
        self.default_model = default_model
        self.allowed_providers = allowed_providers
        self.cost_optimization = cost_optimization
        self.max_cost_per_request = max_cost_per_request
        self.custom_rules: List[RoutingRule] = []
        self.complexity_analyzer = ComplexityAnalyzer()

    def add_model(self, config: ModelConfig):
        """Add or update a model configuration."""
        self.models[config.name] = config

    def remove_model(self, name: str):
        """Remove a model configuration."""
        self.models.pop(name, None)

    def add_rule(self, rule: RoutingRule):
        """Add a custom routing rule."""
        self.custom_rules.append(rule)
        self.custom_rules.sort(key=lambda r: r.priority, reverse=True)

    def remove_rule(self, name: str):
        """Remove a custom routing rule by name."""
        self.custom_rules = [r for r in self.custom_rules if r.name != name]

    def route(
        self,
        prompt: str,
        metadata: Optional[Dict[str, Any]] = None,
        expected_output_tokens: int = 500,
        force_complexity: Optional[TaskComplexity] = None,
        force_provider: Optional[Provider] = None
    ) -> RoutingDecision:
        """
        Route a request to the optimal model.

        Args:
            prompt: The prompt text
            metadata: Optional metadata for routing decisions
            expected_output_tokens: Expected output token count
            force_complexity: Override detected complexity
            force_provider: Force specific provider

        Returns:
            RoutingDecision with selected model and details
        """
        metadata = metadata or {}
        input_tokens = count_tokens(prompt)

        # Check custom rules first
        for rule in self.custom_rules:
            try:
                if rule.condition(prompt, metadata):
                    if rule.model in self.models:
                        model_config = self.models[rule.model]
                        return RoutingDecision(
                            model=rule.model,
                            provider=model_config.provider,
                            complexity=TaskComplexity.MEDIUM,  # Custom rule
                            estimated_cost=model_config.estimate_cost(
                                input_tokens, expected_output_tokens
                            ),
                            reason=f"Custom rule: {rule.name}"
                        )
            except Exception:
                continue

        # Analyze complexity
        complexity = force_complexity or self.complexity_analyzer.analyze(prompt, metadata)

        # Filter eligible models
        eligible = []
        for name, config in self.models.items():
            # Check provider filter
            if force_provider and config.provider != force_provider:
                continue
            if self.allowed_providers and config.provider not in self.allowed_providers:
                continue

            # Check complexity match
            if complexity not in config.complexity_levels:
                continue

            # Check context window
            if input_tokens > config.context_window * 0.9:  # 90% safety margin
                continue

            # Calculate cost
            cost = config.estimate_cost(input_tokens, expected_output_tokens)

            # Check cost limit
            if self.max_cost_per_request and cost > self.max_cost_per_request:
                continue

            eligible.append((name, config, cost))

        if not eligible:
            # Fallback to default
            default_config = self.models.get(self.default_model)
            if default_config:
                return RoutingDecision(
                    model=self.default_model,
                    provider=default_config.provider,
                    complexity=complexity,
                    estimated_cost=default_config.estimate_cost(
                        input_tokens, expected_output_tokens
                    ),
                    reason="Fallback: no eligible models found"
                )
            raise ValueError("No eligible models and no valid default")

        # Sort by preference
        if self.cost_optimization:
            # Sort by cost (ascending), then priority (descending)
            eligible.sort(key=lambda x: (x[2], -x[1].priority))
        else:
            # Sort by priority (descending), then cost (ascending)
            eligible.sort(key=lambda x: (-x[1].priority, x[2]))

        selected_name, selected_config, selected_cost = eligible[0]

        # Build alternatives list
        alternatives = [(name, cost) for name, _, cost in eligible[1:5]]

        return RoutingDecision(
            model=selected_name,
            provider=selected_config.provider,
            complexity=complexity,
            estimated_cost=selected_cost,
            reason=f"Best match for {complexity.value} task",
            alternatives=alternatives
        )

    def get_model_config(self, name: str) -> Optional[ModelConfig]:
        """Get configuration for a specific model."""
        return self.models.get(name)

    def list_models(
        self,
        provider: Optional[Provider] = None,
        complexity: Optional[TaskComplexity] = None
    ) -> List[ModelConfig]:
        """List available models with optional filters."""
        result = []
        for config in self.models.values():
            if provider and config.provider != provider:
                continue
            if complexity and complexity not in config.complexity_levels:
                continue
            result.append(config)
        return sorted(result, key=lambda x: x.priority, reverse=True)

    def estimate_costs(
        self,
        prompt: str,
        expected_output_tokens: int = 500
    ) -> Dict[str, float]:
        """Estimate costs for all available models."""
        input_tokens = count_tokens(prompt)
        costs = {}
        for name, config in self.models.items():
            if self.allowed_providers and config.provider not in self.allowed_providers:
                continue
            costs[name] = config.estimate_cost(input_tokens, expected_output_tokens)
        return dict(sorted(costs.items(), key=lambda x: x[1]))