nexus_os_v2/intelligent_router.py

"""Intelligent Multi-Provider Router for NEXUS OS Space.

Queries ALL available free API providers and picks the best one based on:
- Provider availability (health check)
- Model capability match (coding, reasoning, vision, etc.)
- Estimated latency (provider + model size)
- Cost (free tier vs paid)
- Historical quality score (GSM8K, MMLU benchmarks)

Providers supported:
- HF Inference API (free tier, serverless)
- Together AI (free tier available)
- Cerebras (free tier available)
- Groq (free tier available)
- Fireworks AI (free tier available)
- DeepSeek API (free tier available)
- Ollama relay (user's local models)

Usage:
    router = IntelligentRouter()
    result = router.route(prompt, complexity, required_capabilities)
    # result.provider, result.model, result.latency_ms, result.text
"""
import os
import json
import time
import urllib.request
import urllib.error
from typing import Optional, Dict, Any, List, Tuple
from dataclasses import dataclass, field
from enum import Enum


class Provider(Enum):
    HF_INFERENCE = "hf_inference"
    TOGETHER = "together"
    CEREBRAS = "cerebras"
    GROQ = "groq"
    FIREWORKS = "fireworks"
    DEEPSEEK = "deepseek"
    OLLAMA = "ollama"
    CLOUD = "cloud"
    MOCK = "mock"


@dataclass
class ProviderHealth:
    provider: Provider
    available: bool
    latency_ms: float = 999999.0
    error: str = ""
    models: List[str] = field(default_factory=list)


@dataclass
class RouterResult:
    text: str
    provider: Provider
    model: str
    latency_ms: float
    tokens_input: int = 0
    tokens_output: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)
    fallback_chain: List[str] = field(default_factory=list)


# Provider model mappings — what each provider offers for free
PROVIDER_MODELS = {
    Provider.HF_INFERENCE: {
        "default": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
        "coding": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
        "reasoning": "meta-llama/Llama-3.2-1B-Instruct",
        "fast": "Qwen/Qwen2.5-0.5B-Instruct",
        "vision": None,  # HF Inference API vision is limited
    },
    Provider.TOGETHER: {
        "default": "meta-llama/Llama-3.2-1B-Instruct",
        "coding": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
        "reasoning": "meta-llama/Llama-3.2-1B-Instruct",
        "fast": "meta-llama/Llama-3.2-1B-Instruct",
    },
    Provider.CEREBRAS: {
        "default": "llama-3.2-1b",
        "coding": "llama-3.2-1b",
        "reasoning": "llama-3.2-1b",
        "fast": "llama-3.2-1b",
    },
    Provider.GROQ: {
        "default": "llama-3.2-1b",
        "coding": "qwen-2.5-coder-32b",
        "reasoning": "llama-3.2-1b",
        "fast": "llama-3.2-1b",
    },
    Provider.FIREWORKS: {
        "default": "accounts/fireworks/models/llama-v3p2-1b-instruct",
        "coding": "accounts/fireworks/models/llama-v3p2-1b-instruct",
        "reasoning": "accounts/fireworks/models/llama-v3p2-1b-instruct",
        "fast": "accounts/fireworks/models/llama-v3p2-1b-instruct",
    },
    Provider.DEEPSEEK: {
        "default": "deepseek-chat",
        "coding": "deepseek-chat",
        "reasoning": "deepseek-reasoner",
        "fast": "deepseek-chat",
    },
}

# Provider API endpoints
PROVIDER_ENDPOINTS = {
    Provider.HF_INFERENCE: "https://api-inference.huggingface.co/v1/chat/completions",
    Provider.TOGETHER: "https://api.together.xyz/v1/chat/completions",
    Provider.CEREBRAS: "https://api.cerebras.ai/v1/chat/completions",
    Provider.GROQ: "https://api.groq.com/openai/v1/chat/completions",
    Provider.FIREWORKS: "https://api.fireworks.ai/inference/v1/chat/completions",
    Provider.DEEPSEEK: "https://api.deepseek.com/v1/chat/completions",
}

# Provider API key env vars
PROVIDER_KEYS = {
    Provider.HF_INFERENCE: "HF_TOKEN",
    Provider.TOGETHER: "TOGETHER_API_KEY",
    Provider.CEREBRAS: "CEREBRAS_API_KEY",
    Provider.GROQ: "GROQ_API_KEY",
    Provider.FIREWORKS: "FIREWORKS_API_KEY",
    Provider.DEEPSEEK: "DEEPSEEK_API_KEY",
}


class IntelligentRouter:
    """
    Intelligent multi-provider router for NEXUS OS.
    
    Queries all available providers in parallel, ranks by health + capability match,
    and returns the best response with full fallback chain.
    """
    
    def __init__(self):
        self._health_cache: Dict[Provider, ProviderHealth] = {}
        self._cache_time: float = 0
        self._cache_ttl: float = 60.0  # Cache health for 60 seconds
    
    def _get_api_key(self, provider: Provider) -> Optional[str]:
        """Get API key for provider from env."""
        env_var = PROVIDER_KEYS.get(provider)
        if env_var:
            return os.environ.get(env_var)
        return None
    
    def _check_provider_health(self, provider: Provider) -> ProviderHealth:
        """Check if a provider is available and measure latency."""
        api_key = self._get_api_key(provider)
        if not api_key:
            return ProviderHealth(provider=provider, available=False, error="No API key")
        
        endpoint = PROVIDER_ENDPOINTS.get(provider)
        if not endpoint:
            return ProviderHealth(provider=provider, available=False, error="No endpoint")
        
        # Quick health check: send a minimal request
        messages = [{"role": "user", "content": "Hi"}]
        payload = json.dumps({"model": "test", "messages": messages, "max_tokens": 1}).encode("utf-8")
        
        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
        req = urllib.request.Request(endpoint, data=payload, headers=headers, method="POST")
        
        t0 = time.time()
        try:
            with urllib.request.urlopen(req, timeout=15) as resp:
                _ = resp.read()
                latency = (time.time() - t0) * 1000
                return ProviderHealth(provider=provider, available=True, latency_ms=latency)
        except urllib.error.HTTPError as e:
            # 401 = bad key, 404 = model not found, 429 = rate limit, 503 = overloaded
            if e.code in (401, 403):
                return ProviderHealth(provider=provider, available=False, error=f"Invalid API key ({e.code})")
            elif e.code == 429:
                return ProviderHealth(provider=provider, available=False, error="Rate limited")
            elif e.code == 503:
                return ProviderHealth(provider=provider, available=False, error="Provider overloaded")
            else:
                return ProviderHealth(provider=provider, available=False, error=f"HTTP {e.code}")
        except Exception as e:
            return ProviderHealth(provider=provider, available=False, error=str(e)[:100])
    
    def check_all_providers(self) -> List[ProviderHealth]:
        """Check health of ALL providers. Returns sorted by latency (best first)."""
        now = time.time()
        if now - self._cache_time < self._cache_ttl and self._health_cache:
            return sorted(self._health_cache.values(), key=lambda h: (not h.available, h.latency_ms))
        
        results = []
        for provider in [Provider.HF_INFERENCE, Provider.TOGETHER, Provider.CEREBRAS, 
                          Provider.GROQ, Provider.FIREWORKS, Provider.DEEPSEEK]:
            health = self._check_provider_health(provider)
            self._health_cache[provider] = health
            results.append(health)
        
        self._cache_time = now
        # Sort: available first, then by latency
        return sorted(results, key=lambda h: (not h.available, h.latency_ms))
    
    def _generate_with_provider(self, provider: Provider, prompt: str, model: str,
                                 max_tokens: int = 512, temperature: float = 0.7,
                                 system: Optional[str] = None) -> Optional[RouterResult]:
        """Generate with a specific provider. Returns None on failure."""
        api_key = self._get_api_key(provider)
        if not api_key:
            return None
        
        endpoint = PROVIDER_ENDPOINTS.get(provider)
        if not endpoint:
            return None
        
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        payload = json.dumps({
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": False,
        }).encode("utf-8")
        
        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
        req = urllib.request.Request(endpoint, data=payload, headers=headers, method="POST")
        
        t0 = time.time()
        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                elapsed = (time.time() - t0) * 1000
                
                choice = data.get("choices", [{}])[0]
                message = choice.get("message", {})
                usage = data.get("usage", {})
                
                return RouterResult(
                    text=message.get("content", ""),
                    provider=provider,
                    model=model,
                    latency_ms=elapsed,
                    tokens_input=usage.get("prompt_tokens", 0),
                    tokens_output=usage.get("completion_tokens", 0),
                    metadata={"raw": data},
                )
        except Exception as e:
            return None
    
    def route(self, prompt: str, complexity: float = 0.5,
              required_capabilities: List[str] = None,
              max_tokens: int = 512, temperature: float = 0.7,
              system: Optional[str] = None,
              ollama_relay_url: Optional[str] = None) -> RouterResult:
        """
        Intelligent routing: try all providers in parallel, return best response.
        
        Strategy:
        1. Check health of all providers
        2. Pick the best available provider based on capability match + latency
        3. Generate
        4. If fails, try next best provider
        5. Return full fallback chain
        """
        fallback_chain = []
        
        # Check all providers
        health_results = self.check_all_providers()
        
        # Determine capability need
        capability = "default"
        if required_capabilities:
            if "coding" in required_capabilities:
                capability = "coding"
            elif "reasoning" in required_capabilities:
                capability = "reasoning"
            elif "fast" in required_capabilities:
                capability = "fast"
        
        # Try each available provider in order of health
        for health in health_results:
            if not health.available:
                fallback_chain.append(f"{health.provider.value}: unavailable ({health.error})")
                continue
            
            provider = health.provider
            model = PROVIDER_MODELS.get(provider, {}).get(capability)
            if not model:
                model = PROVIDER_MODELS.get(provider, {}).get("default", "")
            
            fallback_chain.append(f"{provider.value}: trying {model} ({health.latency_ms:.0f}ms health check)")
            
            result = self._generate_with_provider(
                provider=provider,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
                system=system,
            )
            
            if result:
                result.fallback_chain = fallback_chain
                return result
            else:
                fallback_chain.append(f"{provider.value}: generation failed")
        
        # Try Ollama relay if configured
        if ollama_relay_url:
            fallback_chain.append(f"ollama: trying relay at {ollama_relay_url}")
            try:
                from .hf_inference_client import OllamaRelayClient
                client = OllamaRelayClient(relay_url=ollama_relay_url)
                if client.is_connected():
                    text, metadata = client.generate(
                        model_tag="llama3.2:latest",  # Default Ollama model
                        prompt=prompt,
                        system=system,
                        temperature=temperature,
                        max_tokens=max_tokens,
                    )
                    return RouterResult(
                        text=text,
                        provider=Provider.OLLAMA,
                        model=metadata.get("model", "unknown"),
                        latency_ms=metadata.get("latency_ms", 0),
                        fallback_chain=fallback_chain,
                    )
            except Exception as e:
                fallback_chain.append(f"ollama: failed ({str(e)[:80]})")
        
        # All providers failed — return mock
        return RouterResult(
            text=f"[MOCK] All providers unavailable. Fallback chain:\n" + "\n".join(fallback_chain),
            provider=Provider.MOCK,
            model="mock",
            latency_ms=0.0,
            fallback_chain=fallback_chain,
        )