nexus_os_v2/hf_inference_client.py

"""HF Inference API client for NEXUS OS Space.

Uses the free HF Inference API tier — no GPU needed on the Space.
Users just need a HF token (which they already have for Spaces).

Docs: https://huggingface.co/docs/huggingface_hub/en/guides/inference
"""
import os
import json
import time
from typing import Optional, Dict, Any, Tuple, List
from dataclasses import dataclass


@dataclass
class InferenceResult:
    text: str
    model: str
    latency_ms: float
    tokens_generated: int = 0
    tokens_input: int = 0
    raw: Dict[str, Any] = None


class HFInferenceClient:
    """
    Client for HuggingFace Inference API.
    Falls back gracefully if API is unavailable.
    """
    
    # Models that work well on the free tier (small, fast, good quality)
    DEFAULT_MODELS = {
        "smollm2-1.7b": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
        "llama-3.2-1b": "meta-llama/Llama-3.2-1B-Instruct",
        "qwen2.5-0.5b": "Qwen/Qwen2.5-0.5B-Instruct",
        "gemma-2-2b": "google/gemma-2-2b-it",
        "phi-4": "microsoft/Phi-4-mini-instruct",
    }
    
    def __init__(self, token: Optional[str] = None):
        self.token = token or os.environ.get("HF_TOKEN", "")
        self._available = None
    
    def is_available(self) -> bool:
        """Check if Inference API is accessible."""
        if self._available is not None:
            return self._available
        if not self.token:
            self._available = False
            return False
        try:
            import urllib.request
            req = urllib.request.Request(
                "https://api-inference.huggingface.co/models/HuggingFaceTB/SmolLM2-1.7B-Instruct",
                headers={"Authorization": f"Bearer {self.token}"},
                method="GET",
            )
            with urllib.request.urlopen(req, timeout=10) as resp:
                self._available = resp.status == 200
        except Exception:
            self._available = False
        return self._available
    
    def generate(
        self,
        prompt: str,
        model: str = "HuggingFaceTB/SmolLM2-1.7B-Instruct",
        max_tokens: int = 512,
        temperature: float = 0.7,
        system: Optional[str] = None,
    ) -> InferenceResult:
        """Generate text via HF Inference API."""
        import urllib.request
        import urllib.error
        
        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")
        
        req = urllib.request.Request(
            "https://api-inference.huggingface.co/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.token}",
            },
            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 InferenceResult(
                    text=message.get("content", ""),
                    model=data.get("model", model),
                    latency_ms=elapsed,
                    tokens_generated=usage.get("completion_tokens", 0),
                    tokens_input=usage.get("prompt_tokens", 0),
                    raw=data,
                )
        except urllib.error.HTTPError as e:
            error_body = e.read().decode("utf-8")
            raise RuntimeError(f"HF Inference API error {e.code}: {error_body}")
    
    def list_models(self) -> List[str]:
        """List available default models."""
        return list(self.DEFAULT_MODELS.keys())


class MockInferenceClient:
    """Mock client for testing without API access."""
    
    def is_available(self) -> bool:
        return True
    
    def generate(
        self,
        prompt: str,
        model: str = "mock",
        max_tokens: int = 512,
        temperature: float = 0.7,
        system: Optional[str] = None,
    ) -> InferenceResult:
        return InferenceResult(
            text=f"[MOCK] This is a simulated response for: {prompt[:50]}...\n\nIn production, this would be generated by {model} via HF Inference API.",
            model=model,
            latency_ms=100.0,
            tokens_generated=20,
            tokens_input=10,
        )
    
    def list_models(self) -> List[str]:
        return ["mock-model"]


class OllamaRelayClient:
    """
    Connects to user's local Ollama via relay URL.
    The user exposes their local Ollama via ngrok, localtunnel, or Cloudflare Tunnel.
    Set OLLAMA_RELAY_URL env var to the public tunnel endpoint.
    """
    def __init__(self, relay_url: Optional[str] = None):
        self.relay_url = relay_url or os.environ.get("OLLAMA_RELAY_URL", "")
        if not self.relay_url:
            self.relay_url = "http://localhost:11434"
        self.relay_url = self.relay_url.rstrip("/")
        self._available_models: List[str] = []
    
    def is_connected(self) -> bool:
        try:
            import urllib.request
            req = urllib.request.Request(
                f"{self.relay_url}/api/tags",
                headers={"Content-Type": "application/json"},
                method="GET",
            )
            with urllib.request.urlopen(req, timeout=10) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                self._available_models = [m.get("name", m.get("model", "")) for m in data.get("models", [])]
                return True
        except Exception:
            return False
    
    def generate(self, model_tag: str, prompt: str, system: Optional[str] = None,
                 temperature: float = 0.7, max_tokens: int = 2048, stream: bool = False):
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        payload = json.dumps({"model": model_tag, "messages": messages, "stream": stream,
                              "options": {"temperature": temperature, "num_predict": max_tokens}}).encode("utf-8")
        req = urllib.request.Request(f"{self.relay_url}/api/chat", data=payload,
                                     headers={"Content-Type": "application/json"}, method="POST")
        t0 = time.time()
        with urllib.request.urlopen(req, timeout=300) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            elapsed = (time.time() - t0) * 1000
            text = data.get("message", {}).get("content", "") if "message" in data else data.get("response", "")
            return text, {"model": data.get("model", model_tag), "latency_ms": elapsed}
    
    def list_models(self) -> List[str]:
        if not self._available_models:
            self.is_connected()
        return self._available_models