app.py

"""
NEXUS OS v5.0 — Provider Control Center

A multi-provider LLM management dashboard.
Tabs:
  1. Provider Manager — enter API keys, check health, see available models
  2. Side-by-Side Arena — compare 2 providers on same prompt
  3. Experiment Log — track runs, sort by latency/cost
  4. Pinecone Chat — talk to pineosman2 assistant
  5. Model Registry — browse 37+ models

Self-contained. Only dependency: gradio.
"""
import os
import sys
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

try:
    import gradio as gr
    GRADIO_AVAILABLE = True
except ImportError:
    GRADIO_AVAILABLE = False


# ═══════════════════════════════════════════════════════════════
# PROVIDERS
# ═══════════════════════════════════════════════════════════════
class Provider:
    def __init__(self, name, display_name, domain, key_env):
        self.name = name
        self.display_name = display_name
        self.domain = domain
        self.key_env = key_env

PROVIDERS = [
    Provider("hf_router", "HF Inference Providers", "router.huggingface.co", "HF_TOKEN"),
    Provider("groq", "Groq", "api.groq.com", "GROQ_API_KEY"),
    Provider("deepseek", "DeepSeek", "api.deepseek.com", "DEEPSEEK_API_KEY"),
    Provider("openrouter", "OpenRouter", "openrouter.ai", "OPENROUTER_API_KEY"),
    Provider("together", "Together AI", "api.together.xyz", "TOGETHER_API_KEY"),
    Provider("kilocode", "Kilocode", "kilocode.ai", "KILOCODE_API_KEY"),
    Provider("nvidia", "NVIDIA NIM", "integrate.api.nvidia.com", "NVIDIA_API_KEY"),
]

PROVIDER_MAP = {p.name: p for p in PROVIDERS}

ENDPOINTS = {
    "hf_router": "https://router.huggingface.co/v1/chat/completions",
    "groq": "https://api.groq.com/openai/v1/chat/completions",
    "deepseek": "https://api.deepseek.com/v1/chat/completions",
    "openrouter": "https://openrouter.ai/api/v1/chat/completions",
    "together": "https://api.together.xyz/v1/chat/completions",
}

FREE_MODELS = {
    "hf_router": ["HuggingFaceTB/SmolLM2-1.7B-Instruct", "meta-llama/Llama-3.2-1B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct"],
    "groq": ["llama-3.2-1b-preview", "llama-3.2-3b-preview", "mixtral-8x7b-32768"],
    "deepseek": ["deepseek-chat", "deepseek-reasoner"],
    "openrouter": ["meta-llama/llama-3.2-1b-instruct:free"],
    "together": ["meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"],
}

# ═══════════════════════════════════════════════════════════════
# API HELPERS
# ═══════════════════════════════════════════════════════════════
def _api_call(endpoint: str, api_key: str, payload: Dict[str, Any], timeout: int = 120):
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
    if "openrouter" in endpoint:
        headers["HTTP-Referer"] = "https://huggingface.co/spaces/specimba/nexus-os-space"
        headers["X-Title"] = "NEXUS OS"
    req = urllib.request.Request(endpoint, data=body, headers=headers, method="POST")
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return True, json.loads(resp.read().decode("utf-8")), (time.time() - t0) * 1000, ""
    except urllib.error.HTTPError as e:
        err = e.read().decode("utf-8", errors="replace")[:300]
        return False, {}, (time.time() - t0) * 1000, f"HTTP {e.code}: {err}"
    except Exception as e:
        return False, {}, (time.time() - t0) * 1000, str(e)[:200]


def check_provider_health(provider_name: str, api_key: str) -> Dict[str, Any]:
    """Check provider health. Returns dict with status, latency, models, error."""
    provider = PROVIDER_MAP.get(provider_name)
    if not provider:
        return {"provider": provider_name, "status": "unknown", "latency_ms": 0, "error": "Unknown provider", "models": []}
    
    if not api_key:
        return {"provider": provider.display_name, "status": "no_key", "latency_ms": 0, 
                "error": "", "models": FREE_MODELS.get(provider_name, [])}
    
    endpoint = ENDPOINTS.get(provider_name)
    if not endpoint:
        return {"provider": provider.display_name, "status": "no_endpoint", "latency_ms": 0,
                "error": "No endpoint configured", "models": []}
    
    models = FREE_MODELS.get(provider_name, [])
    model = models[0] if models else ""
    if not model:
        return {"provider": provider.display_name, "status": "no_models", "latency_ms": 0,
                "error": "No free models configured", "models": []}
    
    payload = {"model": model, "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 5, "temperature": 0.1}
    success, data, latency, error = _api_call(endpoint, api_key, payload, timeout=15)
    
    if success and data.get("choices"):
        return {"provider": provider.display_name, "status": "online", "latency_ms": round(latency, 1),
                "error": "", "models": models}
    elif "429" in error or "rate limit" in error.lower():
        return {"provider": provider.display_name, "status": "rate_limited", "latency_ms": round(latency, 1),
                "error": error[:100], "models": models}
    else:
        return {"provider": provider.display_name, "status": "offline", "latency_ms": round(latency, 1),
                "error": error[:100], "models": models}


def generate_with_provider(provider_name: str, api_key: str, model: str, prompt: str,
                            system: Optional[str] = None, max_tokens: int = 512, temperature: float = 0.7) -> Dict[str, Any]:
    """Generate text with a provider. Returns dict with text, latency, tokens, error."""
    endpoint = ENDPOINTS.get(provider_name)
    if not endpoint:
        return {"text": "", "latency_ms": 0, "tokens_input": 0, "tokens_output": 0, "error": "No endpoint"}
    
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    
    payload = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature}
    success, data, latency, error = _api_call(endpoint, api_key, payload)
    
    if not success:
        return {"text": "", "latency_ms": round(latency, 1), "tokens_input": 0, "tokens_output": 0, "error": error}
    
    choice = data.get("choices", [{}])[0]
    message = choice.get("message", {})
    usage = data.get("usage", {})
    
    return {
        "text": message.get("content", ""),
        "latency_ms": round(latency, 1),
        "tokens_input": usage.get("prompt_tokens", 0),
        "tokens_output": usage.get("completion_tokens", 0),
        "error": "",
    }


# ═══════════════════════════════════════════════════════════════
# MODEL REGISTRY
# ═══════════════════════════════════════════════════════════════
REGISTRY = {
    "deepseek-r1-8b": {"name": "DeepSeek-R1 8B", "family": "deepseek", "tier": "16GB", "size_gb": 5.2, "params_b": 8.0, "caps": "reasoning, coding, long_context", "max_context": 128000, "temp": 0.6},
    "qwen2.5-coder-7b": {"name": "Qwen 2.5 Coder 7B", "family": "qwen", "tier": "16GB", "size_gb": 4.7, "params_b": 7.0, "caps": "coding, fast", "max_context": 32768, "temp": 0.3},
    "l3.1-dark-reasoning-8b": {"name": "L3.1 Dark Reasoning 8B", "family": "llama", "tier": "16GB", "size_gb": 5.7, "params_b": 8.0, "caps": "reasoning, coding", "max_context": 32768, "temp": 0.7},
    "omega-evolution-9b": {"name": "Omega Evolution 9B", "family": "omega", "tier": "16GB", "size_gb": 6.6, "params_b": 9.0, "caps": "reasoning, coding, vision", "max_context": 32768, "temp": 0.7},
    "darwin-9b-opus": {"name": "Darwin 9B Opus", "family": "darwin", "tier": "16GB", "size_gb": 6.3, "params_b": 9.0, "caps": "reasoning, coding, long_context", "max_context": 65536, "temp": 0.7},
    "qwopus-3.5-9b": {"name": "Qwopus 3.5 9B", "family": "qwopus", "tier": "16GB", "size_gb": 5.6, "params_b": 9.0, "caps": "reasoning, coding", "max_context": 32768, "temp": 0.7},
    "carnice-9b": {"name": "Carnice 9B", "family": "carnice", "tier": "16GB", "size_gb": 5.6, "params_b": 9.0, "caps": "reasoning, coding, vision", "max_context": 32768, "temp": 0.7},
    "open-search-vl-8b": {"name": "OpenSearch VL 8B", "family": "opensearch", "tier": "16GB", "size_gb": 6.6, "params_b": 8.0, "caps": "vision, reasoning, long_context", "max_context": 65536, "temp": 0.7},
    "granite-4.1-8b-abliterated": {"name": "Granite 4.1 8B Abliterated", "family": "granite", "tier": "16GB", "size_gb": 5.1, "params_b": 8.0, "caps": "reasoning, coding, long_context", "max_context": 128000, "temp": 0.7},
    "jaahas-qwen3.5-9b": {"name": "Jaahas Qwen 3.5 9B", "family": "qwen", "tier": "16GB", "size_gb": 7.4, "params_b": 9.0, "caps": "reasoning, coding, multilingual", "max_context": 32768, "temp": 0.7},
    "lfm2-12b-deckard": {"name": "LFM2 12B Deckard", "family": "lfm", "tier": "24GB", "size_gb": 5.8, "params_b": 12.0, "caps": "reasoning, coding, long_context, fast", "max_context": 128000, "temp": 0.7},
    "gemma4-e2b-opus": {"name": "Gemma4 E2B Opus", "family": "gemma", "tier": "24GB", "size_gb": 5.5, "params_b": 4.0, "caps": "reasoning, coding, long_context", "max_context": 128000, "temp": 0.7},
    "gemma4-uncensored": {"name": "Gemma 4 Uncensored", "family": "gemma", "tier": "24GB", "size_gb": 4.9, "params_b": 4.0, "caps": "reasoning, coding, vision", "max_context": 32768, "temp": 0.7},
    "gemma4-obliterated": {"name": "Gemma 4 OBLITERATED", "family": "gemma", "tier": "24GB", "size_gb": 6.3, "params_b": 4.0, "caps": "reasoning, coding, vision", "max_context": 32768, "temp": 0.7},
    "qwen3.6-27b-dflash": {"name": "Qwen 3.6 27B DFlash", "family": "qwen", "tier": "24GB", "size_gb": 1.0, "params_b": 27.0, "caps": "reasoning, coding, long_context, fast", "max_context": 128000, "temp": 0.7},
    "gemma4-31b-cloud": {"name": "Gemma4 31B Cloud", "family": "gemma", "tier": "48GB", "size_gb": 18.0, "params_b": 31.0, "caps": "reasoning, coding, vision, long_context, multilingual", "max_context": 128000, "temp": 0.7},
    "nemotron-3-nano-omni-30b": {"name": "Nemotron-3 Nano-Omni 30B", "family": "nemotron", "tier": "48GB", "size_gb": 18.0, "params_b": 30.0, "caps": "reasoning, coding, vision, long_context, safety, tool_use", "max_context": 256000, "temp": 0.6},
    "opensonnet-lite-max": {"name": "OpenSonnet-Lite-MAX", "family": "qwen", "tier": "8GB", "size_gb": 2.5, "params_b": 4.0, "caps": "reasoning, coding, fast, long_context", "max_context": 262144, "temp": 0.6},
    "deepseek-v4-pro": {"name": "DeepSeek V4 Pro", "family": "deepseek", "tier": "cloud", "size_gb": 0.0, "params_b": 671.0, "caps": "reasoning, coding, long_context, multilingual, tool_use", "max_context": 64000, "temp": 0.6},
    "qwen3-coder-next": {"name": "Qwen 3 Coder Next", "family": "qwen", "tier": "cloud", "size_gb": 0.0, "params_b": 32.0, "caps": "coding, reasoning, fast, long_context, tool_use", "max_context": 128000, "temp": 0.3},
    "kimi-k2.6": {"name": "Kimi K2.6", "family": "kimi", "tier": "cloud", "size_gb": 0.0, "params_b": 32.0, "caps": "reasoning, coding, long_context, multilingual, vision", "max_context": 200000, "temp": 0.7},
    "glm-5.1": {"name": "GLM 5.1", "family": "glm", "tier": "cloud", "size_gb": 0.0, "params_b": 32.0, "caps": "reasoning, coding, multilingual, tool_use, vision", "max_context": 128000, "temp": 0.7},
}


# ═══════════════════════════════════════════════════════════════
# EXPERIMENT LOG (session state)
# ═══════════════════════════════════════════════════════════════
experiment_log: List[Dict[str, Any]] = []


# ═══════════════════════════════════════════════════════════════
# GRADIO INTERFACE
# ═══════════════════════════════════════════════════════════════
def build_control_center():
    with gr.Blocks(title="NEXUS OS — Provider Control Center") as demo:
        
        gr.Markdown("""
        # 🔥 NEXUS OS — Provider Control Center
        
        **Manage API providers, compare models, log experiments, chat with your knowledge base.**
        """)
        
        with gr.Tabs():
            
            # ═══════════════════════════════════════════════════
            # TAB 1: Provider Manager
            # ═══════════════════════════════════════════════════
            with gr.TabItem("🔌 Provider Manager"):
                gr.Markdown("""
                ### Enter API keys and check provider health
                
                Keys are stored **in this session only** (not saved).
                """)
                
                # Provider key inputs
                key_inputs = {}
                for provider in PROVIDERS:
                    key_inputs[provider.name] = gr.Textbox(
                        label=f"{provider.display_name} API Key",
                        placeholder=f"Paste {provider.key_env} here...",
                        type="password",
                        value=os.environ.get(provider.key_env, ""),
                    )
                
                check_btn = gr.Button("🔍 Check All Providers", variant="primary")
                
                health_table = gr.DataFrame(
                    label="Provider Health Dashboard",
                    headers=["Provider", "Status", "Latency (ms)", "Free Models", "Error"],
                    interactive=False,
                )
                
                def check_all(*keys):
                    results = []
                    for provider, key in zip(PROVIDERS, keys):
                        h = check_provider_health(provider.name, key)
                        emoji = {"online": "🟢", "rate_limited": "🟡", "offline": "🔴", 
                                "no_key": "⚪", "no_endpoint": "⚫", "no_models": "⚫"}.get(h["status"], "⚪")
                        models = ", ".join(h["models"][:3]) if h["models"] else "N/A"
                        results.append({
                            "Provider": f"{emoji} {h['provider']}",
                            "Status": h["status"],
                            "Latency (ms)": str(h["latency_ms"]) if h["latency_ms"] > 0 else "N/A",
                            "Free Models": models,
                            "Error": h["error"],
                        })
                    return results
                
                check_btn.click(
                    fn=check_all,
                    inputs=list(key_inputs.values()),
                    outputs=[health_table],
                )
            
            # ═══════════════════════════════════════════════════
            # TAB 2: Side-by-Side Arena
            # ═══════════════════════════════════════════════════
            with gr.TabItem("⚔️ Side-by-Side Arena"):
                gr.Markdown("""
                ### Send the same prompt to 2 providers and compare
                """)
                
                with gr.Row():
                    arena_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Write a Python function to reverse a linked list...",
                        lines=4,
                        scale=2,
                    )
                    arena_system = gr.Textbox(
                        label="System Prompt (optional)",
                        placeholder="You are a helpful coding assistant...",
                        lines=2,
                        scale=1,
                    )
                
                with gr.Row():
                    left_provider = gr.Dropdown(
                        label="Left Provider",
                        choices=[(p.display_name, p.name) for p in PROVIDERS if p.name in ENDPOINTS],
                        value="hf_router",
                    )
                    right_provider = gr.Dropdown(
                        label="Right Provider",
                        choices=[(p.display_name, p.name) for p in PROVIDERS if p.name in ENDPOINTS],
                        value="groq",
                    )
                
                with gr.Row():
                    left_model = gr.Dropdown(label="Left Model", choices=[], value="")
                    right_model = gr.Dropdown(label="Right Model", choices=[], value="")
                
                with gr.Row():
                    arena_max_tokens = gr.Slider(minimum=64, maximum=2048, value=512, step=64, label="Max Tokens")
                    arena_temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature")
                
                arena_go = gr.Button("🚀 Run Arena", variant="primary")
                
                with gr.Row():
                    with gr.Column():
                        left_output = gr.Textbox(label="Left Output", lines=12, interactive=False)
                        left_metrics = gr.Textbox(label="Left Metrics", interactive=False)
                    with gr.Column():
                        right_output = gr.Textbox(label="Right Output", lines=12, interactive=False)
                        right_metrics = gr.Textbox(label="Right Metrics", interactive=False)
                
                # Update model lists when provider changes
                def update_models(provider_name):
                    models = FREE_MODELS.get(provider_name, [])
                    choices = [(m, m) for m in models]
                    default = models[0] if models else ""
                    return gr.Dropdown(choices=choices, value=default)
                
                left_provider.change(fn=update_models, inputs=[left_provider], outputs=[left_model])
                right_provider.change(fn=update_models, inputs=[right_provider], outputs=[right_model])
                
                # Run arena
                def run_arena(prompt, system, left_prov, right_prov, left_mod, right_mod, 
                              max_tokens, temperature, *all_keys):
                    if not prompt.strip():
                        return "Enter a prompt", "", "Enter a prompt", ""
                    
                    key_map = {p.name: k for p, k in zip(PROVIDERS, all_keys)}
                    
                    # Left
                    left_key = key_map.get(left_prov, "")
                    left_result = generate_with_provider(left_prov, left_key, left_mod, prompt, system, max_tokens, temperature) if left_key else {"text": "❌ No API key", "latency_ms": 0, "tokens_output": 0, "error": "No key"}
                    left_text = left_result["text"] if not left_result["error"] else f"❌ {left_result['error']}"
                    left_met = f"⏱️ {left_result['latency_ms']}ms | 📝 {left_result['tokens_output']} tokens | 🎲 {left_mod}"
                    
                    # Right
                    right_key = key_map.get(right_prov, "")
                    right_result = generate_with_provider(right_prov, right_key, right_mod, prompt, system, max_tokens, temperature) if right_key else {"text": "❌ No API key", "latency_ms": 0, "tokens_output": 0, "error": "No key"}
                    right_text = right_result["text"] if not right_result["error"] else f"❌ {right_result['error']}"
                    right_met = f"⏱️ {right_result['latency_ms']}ms | 📝 {right_result['tokens_output']} tokens | 🎲 {right_mod}"
                    
                    # Log to experiment log
                    global experiment_log
                    import datetime
                    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    experiment_log.append({
                        "Time": now,
                        "Provider": f"{left_prov} vs {right_prov}",
                        "Model": f"{left_mod} vs {right_mod}",
                        "Prompt (first 50 chars)": prompt[:50],
                        "Latency Left (ms)": left_result["latency_ms"],
                        "Latency Right (ms)": right_result["latency_ms"],
                        "Tokens Left": left_result["tokens_output"],
                        "Tokens Right": right_result["tokens_output"],
                        "Status": "complete" if not left_result["error"] and not right_result["error"] else "error",
                    })
                    
                    return left_text, left_met, right_text, right_met
                
                arena_go.click(
                    fn=run_arena,
                    inputs=[arena_prompt, arena_system, left_provider, right_provider, 
                            left_model, right_model, arena_max_tokens, arena_temperature] + list(key_inputs.values()),
                    outputs=[left_output, left_metrics, right_output, right_metrics],
                )
            
            # ═══════════════════════════════════════════════════
            # TAB 3: Experiment Log
            # ═══════════════════════════════════════════════════
            with gr.TabItem("📊 Experiment Log"):
                gr.Markdown("""
                ### Track and compare your arena runs
                """)
                
                log_table = gr.DataFrame(
                    label="Experiment History",
                    headers=["Time", "Provider", "Model", "Prompt (first 50 chars)", 
                            "Latency Left (ms)", "Latency Right (ms)", "Tokens Left", "Tokens Right", "Status"],
                    interactive=False,
                )
                
                refresh_log_btn = gr.Button("🔄 Refresh Log")
                clear_log_btn = gr.Button("🗑️ Clear Log")
                
                def refresh_log():
                    return experiment_log
                
                def clear_log():
                    global experiment_log
                    experiment_log = []
                    return []
                
                refresh_log_btn.click(fn=refresh_log, outputs=[log_table])
                clear_log_btn.click(fn=clear_log, outputs=[log_table])
            
            # ═══════════════════════════════════════════════════
            # TAB 4: Pinecone Chat
            # ═══════════════════════════════════════════════════
            with gr.TabItem("🌲 Pinecone Chat"):
                gr.Markdown("""
                ### Chat with your Pinecone Assistant `pineosman2`
                
                Uses Pinecone's conversational retrieval over your uploaded documents.
                """)
                
                pinecone_key = gr.Textbox(
                    label="Pinecone API Key",
                    placeholder="pcsk_...",
                    type="password",
                    value=os.environ.get("PINECONE_API_KEY", ""),
                )
                
                pinecone_chat = gr.Chatbot(label="Conversation with pineosman2", height=400)
                pinecone_msg = gr.Textbox(label="Your message", placeholder="Ask about your documents...")
                pinecone_send = gr.Button("Send", variant="primary")
                pinecone_status = gr.Textbox(label="Status", interactive=False)
                
                def pinecone_chat_fn(message, history, api_key):
                    if not api_key:
                        return history + [(message, "❌ Please enter your Pinecone API key")], "No key"
                    if not message.strip():
                        return history, "Empty message"
                    
                    # Use Pinecone REST API
                    try:
                        payload = json.dumps({
                            "messages": [{"role": "user", "content": message}],
                        }).encode("utf-8")
                        req = urllib.request.Request(
                            "https://api.pinecone.io/assistant/chat/pineosman2",
                            data=payload,
                            headers={
                                "Content-Type": "application/json",
                                "Api-Key": api_key,
                            },
                            method="POST",
                        )
                        with urllib.request.urlopen(req, timeout=60) as resp:
                            data = json.loads(resp.read().decode("utf-8"))
                            reply = data.get("message", {}).get("content", "No response")
                            return history + [(message, reply)], f"✓ Response received ({len(reply)} chars)"
                    except Exception as e:
                        return history + [(message, f"❌ Error: {str(e)[:200]}")], f"Error: {str(e)[:100]}"
                
                pinecone_send.click(
                    fn=pinecone_chat_fn,
                    inputs=[pinecone_msg, pinecone_chat, pinecone_key],
                    outputs=[pinecone_chat, pinecone_status],
                ).then(lambda: "", outputs=[pinecone_msg])
            
            # ═══════════════════════════════════════════════════
            # TAB 5: Model Registry
            # ═══════════════════════════════════════════════════
            with gr.TabItem("📋 Model Registry"):
                gr.Markdown("""
                ### Browse all models in the NEXUS OS registry
                """)
                
                registry_table = gr.DataFrame(
                    label="Registered Models",
                    headers=["ID", "Name", "Family", "Tier", "Size (GB)", "Params (B)", 
                            "Capabilities", "Context", "Temp"],
                    interactive=False,
                )
                
                def load_registry():
                    return [{
                        "ID": k,
                        "Name": v["name"],
                        "Family": v["family"],
                        "Tier": v["tier"],
                        "Size (GB)": v["size_gb"],
                        "Params (B)": v["params_b"],
                        "Capabilities": v["caps"],
                        "Context": v["max_context"],
                        "Temp": v["temp"],
                    } for k, v in REGISTRY.items()]
                
                demo.load(fn=load_registry, outputs=[registry_table])
    
    return demo


if __name__ == "__main__":
    if not GRADIO_AVAILABLE:
        print("ERROR: Gradio is required.")
        sys.exit(1)
    demo = build_control_center()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)