stratifyai 0.1.0__py3-none-any.whl

stratifyai/router.py

@@ -0,0 +1,643 @@
+"""Intelligent model router for selecting optimal LLM providers."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+import re
+
+from .config import MODEL_CATALOG
+from .models import Message
+from .utils.file_analyzer import FileType
+
+
+class RoutingStrategy(str, Enum):
+    """Routing strategy types."""
+    COST = "cost"
+    QUALITY = "quality"
+    LATENCY = "latency"
+    HYBRID = "hybrid"
+
+
+@dataclass
+class ModelMetadata:
+    """Metadata for a specific model."""
+    provider: str
+    model: str
+    quality_score: float  # 0.0 - 1.0
+    cost_per_1m_input: float  # USD per 1M tokens
+    cost_per_1m_output: float  # USD per 1M tokens
+    avg_latency_ms: float  # Average response time in milliseconds
+    context_window: int  # Maximum context tokens
+    capabilities: List[str]  # e.g., ["reasoning", "vision", "tools"]
+    reasoning_model: bool = False
+    supports_streaming: bool = True
+
+
+class Router:
+    """Intelligent model router for selecting optimal providers."""
+    
+    def __init__(
+        self,
+        strategy: RoutingStrategy = RoutingStrategy.HYBRID,
+        preferred_providers: Optional[List[str]] = None,
+        excluded_providers: Optional[List[str]] = None,
+    ):
+        """
+        Initialize router.
+        
+        Args:
+            strategy: Routing strategy to use
+            preferred_providers: List of preferred providers (prioritized)
+            excluded_providers: List of providers to exclude
+        """
+        self.strategy = strategy
+        self.preferred_providers = preferred_providers or []
+        self.excluded_providers = excluded_providers or []
+        self._load_model_metadata()
+    
+    def _load_model_metadata(self) -> None:
+        """Load model metadata from MODEL_CATALOG."""
+        self.model_metadata: Dict[str, ModelMetadata] = {}
+        
+        # Quality scores (estimated based on benchmarks)
+        # These are rough estimates - should be updated with real benchmark data
+        quality_scores = {
+            # OpenAI
+            "gpt-5": 0.98,
+            "o3-mini": 0.95,
+            "gpt-4.5-turbo-20250205": 0.93,
+            "gpt-4.1-turbo": 0.90,
+            "gpt-4.1-mini": 0.82,
+            "o1-mini": 0.88,
+            "o1": 0.96,
+            
+            # Anthropic
+            "claude-sonnet-4-5-20250929": 0.94,
+            "claude-3-5-sonnet-20241022": 0.92,
+            "claude-3-5-haiku-20241022": 0.80,
+            
+            # Google
+            "gemini-2.5-pro": 0.91,
+            "gemini-2.5-flash": 0.85,
+            "gemini-2.5-flash-lite": 0.78,
+            
+            # DeepSeek
+            "deepseek-chat": 0.85,
+            "deepseek-reasoner": 0.90,
+            
+            # Groq
+            "llama-3.1-70b-versatile": 0.83,
+            "llama-3.1-8b-instant": 0.75,
+            "mixtral-8x7b-32768": 0.80,
+            
+            # Grok
+            "grok-beta": 0.87,
+            
+            # OpenRouter (using same as original providers)
+            "anthropic/claude-3-5-sonnet": 0.92,
+            "openai/gpt-4-turbo": 0.90,
+            
+            # Ollama (local models)
+            "llama3.2": 0.70,
+            "mistral": 0.68,
+            "codellama": 0.72,
+        }
+        
+        # Average latency (ms) - rough estimates
+        latency_estimates = {
+            # OpenAI
+            "gpt-5": 3500,
+            "o3-mini": 8000,
+            "gpt-4.5-turbo-20250205": 2500,
+            "gpt-4.1-turbo": 2000,
+            "gpt-4.1-mini": 800,
+            "o1-mini": 5000,
+            "o1": 10000,
+            
+            # Anthropic
+            "claude-sonnet-4-5-20250929": 2800,
+            "claude-3-5-sonnet-20241022": 2200,
+            "claude-3-5-haiku-20241022": 1200,
+            
+            # Google
+            "gemini-2.5-pro": 2000,
+            "gemini-2.5-flash": 1000,
+            "gemini-2.5-flash-lite": 600,
+            
+            # DeepSeek
+            "deepseek-chat": 1500,
+            "deepseek-reasoner": 6000,
+            
+            # Groq (known for speed)
+            "llama-3.1-70b-versatile": 400,
+            "llama-3.1-8b-instant": 200,
+            "mixtral-8x7b-32768": 350,
+            
+            # Grok
+            "grok-beta": 2500,
+            
+            # OpenRouter
+            "anthropic/claude-3-5-sonnet": 2500,
+            "openai/gpt-4-turbo": 2200,
+            
+            # Ollama (local, very fast)
+            "llama3.2": 100,
+            "mistral": 120,
+            "codellama": 110,
+        }
+        
+        # Build metadata from MODEL_CATALOG
+        for provider, models in MODEL_CATALOG.items():
+            # Skip if provider is excluded
+            if provider in self.excluded_providers:
+                continue
+                
+            for model_name, model_info in models.items():
+                key = f"{provider}/{model_name}"
+                
+                # Extract capabilities
+                capabilities = []
+                if model_info.get("supports_vision"):
+                    capabilities.append("vision")
+                if model_info.get("supports_tools"):
+                    capabilities.append("tools")
+                if model_info.get("reasoning_model"):
+                    capabilities.append("reasoning")
+                
+                self.model_metadata[key] = ModelMetadata(
+                    provider=provider,
+                    model=model_name,
+                    quality_score=quality_scores.get(model_name, 0.75),
+                    cost_per_1m_input=model_info.get("cost_input", 0.0),
+                    cost_per_1m_output=model_info.get("cost_output", 0.0),
+                    avg_latency_ms=latency_estimates.get(model_name, 2000),
+                    context_window=model_info.get("context", 8192),
+                    capabilities=capabilities,
+                    reasoning_model=model_info.get("reasoning_model", False),
+                    supports_streaming=True,
+                )
+    
+    def route(
+        self,
+        messages: List[Message],
+        required_capabilities: Optional[List[str]] = None,
+        max_cost_per_1k_tokens: Optional[float] = None,
+        max_latency_ms: Optional[float] = None,
+        min_context_window: Optional[int] = None,
+    ) -> Tuple[str, str]:
+        """
+        Select the best model for the given request.
+        
+        Args:
+            messages: Conversation messages
+            required_capabilities: Required model capabilities (e.g., ["vision", "tools"])
+            max_cost_per_1k_tokens: Maximum acceptable cost per 1k tokens
+            max_latency_ms: Maximum acceptable latency in milliseconds
+            min_context_window: Minimum required context window
+        
+        Returns:
+            Tuple of (provider, model) for the selected model
+        """
+        # Analyze prompt complexity
+        complexity = self._analyze_complexity(messages)
+        
+        # Filter models by requirements
+        candidates = self._filter_candidates(
+            required_capabilities,
+            max_cost_per_1k_tokens,
+            max_latency_ms,
+            min_context_window,
+        )
+        
+        if not candidates:
+            raise ValueError(
+                "No models meet the specified requirements. "
+                "Consider relaxing constraints."
+            )
+        
+        # Select model based on strategy
+        if self.strategy == RoutingStrategy.COST:
+            selected_key = self._select_by_cost(candidates)
+        elif self.strategy == RoutingStrategy.QUALITY:
+            selected_key = self._select_by_quality(candidates, complexity)
+        elif self.strategy == RoutingStrategy.LATENCY:
+            selected_key = self._select_by_latency(candidates)
+        elif self.strategy == RoutingStrategy.HYBRID:
+            selected_key = self._select_hybrid(candidates, complexity)
+        else:
+            selected_key = list(candidates.keys())[0]
+        
+        metadata = self.model_metadata[selected_key]
+        return metadata.provider, metadata.model
+    
+    def _analyze_complexity(self, messages: List[Message]) -> float:
+        """
+        Analyze prompt complexity (0.0 - 1.0).
+        
+        Higher complexity scores indicate more difficult tasks requiring
+        better models.
+        
+        Args:
+            messages: Conversation messages
+        
+        Returns:
+            Complexity score from 0.0 (simple) to 1.0 (complex)
+        """
+        # Combine all message content
+        text = " ".join(msg.content for msg in messages)
+        
+        # Initialize score
+        complexity_score = 0.0
+        
+        # 1. Check for reasoning keywords (40% weight)
+        reasoning_keywords = [
+            r'\banalyze\b', r'\bexplain\b', r'\breasoning\b', r'\bproof\b',
+            r'\bstep by step\b', r'\bcomplex\b', r'\bcalculate\b',
+            r'\bderive\b', r'\bsolve\b', r'\bprove\b', r'\bthink\b',
+            r'\bcompare\b', r'\bevaluate\b', r'\bdetailed\b',
+        ]
+        
+        reasoning_matches = sum(
+            1 for pattern in reasoning_keywords
+            if re.search(pattern, text.lower())
+        )
+        complexity_score += min(reasoning_matches / len(reasoning_keywords), 1.0) * 0.4
+        
+        # 2. Length-based complexity (20% weight)
+        # Longer prompts often indicate more complex tasks
+        text_length = len(text)
+        length_score = min(text_length / 2000, 1.0)  # Normalize to 2000 chars
+        complexity_score += length_score * 0.2
+        
+        # 3. Check for code/technical content (20% weight)
+        code_indicators = [
+            r'```', r'function\s+\w+', r'class\s+\w+', r'def\s+\w+',
+            r'import\s+\w+', r'\/\/.*', r'\/\*.*\*\/', r'\bcode\b',
+        ]
+        
+        code_matches = sum(
+            1 for pattern in code_indicators
+            if re.search(pattern, text, re.IGNORECASE)
+        )
+        complexity_score += min(code_matches / len(code_indicators), 1.0) * 0.2
+        
+        # 4. Multi-turn conversation (10% weight)
+        # More messages can indicate more complex context
+        message_count = len(messages)
+        conversation_score = min(message_count / 10, 1.0)
+        complexity_score += conversation_score * 0.1
+        
+        # 5. Mathematical content (10% weight)
+        math_indicators = [
+            r'\d+\s*[+\-*/]\s*\d+', r'\bequation\b', r'\bformula\b',
+            r'\bcalculus\b', r'\balgebra\b', r'\bintegral\b',
+        ]
+        
+        math_matches = sum(
+            1 for pattern in math_indicators
+            if re.search(pattern, text, re.IGNORECASE)
+        )
+        complexity_score += min(math_matches / len(math_indicators), 1.0) * 0.1
+        
+        return min(complexity_score, 1.0)
+    
+    def _filter_candidates(
+        self,
+        required_capabilities: Optional[List[str]],
+        max_cost_per_1k: Optional[float],
+        max_latency_ms: Optional[float],
+        min_context_window: Optional[int],
+    ) -> Dict[str, ModelMetadata]:
+        """Filter models based on requirements."""
+        candidates = {}
+        
+        for key, meta in self.model_metadata.items():
+            # Check capabilities
+            if required_capabilities:
+                if not all(cap in meta.capabilities for cap in required_capabilities):
+                    continue
+            
+            # Check cost constraint
+            if max_cost_per_1k:
+                avg_cost_per_1k = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 1000
+                if avg_cost_per_1k > max_cost_per_1k:
+                    continue
+            
+            # Check latency constraint
+            if max_latency_ms and meta.avg_latency_ms > max_latency_ms:
+                continue
+            
+            # Check context window constraint
+            if min_context_window and meta.context_window < min_context_window:
+                continue
+            
+            candidates[key] = meta
+        
+        # Prioritize preferred providers
+        if self.preferred_providers:
+            prioritized = {}
+            for key, meta in candidates.items():
+                if meta.provider in self.preferred_providers:
+                    prioritized[key] = meta
+            if prioritized:
+                return prioritized
+        
+        return candidates
+    
+    def _select_by_cost(self, candidates: Dict[str, ModelMetadata]) -> str:
+        """Select cheapest model."""
+        def cost_score(meta: ModelMetadata) -> float:
+            # Average cost per 1M tokens (input + output weighted equally)
+            return (meta.cost_per_1m_input + meta.cost_per_1m_output) / 2
+        
+        return min(candidates.keys(), key=lambda k: cost_score(candidates[k]))
+    
+    def _select_by_quality(
+        self,
+        candidates: Dict[str, ModelMetadata],
+        complexity: float
+    ) -> str:
+        """
+        Select highest quality model.
+        
+        For high complexity tasks, prioritize models with reasoning capabilities.
+        """
+        def quality_score(meta: ModelMetadata) -> float:
+            base_score = meta.quality_score
+            
+            # Boost reasoning models for complex tasks
+            if complexity > 0.6 and meta.reasoning_model:
+                base_score += 0.05
+            
+            return base_score
+        
+        return max(candidates.keys(), key=lambda k: quality_score(candidates[k]))
+    
+    def _select_by_latency(self, candidates: Dict[str, ModelMetadata]) -> str:
+        """Select fastest model."""
+        return min(candidates.keys(), key=lambda k: candidates[k].avg_latency_ms)
+    
+    def _select_hybrid(
+        self,
+        candidates: Dict[str, ModelMetadata],
+        complexity: float
+    ) -> str:
+        """
+        Hybrid selection balancing cost, quality, and latency.
+        
+        Adjusts weights based on task complexity:
+        - Low complexity: Prioritize cost (60%) and speed (30%), quality (10%)
+        - High complexity: Prioritize quality (60%) and cost (30%), speed (10%)
+        """
+        scores = {}
+        
+        # Adjust weights based on complexity
+        quality_weight = 0.1 + (complexity * 0.5)  # 0.1 to 0.6
+        cost_weight = 0.6 - (complexity * 0.3)     # 0.6 to 0.3
+        latency_weight = 0.3 - (complexity * 0.2)  # 0.3 to 0.1
+        
+        for key, meta in candidates.items():
+            # Quality score (0-1, higher is better)
+            quality_score = meta.quality_score
+            
+            # Cost score (0-1, lower cost is better)
+            # Normalize to typical range: $0.001 to $0.050 per 1k tokens
+            avg_cost_per_1k = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 1000
+            cost_score = max(0, 1 - (avg_cost_per_1k / 0.050))
+            
+            # Latency score (0-1, lower latency is better)
+            # Normalize to typical range: 100ms to 10000ms
+            latency_score = max(0, 1 - (meta.avg_latency_ms / 10000))
+            
+            # Calculate weighted score
+            scores[key] = (
+                quality_weight * quality_score +
+                cost_weight * cost_score +
+                latency_weight * latency_score
+            )
+        
+        return max(scores, key=scores.get)
+    
+    def get_model_info(self, provider: str, model: str) -> Optional[ModelMetadata]:
+        """Get metadata for a specific model."""
+        key = f"{provider}/{model}"
+        return self.model_metadata.get(key)
+    
+    def route_for_extraction(
+        self,
+        file_type: FileType,
+        extraction_mode: str = "schema",
+        max_cost_per_1k_tokens: Optional[float] = None,
+    ) -> Tuple[str, str]:
+        """
+        Select the best model for file extraction tasks.
+        
+        Prioritizes quality over cost/latency for extraction tasks,
+        with task-specific optimizations.
+        
+        Args:
+            file_type: Type of file being extracted
+            extraction_mode: Mode of extraction ("schema", "errors", "structure", "summary")
+            max_cost_per_1k_tokens: Optional cost constraint
+        
+        Returns:
+            Tuple of (provider, model) optimized for extraction
+        """
+        # Define extraction-specific quality weights
+        if extraction_mode == "schema":
+            quality_weight = 0.90
+            cost_weight = 0.10
+        elif extraction_mode == "errors":
+            quality_weight = 0.80
+            cost_weight = 0.20
+        elif extraction_mode == "structure":
+            quality_weight = 0.85
+            cost_weight = 0.15
+        else:  # summary
+            quality_weight = 0.70
+            cost_weight = 0.30
+        
+        # Filter candidates
+        candidates = {}
+        for key, meta in self.model_metadata.items():
+            # Apply cost constraint if specified
+            if max_cost_per_1k_tokens:
+                avg_cost_per_1k = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 1000
+                if avg_cost_per_1k > max_cost_per_1k_tokens:
+                    continue
+            
+            candidates[key] = meta
+        
+        if not candidates:
+            raise ValueError(
+                "No models meet cost constraints. "
+                "Consider increasing max_cost_per_1k_tokens."
+            )
+        
+        # Score candidates with extraction-focused weights
+        scores = {}
+        for key, meta in candidates.items():
+            # Quality score (0-1, higher is better)
+            quality_score = meta.quality_score
+            
+            # Boost for reasoning models on error extraction
+            if extraction_mode == "errors" and meta.reasoning_model:
+                quality_score += 0.05
+            
+            # Cost score (0-1, lower cost is better)
+            avg_cost_per_1k = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 1000
+            cost_score = max(0, 1 - (avg_cost_per_1k / 0.050))
+            
+            # Calculate weighted score (no latency for extraction)
+            scores[key] = (
+                quality_weight * quality_score +
+                cost_weight * cost_score
+            )
+        
+        # Select highest scoring model
+        best_key = max(scores, key=scores.get)
+        metadata = self.model_metadata[best_key]
+        return metadata.provider, metadata.model
+    
+    def get_fallback_chain(
+        self,
+        messages: List[Message],
+        count: int = 3,
+        required_capabilities: Optional[List[str]] = None,
+        max_cost_per_1k_tokens: Optional[float] = None,
+        max_latency_ms: Optional[float] = None,
+        min_context_window: Optional[int] = None,
+        exclude_models: Optional[List[str]] = None,
+    ) -> List[Tuple[str, str]]:
+        """
+        Get a ranked list of fallback models for resilient routing.
+        
+        Returns models ranked by the current strategy, allowing automatic
+        fallback if the primary model fails.
+        
+        Args:
+            messages: Conversation messages (used for complexity analysis)
+            count: Number of fallback models to return (default: 3)
+            required_capabilities: Required model capabilities
+            max_cost_per_1k_tokens: Maximum acceptable cost
+            max_latency_ms: Maximum acceptable latency
+            min_context_window: Minimum required context window
+            exclude_models: Models to exclude from results
+        
+        Returns:
+            List of (provider, model) tuples ranked by strategy
+        
+        Example:
+            >>> router = Router(strategy=RoutingStrategy.HYBRID)
+            >>> fallbacks = router.get_fallback_chain(messages, count=3)
+            >>> # Returns: [("openai", "gpt-4o"), ("anthropic", "claude-3-5-sonnet"), ...]
+        """
+        # Analyze complexity for scoring
+        complexity = self._analyze_complexity(messages)
+        
+        # Filter candidates
+        candidates = self._filter_candidates(
+            required_capabilities,
+            max_cost_per_1k_tokens,
+            max_latency_ms,
+            min_context_window,
+        )
+        
+        # Exclude specified models
+        if exclude_models:
+            candidates = {
+                k: v for k, v in candidates.items()
+                if v.model not in exclude_models
+            }
+        
+        if not candidates:
+            return []
+        
+        # Score all candidates based on strategy
+        scores = self._score_candidates(candidates, complexity)
+        
+        # Sort by score (descending) and return top N
+        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+        
+        result = []
+        for key, _ in ranked[:count]:
+            meta = self.model_metadata[key]
+            result.append((meta.provider, meta.model))
+        
+        return result
+    
+    def _score_candidates(
+        self,
+        candidates: Dict[str, ModelMetadata],
+        complexity: float,
+    ) -> Dict[str, float]:
+        """
+        Score all candidates based on current routing strategy.
+        
+        Args:
+            candidates: Filtered model candidates
+            complexity: Task complexity score (0.0 - 1.0)
+        
+        Returns:
+            Dictionary mapping model keys to scores
+        """
+        scores = {}
+        
+        for key, meta in candidates.items():
+            if self.strategy == RoutingStrategy.COST:
+                # Lower cost = higher score
+                avg_cost = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 2
+                scores[key] = max(0, 1 - (avg_cost / 100))  # Normalize to $100/1M
+                
+            elif self.strategy == RoutingStrategy.QUALITY:
+                scores[key] = meta.quality_score
+                if complexity > 0.6 and meta.reasoning_model:
+                    scores[key] += 0.05
+                    
+            elif self.strategy == RoutingStrategy.LATENCY:
+                # Lower latency = higher score
+                scores[key] = max(0, 1 - (meta.avg_latency_ms / 10000))
+                
+            else:  # HYBRID
+                # Dynamic weights based on complexity
+                quality_weight = 0.1 + (complexity * 0.5)
+                cost_weight = 0.6 - (complexity * 0.3)
+                latency_weight = 0.3 - (complexity * 0.2)
+                
+                quality_score = meta.quality_score
+                avg_cost = (meta.cost_per_1m_input + meta.cost_per_1m_output) / 1000
+                cost_score = max(0, 1 - (avg_cost / 0.050))
+                latency_score = max(0, 1 - (meta.avg_latency_ms / 10000))
+                
+                scores[key] = (
+                    quality_weight * quality_score +
+                    cost_weight * cost_score +
+                    latency_weight * latency_score
+                )
+        
+        return scores
+    
+    def list_models(
+        self,
+        required_capabilities: Optional[List[str]] = None,
+    ) -> List[Tuple[str, str, ModelMetadata]]:
+        """
+        List all available models with their metadata.
+        
+        Args:
+            required_capabilities: Filter by required capabilities
+        
+        Returns:
+            List of (provider, model, metadata) tuples
+        """
+        results = []
+        
+        for key, meta in self.model_metadata.items():
+            # Check capabilities
+            if required_capabilities:
+                if not all(cap in meta.capabilities for cap in required_capabilities):
+                    continue
+            
+            results.append((meta.provider, meta.model, meta))
+        
+        return results