app.py

#!/usr/bin/env python3
"""
Agent Zero — HF Spaces Native Version
Loads your actual ScottzillaSystems model weights directly via transformers.
No TGE endpoints, no LiteLLM proxy, no Docker Compose — works on any HF Space.

Models are loaded on-demand and cached. Switch between models via dropdown.
Uses @spaces.GPU for ZeroGPU compatibility on zero-a10g hardware.
"""

import os
import re
import json
import asyncio
from pathlib import Path
from typing import List, Dict, Optional, Any
from threading import Thread

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


# ─── Configuration ───────────────────────────────────────────────────────────

AVAILABLE_MODELS = {
    "cydonia-24b": {
        "repo": "ScottzillaSystems/Cydonia-24B-v4.1",
        "description": "Cydonia 24B — Mistral-based general purpose",
        "tier": "T2",
        "device_map": "auto",
        "max_new_tokens": 2048,
    },
    "qwen3.5-27b": {
        "repo": "ScottzillaSystems/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled",
        "description": "Qwen3.5 27B — Claude Opus distilled reasoning",
        "tier": "T3",
        "device_map": "auto",
        "max_new_tokens": 4096,
    },
    "qwen3.5-9b": {
        "repo": "ScottzillaSystems/Qwen3.5-9B-Chat",
        "description": "Qwen3.5 9B — Fast general purpose, daily driver",
        "tier": "T1",
        "device_map": "auto",
        "max_new_tokens": 2048,
    },
    "chatgpt5": {
        "repo": "ScottzillaSystems/ChatGPT-5-Chat",
        "description": "ChatGPT-5 494M — Ultra-fast router/classification",
        "tier": "T0",
        "device_map": "auto",
        "max_new_tokens": 1024,
    },
    "fallen-command": {
        "repo": "ScottzillaSystems/Fallen-Command-A-111B-Chat",
        "description": "Fallen Command 111B — Flagship reasoning",
        "tier": "T4",
        "device_map": "auto",
        "load_in_8bit": True,
        "max_new_tokens": 4096,
    },
}

DEFAULT_MODEL = "qwen3.5-9b"

# Global model cache (persists across requests on paid hardware)
_model_cache: Dict[str, Any] = {}
_tokenizer_cache: Dict[str, Any] = {}


# ─── Model Loading ───────────────────────────────────────────────────────────

def load_model(model_key: str):
    """Load model and tokenizer, caching in memory."""
    if model_key in _model_cache:
        return _model_cache[model_key], _tokenizer_cache[model_key]

    config = AVAILABLE_MODELS.get(model_key)
    if not config:
        raise ValueError(f"Unknown model: {model_key}. Available: {list(AVAILABLE_MODELS.keys())}")

    repo_id = config["repo"]
    print(f"[AgentZero] ⏳ Loading {model_key} from {repo_id}...")

    tokenizer = AutoTokenizer.from_pretrained(
        repo_id,
        trust_remote_code=True,
        token=os.getenv("HF_TOKEN"),
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    load_kwargs = {
        "pretrained_model_name_or_path": repo_id,
        "trust_remote_code": True,
        "token": os.getenv("HF_TOKEN"),
        "torch_dtype": torch.bfloat16,
        "device_map": config.get("device_map", "auto"),
    }
    if config.get("load_in_8bit"):
        load_kwargs["load_in_8bit"] = True

    model = AutoModelForCausalLM.from_pretrained(**load_kwargs)

    _model_cache[model_key] = model
    _tokenizer_cache[model_key] = tokenizer

    print(f"[AgentZero] ✅ {model_key} loaded successfully")
    return model, tokenizer


def unload_model(model_key: str):
    """Free GPU memory."""
    if model_key in _model_cache:
        del _model_cache[model_key]
        del _tokenizer_cache[model_key]
        torch.cuda.empty_cache()
        print(f"[AgentZero] 🔄 Unloaded {model_key}")
        return f"✅ {model_key} unloaded — memory freed"
    return f"ℹ️ {model_key} was not loaded"


def get_model_status():
    """Report which models are loaded."""
    loaded = list(_model_cache.keys())
    if not loaded:
        return "No models loaded"
    return f"Loaded: {', '.join(loaded)} | GPU memory: {torch.cuda.memory_allocated() // 1024**3 if torch.cuda.is_available() else 0}GB used"


# ─── Inference ───────────────────────────────────────────────────────────────

@spaces.GPU(duration=120)
def generate_stream(model_key: str, messages: List[Dict[str, str]], max_new_tokens: int = None, temperature: float = 0.7):
    """Stream tokens from the model."""
    model, tokenizer = load_model(model_key)
    config = AVAILABLE_MODELS[model_key]

    if max_new_tokens is None:
        max_new_tokens = config.get("max_new_tokens", 2048)

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
    )

    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True,
    )

    gen_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=0.9,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    for text in streamer:
        yield text

    thread.join()


# ─── Gradio UI ───────────────────────────────────────────────────────────────

CSS = """
.agent-zero-header { text-align: center; padding: 20px; background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); border-radius: 12px; margin-bottom: 16px; }
.agent-zero-header h1 { color: #e94560; margin: 0; font-size: 2em; }
.agent-zero-header p { color: #a0a0b0; margin: 8px 0 0 0; }
.model-info { background: #0f0f23; padding: 12px; border-radius: 8px; border-left: 4px solid #e94560; margin-bottom: 8px; }
.tier-badge { display: inline-block; padding: 2px 8px; border-radius: 4px; font-size: 0.8em; font-weight: bold; margin-left: 6px; }
.tier-T0 { background: #00d4aa; color: #000; }
.tier-T1 { background: #00a8e8; color: #000; }
.tier-T2 { background: #f7b731; color: #000; }
.tier-T3 { background: #e94560; color: #fff; }
.tier-T4 { background: #9b59b6; color: #fff; }
.status-bar { font-size: 0.85em; color: #6c6c8a; padding: 8px; background: #0f0f23; border-radius: 6px; }
"""


def create_ui():
    with gr.Blocks(css=CSS, title="Agent Zero — Native") as demo:
        with gr.Column(elem_classes="agent-zero-header"):
            gr.HTML("""
                <h1>🤖 Agent Zero</h1>
                <p>Autonomous multi-model agent — loading YOUR weights directly via transformers</p>
            """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### ⚙️ Model")
                model_dropdown = gr.Dropdown(
                    choices=list(AVAILABLE_MODELS.keys()),
                    value=DEFAULT_MODEL,
                    label="Active Model",
                )
                model_info = gr.Markdown("Select a model to see details")

                with gr.Accordion("🧠 Catalog", open=False):
                    catalog_html = "<table style='width:100%'>"
                    for k, v in AVAILABLE_MODELS.items():
                        catalog_html += f"<tr><td><b>{k}</b> <span class='tier-badge tier-{v['tier']}'>{v['tier']}</span></td><td style='font-size:0.9em'>{v['description']}</td></tr>"
                    catalog_html += "</table>"
                    gr.HTML(catalog_html)

                with gr.Accordion("🔧 Settings", open=False):
                    max_tokens_slider = gr.Slider(128, 4096, value=2048, step=128, label="Max New Tokens")
                    temperature_slider = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")

                status_bar = gr.Textbox(label="System Status", value="Ready — no models loaded", interactive=False, elem_classes="status-bar")

            with gr.Column(scale=3):
                chatbot = gr.Chatbot(label="Agent Zero", type="messages", height=550)
                with gr.Row():
                    msg_input = gr.Textbox(placeholder="Ask anything... model loads on first send", show_label=False, scale=8)
                    send_btn = gr.Button("Send", scale=1, variant="primary")
                with gr.Row():
                    clear_btn = gr.Button("🗑 Clear")
                    unload_btn = gr.Button("🔄 Unload Model")
                    status_btn = gr.Button("📊 Status")

        # ─── Callbacks ───

        def update_model_info(model_key):
            config = AVAILABLE_MODELS.get(model_key, {})
            return f"""<div class="model-info">
<b>{config.get('description', 'Unknown')}</b><br>
Tier: <span class="tier-badge tier-{config.get('tier', 'T0')}">{config.get('tier', 'T0')}</span> | 
Max tokens: {config.get('max_new_tokens', 'N/A')}<br>
<code>{config.get('repo', 'N/A')}</code>
</div>"""

        model_dropdown.change(update_model_info, inputs=model_dropdown, outputs=model_info)

        async def chat_fn(message, history, model_key, max_tok, temp):
            if not message.strip():
                yield history, "", ""

            history = history or []
            history.append({"role": "user", "content": message})
            yield history, "", f"⏳ Loading {model_key}..."

            try:
                messages = [{"role": h["role"], "content": h["content"]} for h in history]

                response_text = ""
                for chunk in generate_stream(model_key, messages, max_tok, temp):
                    response_text += chunk
                    if history and history[-1]["role"] == "assistant":
                        history[-1]["content"] = response_text
                    else:
                        history.append({"role": "assistant", "content": response_text})
                    yield history, "", get_model_status()

            except Exception as e:
                error_msg = f"❌ Error: {str(e)}\n\nTry a smaller model or check status."
                history.append({"role": "assistant", "content": error_msg})
                yield history, "", get_model_status()

        send_btn.click(
            chat_fn,
            inputs=[msg_input, chatbot, model_dropdown, max_tokens_slider, temperature_slider],
            outputs=[chatbot, msg_input, status_bar],
        )
        msg_input.submit(
            chat_fn,
            inputs=[msg_input, chatbot, model_dropdown, max_tokens_slider, temperature_slider],
            outputs=[chatbot, msg_input, status_bar],
        )

        clear_btn.click(lambda: ([], "", "Ready"), outputs=[chatbot, msg_input, status_bar])
        unload_btn.click(
            lambda m: (unload_model(m), get_model_status()),
            inputs=model_dropdown, outputs=[status_bar, status_bar],
        )
        status_btn.click(lambda: get_model_status(), outputs=status_bar)

    return demo


if __name__ == "__main__":
    demo = create_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        share=False,
    )