stravinsky 0.4.18__py3-none-any.whl → 0.4.66__py3-none-any.whl

mcp_bridge/routing/provider_state.py

@@ -0,0 +1,261 @@
+"""
+Provider State Tracking for Multi-Provider Routing.
+
+Tracks the availability and health of each provider (Claude, OpenAI, Gemini)
+to enable intelligent fallback when providers are rate-limited or unavailable.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ProviderState:
+    """Tracks the state of a single provider."""
+
+    name: str
+    is_available: bool = True
+    cooldown_until: float | None = None
+    error_count: int = 0
+    last_success: float | None = None
+    last_error: str | None = None
+    total_requests: int = 0
+    total_failures: int = 0
+
+    def mark_rate_limited(self, duration: int = 300, reason: str = "429 rate limit") -> None:
+        """
+        Mark provider as rate-limited with cooldown.
+
+        Args:
+            duration: Cooldown duration in seconds (default 5 minutes)
+            reason: Reason for rate limiting (for logging)
+        """
+        self.cooldown_until = time.time() + duration
+        self.is_available = False
+        self.error_count += 1
+        self.total_failures += 1
+        self.last_error = reason
+
+        logger.warning(
+            f"[ProviderState] {self.name} rate-limited: {reason}. "
+            f"Cooldown until {time.strftime('%H:%M:%S', time.localtime(self.cooldown_until))}"
+        )
+
+        # User-visible notification
+        print(
+            f"⚠️ {self.name.upper()}: Rate-limited ({reason}). Cooldown for {duration}s.",
+            file=sys.stderr,
+        )
+
+    def mark_success(self) -> None:
+        """Mark a successful request to this provider."""
+        self.last_success = time.time()
+        self.error_count = 0  # Reset consecutive error count
+        self.total_requests += 1
+
+        # If we were in cooldown but succeeded, clear it
+        if self.cooldown_until is not None:
+            logger.info(f"[ProviderState] {self.name} recovered from cooldown")
+            self.cooldown_until = None
+            self.is_available = True
+
+    def mark_error(self, error: str) -> None:
+        """Mark a non-rate-limit error."""
+        self.error_count += 1
+        self.total_failures += 1
+        self.last_error = error
+        logger.warning(f"[ProviderState] {self.name} error ({self.error_count}): {error}")
+
+    def check_availability(self) -> bool:
+        """
+        Check if provider is available (cooldown expired?).
+
+        Returns:
+            True if provider is available, False if in cooldown
+        """
+        if self.cooldown_until is None:
+            return True
+
+        if time.time() > self.cooldown_until:
+            # Cooldown expired - reset
+            logger.info(f"[ProviderState] {self.name} cooldown expired. Marking available.")
+            self.cooldown_until = None
+            self.is_available = True
+            return True
+
+        # Still in cooldown
+        remaining = self.cooldown_until - time.time()
+        logger.debug(f"[ProviderState] {self.name} still in cooldown ({remaining:.0f}s remaining)")
+        return False
+
+    def get_cooldown_remaining(self) -> float | None:
+        """Get remaining cooldown time in seconds, or None if not in cooldown."""
+        if self.cooldown_until is None:
+            return None
+        remaining = self.cooldown_until - time.time()
+        return max(0, remaining)
+
+    def reset(self) -> None:
+        """Reset provider state (clear cooldown and errors)."""
+        self.is_available = True
+        self.cooldown_until = None
+        self.error_count = 0
+        self.last_error = None
+        logger.info(f"[ProviderState] {self.name} state reset")
+
+
+class ProviderStateTracker:
+    """
+    Tracks availability of all providers with thread-safe access.
+
+    Provides fallback chain logic when primary providers are unavailable.
+    """
+
+    # Default fallback chains for each provider
+    DEFAULT_FALLBACK_CHAINS: dict[str, list[str]] = {
+        "claude": ["openai", "gemini"],
+        "openai": ["gemini", "claude"],
+        "gemini": ["openai", "claude"],
+    }
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self.providers: dict[str, ProviderState] = {
+            "claude": ProviderState("claude"),
+            "openai": ProviderState("openai"),
+            "gemini": ProviderState("gemini"),
+        }
+        self._fallback_chains = self.DEFAULT_FALLBACK_CHAINS.copy()
+
+    def get_provider(self, name: str) -> ProviderState:
+        """Get the state for a specific provider."""
+        with self._lock:
+            if name not in self.providers:
+                # Create new provider state if not exists
+                self.providers[name] = ProviderState(name)
+            return self.providers[name]
+
+    def mark_rate_limited(
+        self, provider: str, duration: int = 300, reason: str = "429 rate limit"
+    ) -> None:
+        """Mark a provider as rate-limited."""
+        with self._lock:
+            self.get_provider(provider).mark_rate_limited(duration, reason)
+
+    def mark_success(self, provider: str) -> None:
+        """Mark a successful request to a provider."""
+        with self._lock:
+            self.get_provider(provider).mark_success()
+
+    def mark_error(self, provider: str, error: str) -> None:
+        """Mark an error for a provider."""
+        with self._lock:
+            self.get_provider(provider).mark_error(error)
+
+    def is_available(self, provider: str) -> bool:
+        """Check if a provider is available."""
+        with self._lock:
+            return self.get_provider(provider).check_availability()
+
+    def get_fallback_provider(self, preferred: str) -> str:
+        """
+        Get best available provider, falling back as needed.
+
+        Args:
+            preferred: The preferred provider to use
+
+        Returns:
+            The best available provider (preferred if available, otherwise fallback)
+        """
+        with self._lock:
+            # Check preferred first
+            if self.get_provider(preferred).check_availability():
+                return preferred
+
+            # Try fallback chain
+            fallback_chain = self._fallback_chains.get(preferred, [])
+            for fallback in fallback_chain:
+                if self.get_provider(fallback).check_availability():
+                    logger.info(
+                        f"[ProviderStateTracker] Falling back from {preferred} to {fallback}"
+                    )
+                    # Notify user
+                    print(
+                        f"⚠️ {preferred.title()} unavailable → Routing to {fallback.title()}",
+                        file=sys.stderr,
+                    )
+                    return fallback
+
+            # All providers unavailable, return preferred anyway (will likely fail)
+            logger.warning(
+                f"[ProviderStateTracker] All providers unavailable. Using {preferred} anyway."
+            )
+            return preferred
+
+    def get_status(self) -> dict[str, dict[str, Any]]:
+        """Get status of all providers for dashboard/CLI."""
+        with self._lock:
+            status = {}
+            for name, state in self.providers.items():
+                state.check_availability()  # Update availability
+                cooldown_remaining = state.get_cooldown_remaining()
+                status[name] = {
+                    "available": state.is_available,
+                    "cooldown_remaining": cooldown_remaining,
+                    "error_count": state.error_count,
+                    "last_success": state.last_success,
+                    "last_error": state.last_error,
+                    "total_requests": state.total_requests,
+                    "total_failures": state.total_failures,
+                }
+            return status
+
+    def reset_all(self) -> None:
+        """Reset all provider states."""
+        with self._lock:
+            for state in self.providers.values():
+                state.reset()
+            logger.info("[ProviderStateTracker] All provider states reset")
+
+    def reset_provider(self, provider: str) -> None:
+        """Reset a specific provider's state."""
+        with self._lock:
+            if provider in self.providers:
+                self.providers[provider].reset()
+
+    def set_fallback_chain(self, provider: str, chain: list[str]) -> None:
+        """Set custom fallback chain for a provider."""
+        with self._lock:
+            self._fallback_chains[provider] = chain
+
+
+# Global singleton instance
+_provider_tracker: ProviderStateTracker | None = None
+_tracker_lock = threading.Lock()
+
+
+def get_provider_tracker() -> ProviderStateTracker:
+    """Get or create the global ProviderStateTracker instance."""
+    global _provider_tracker
+    if _provider_tracker is None:
+        with _tracker_lock:
+            if _provider_tracker is None:
+                _provider_tracker = ProviderStateTracker()
+                logger.info("[ProviderStateTracker] Created global instance")
+    return _provider_tracker
+
+
+def reset_provider_tracker() -> None:
+    """Reset the global provider tracker (mainly for testing)."""
+    global _provider_tracker
+    with _tracker_lock:
+        if _provider_tracker is not None:
+            _provider_tracker.reset_all()

mcp_bridge/routing/task_classifier.py

@@ -0,0 +1,190 @@
+"""
+Task Classifier for Intelligent Routing.
+
+Classifies incoming tasks to determine the optimal provider and model.
+Uses pattern matching and heuristics to categorize tasks.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from enum import Enum, auto
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class TaskType(Enum):
+    """Types of tasks for routing purposes."""
+
+    CODE_GENERATION = auto()  # Creating new code
+    CODE_REFACTORING = auto()  # Improving existing code structure
+    DEBUGGING = auto()  # Fixing bugs and errors
+    ARCHITECTURE = auto()  # System design and planning
+    DOCUMENTATION = auto()  # Writing docs, comments, READMEs
+    CODE_SEARCH = auto()  # Finding code patterns
+    SECURITY_REVIEW = auto()  # Security analysis
+    GENERAL = auto()  # Default fallback
+
+
+# Patterns for task classification
+TASK_PATTERNS: dict[TaskType, list[str]] = {
+    TaskType.CODE_GENERATION: [
+        r"\b(generate|create|implement|build|write|add|make|develop)\b.*\b(code|function|class|module|component|api|endpoint|feature)\b",
+        r"\b(new|fresh)\b.*\b(implementation|feature|module)\b",
+        r"\bimplement\b",
+        r"\bcreate a\b",
+    ],
+    TaskType.CODE_REFACTORING: [
+        r"\b(refactor|restructure|reorganize|clean\s*up|simplify|optimize|improve)\b",
+        r"\b(extract|inline|rename|move)\b.*\b(method|function|class|variable)\b",
+        r"\bcode\s*(cleanup|quality)\b",
+        r"\breduce\s*(complexity|duplication)\b",
+    ],
+    TaskType.DEBUGGING: [
+        r"\b(debug|fix|solve|resolve|troubleshoot|diagnose)\b",
+        r"\b(bug|error|issue|problem|failing|broken|crash)\b",
+        r"\b(not\s*working|doesn't\s*work|won't\s*work)\b",
+        r"\b(exception|traceback|stack\s*trace)\b",
+        r"\bwhy\s*(is|does|doesn't)\b.*\b(fail|error|crash)\b",
+    ],
+    TaskType.ARCHITECTURE: [
+        r"\b(architect|design|structure|pattern|system)\b",
+        r"\b(high\s*level|overall|big\s*picture)\b",
+        r"\b(scalability|maintainability|extensibility)\b",
+        r"\b(trade\s*off|decision|approach|strategy)\b",
+        r"\bhow\s*should\s*(we|i)\s*(design|structure|organize)\b",
+    ],
+    TaskType.DOCUMENTATION: [
+        r"\b(document|readme|docstring|comment|explain|describe)\b",
+        r"\b(api\s*docs|documentation|jsdoc|pydoc)\b",
+        r"\bwrite\s*(up|docs|documentation)\b",
+        r"\badd\s*comments?\b",
+    ],
+    TaskType.CODE_SEARCH: [
+        r"\b(find|search|locate|where\s*is|look\s*for)\b.*\b(code|function|class|implementation)\b",
+        r"\b(grep|ripgrep|search)\b",
+        r"\bhow\s*is\b.*\b(implemented|used|called)\b",
+        r"\bshow\s*me\b.*\b(code|implementation)\b",
+    ],
+    TaskType.SECURITY_REVIEW: [
+        r"\b(security|vulnerability|exploit|attack|injection)\b",
+        r"\b(auth|authentication|authorization|permission)\b.*\b(check|review|audit)\b",
+        r"\b(secure|harden|protect)\b",
+        r"\b(xss|csrf|sql\s*injection|rce)\b",
+    ],
+}
+
+# Default routing for each task type
+DEFAULT_TASK_ROUTING: dict[TaskType, tuple[str, str | None]] = {
+    TaskType.CODE_GENERATION: ("openai", "gpt-5-codex"),
+    TaskType.CODE_REFACTORING: ("openai", "gpt-5-codex"),
+    TaskType.DEBUGGING: ("openai", "gpt-5-codex"),
+    TaskType.ARCHITECTURE: ("openai", "gpt-5.2-medium"),  # Delphi-style
+    TaskType.DOCUMENTATION: ("gemini", "gemini-3-flash"),
+    TaskType.CODE_SEARCH: ("gemini", "gemini-3-flash"),
+    TaskType.SECURITY_REVIEW: ("claude", None),  # Keep in Claude
+    TaskType.GENERAL: ("claude", None),  # Default to Claude
+}
+
+
+def classify_task(prompt: str, context: dict[str, Any] | None = None) -> TaskType:
+    """
+    Classify a task based on prompt content and optional context.
+
+    Uses pattern matching against known task type indicators.
+
+    Args:
+        prompt: The user's prompt or request
+        context: Optional context dict with additional signals
+
+    Returns:
+        TaskType enum indicating the classification
+    """
+    if not prompt:
+        return TaskType.GENERAL
+
+    prompt_lower = prompt.lower()
+
+    # Check patterns for each task type
+    # Priority order matters - first match wins
+    priority_order = [
+        TaskType.DEBUGGING,  # Most specific - error fixing
+        TaskType.SECURITY_REVIEW,  # Security concerns
+        TaskType.CODE_REFACTORING,  # Improvement tasks
+        TaskType.ARCHITECTURE,  # Design decisions
+        TaskType.DOCUMENTATION,  # Doc writing
+        TaskType.CODE_SEARCH,  # Finding code
+        TaskType.CODE_GENERATION,  # Creating code (broad)
+    ]
+
+    for task_type in priority_order:
+        patterns = TASK_PATTERNS.get(task_type, [])
+        for pattern in patterns:
+            if re.search(pattern, prompt_lower, re.IGNORECASE):
+                logger.debug(f"[TaskClassifier] Matched {task_type.name} with pattern: {pattern}")
+                return task_type
+
+    # Check context for additional signals
+    if context:
+        # If there's an error in context, likely debugging
+        if context.get("error") or context.get("exception"):
+            return TaskType.DEBUGGING
+
+        # If there's existing code to modify
+        if context.get("existing_code") and not context.get("create_new"):
+            return TaskType.CODE_REFACTORING
+
+    return TaskType.GENERAL
+
+
+def get_routing_for_task(
+    task_type: TaskType,
+    config: dict[str, Any] | None = None,
+) -> tuple[str, str | None]:
+    """
+    Get the recommended (provider, model) for a task type.
+
+    Args:
+        task_type: The classified task type
+        config: Optional routing config override
+
+    Returns:
+        Tuple of (provider, model) where model may be None
+    """
+    # Use config if provided
+    if config:
+        task_name = task_type.name.lower()
+        if task_name in config:
+            rule = config[task_name]
+            return (rule.get("provider", "claude"), rule.get("model"))
+
+    # Fall back to defaults
+    return DEFAULT_TASK_ROUTING.get(task_type, ("claude", None))
+
+
+def classify_and_route(
+    prompt: str,
+    context: dict[str, Any] | None = None,
+    config: dict[str, Any] | None = None,
+) -> tuple[TaskType, str, str | None]:
+    """
+    Classify a task and get routing recommendation in one call.
+
+    Args:
+        prompt: The user's prompt
+        context: Optional context dict
+        config: Optional routing config override
+
+    Returns:
+        Tuple of (task_type, provider, model)
+    """
+    task_type = classify_task(prompt, context)
+    provider, model = get_routing_for_task(task_type, config)
+
+    logger.info(
+        f"[TaskClassifier] Classified as {task_type.name} → {provider}/{model or 'default'}"
+    )
+
+    return task_type, provider, model