app.py

import os
import uuid
from enum import Enum
from html import escape
from typing import Optional

from fastapi import FastAPI, Form, Header, HTTPException
from fastapi.responses import HTMLResponse, RedirectResponse
from pydantic import BaseModel


# =============================================================================
# Configuration
# =============================================================================

BROKER_TOKEN = os.environ.get("BROKER_TOKEN")
UI_TOKEN = os.environ.get("UI_TOKEN")

if not BROKER_TOKEN:
    raise RuntimeError(
        "Missing BROKER_TOKEN. Add it in Hugging Face Space "
        "Settings → Variables and secrets → New secret."
    )

if not UI_TOKEN:
    raise RuntimeError(
        "Missing UI_TOKEN. Add it in Hugging Face Space "
        "Settings → Variables and secrets → New secret."
    )


# =============================================================================
# Safe predefined options
# =============================================================================

MODEL_OPTIONS = {
    "granite-4.1-8b": "ibm-granite/granite-4.1-8b",
    "granite-4.1-30b": "ibm-granite/granite-4.1-30b",
    "llama-3-3-70b-instruct": "meta-llama/llama-3-3-70b-instruct",
}

QUERY_LLM_COMMAND = "query_llm"


# =============================================================================
# Data model
# =============================================================================

class JobStatus(str, Enum):
    queued = "queued"
    running = "running"
    done = "done"
    failed = "failed"


class Job(BaseModel):
    id: str
    command: str
    status: JobStatus = JobStatus.queued
    result: Optional[str] = None

    # Parameters for query_llm
    model: Optional[str] = None
    gpus: Optional[int] = None
    user_text: Optional[str] = None


app = FastAPI(title="Fury Broker")

# In-memory storage. Jobs disappear if the Space restarts.
jobs: dict[str, Job] = {}


# =============================================================================
# Security helpers
# =============================================================================

def verify_broker_token(x_broker_token: Optional[str]) -> None:
    """
    Used by the Fury worker and command-line API calls.

    Header:
        X-Broker-Token: BROKER_TOKEN
    """
    if x_broker_token != BROKER_TOKEN:
        raise HTTPException(status_code=401, detail="Unauthorized")


def verify_ui_token(ui_token: Optional[str]) -> None:
    """
    Used by browser form submissions.

    Form field:
        ui_token
    """
    if ui_token != UI_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid UI token")



def validate_query_llm_args(
    model: str,
    gpus: int,
    user_text: str,
) -> None:
    if model not in MODEL_OPTIONS:
        raise HTTPException(status_code=400, detail=f"Model is not allowed: {model}")

    if gpus < 1 or gpus > 16:
        raise HTTPException(status_code=400, detail="GPUs must be between 1 and 16")

    if user_text is None:
        raise HTTPException(status_code=400, detail="Prompt is required")

    if len(user_text.strip()) == 0:
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")

    if len(user_text) > 10_000:
        raise HTTPException(
            status_code=400,
            detail="Prompt is too long; max 10,000 characters",
        )


# =============================================================================
# Health and inspection APIs
# =============================================================================

@app.get("/health")
def health():
    return {
        "status": "ok",
        "service": "fury-broker",
        "jobs_count": len(jobs),
    }


@app.get("/api/commands")
def list_commands(x_broker_token: Optional[str] = Header(default=None)):
    verify_broker_token(x_broker_token)

    return {
        "query_llm": {
            "command": QUERY_LLM_COMMAND,
            "model_options": MODEL_OPTIONS,
            "gpu_range": [1, 16],
            "default_gpus": 1,
            "default_prompt": "hello",
        },
    }


# =============================================================================
# Browser UI
# =============================================================================

@app.get("/", response_class=HTMLResponse)
def home():
    model_options_html = "\n".join(
        f'<option value="{escape(key)}">{escape(value)}</option>'
        for key, value in MODEL_OPTIONS.items()
    )

    gpu_options_html = "\n".join(
        f'<option value="{i}" {"selected" if i == 1 else ""}>{i}</option>'
        for i in range(1, 17)
    )

    rows = ""

    for job in reversed(list(jobs.values())):
        details = []

        if job.model:
            details.append(f"model={job.model}")

        if job.gpus:
            details.append(f"gpus={job.gpus}")

        if job.user_text:
            preview = job.user_text[:500]
            if len(job.user_text) > 500:
                preview += "..."
            details.append(f"prompt={preview}")

        safe_details = escape("\n".join(details))
        safe_result = escape(job.result or "")

        rows += f"""
        <tr>
            <td><code>{escape(job.id)}</code></td>
            <td>{escape(job.command)}</td>
            <td>{escape(job.status.value)}</td>
            <td><pre>{safe_details}</pre></td>
            <td><pre>{safe_result}</pre></td>
        </tr>
        """

    return f"""
    <!doctype html>
    <html>
        <head>
            <title>Fury Broker</title>
            <style>
                body {{
                    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
                    margin: 40px;
                    max-width: 1300px;
                    line-height: 1.45;
                }}
                h1 {{
                    margin-bottom: 0.2rem;
                }}
                .subtitle {{
                    color: #555;
                    margin-bottom: 2rem;
                }}
                form {{
                    margin: 1.5rem 0;
                    padding: 1rem;
                    border: 1px solid #ddd;
                    border-radius: 8px;
                    background: #fafafa;
                }}
                input, select, textarea, button {{
                    padding: 6px;
                    margin: 4px;
                }}
                textarea {{
                    width: 95%;
                    font-family: monospace;
                }}
                table {{
                    border-collapse: collapse;
                    width: 100%;
                    margin-top: 20px;
                }}
                th, td {{
                    border: 1px solid #ccc;
                    padding: 8px;
                    vertical-align: top;
                }}
                th {{
                    background: #f5f5f5;
                    text-align: left;
                }}
                pre {{
                    white-space: pre-wrap;
                    max-width: 600px;
                    max-height: 500px;
                    overflow: auto;
                    margin: 0;
                }}
                .warning {{
                    color: #8a4b00;
                    background: #fff4dd;
                    border: 1px solid #f0c36d;
                    padding: 0.8rem;
                    border-radius: 6px;
                }}
            </style>
        </head>
        <body>
            <h1>Fury Broker</h1>
            <div class="subtitle">
                Submit approved jobs from Hugging Face Spaces to the worker running on Fury.
            </div>

            <div class="warning">
                This UI submits only validated query_llm jobs.
                Fury validates the model, GPU count, and prompt again locally before execution.
            </div>

            <form method="post" action="/submit-query-llm">
                <h2>query_llm</h2>

                <div>
                    <label><strong>UI token:</strong></label>
                    <input type="password" name="ui_token" placeholder="Enter UI_TOKEN" required>
                </div>

                <div>
                    <label><strong>Model:</strong></label>
                    <select name="model">
                        {model_options_html}
                    </select>
                </div>

                <div>
                    <label><strong>Number of GPUs:</strong></label>
                    <select name="gpus">
                        {gpu_options_html}
                    </select>
                </div>

                <div>
                    <label><strong>Prompt to send to the LLM:</strong></label><br>
                    <textarea
                        name="user_text"
                        rows="8"
                        required
                    >hello</textarea>
                </div>

                <button type="submit">Submit query_llm job</button>
            </form>

            <p>
                Refresh the page after a few seconds to see updated results.
            </p>

            <h2>Jobs</h2>

            <table>
                <tr>
                    <th>ID</th>
                    <th>Command</th>
                    <th>Status</th>
                    <th>Parameters</th>
                    <th>Result</th>
                </tr>
                {rows}
            </table>
        </body>
    </html>
    """


# =============================================================================
# Browser submit endpoints
# =============================================================================

@app.post("/submit-query-llm")
def submit_query_llm_from_ui(
    model: str = Form(...),
    gpus: int = Form(1),
    user_text: str = Form("hello"),
    ui_token: str = Form(...),
):
    verify_ui_token(ui_token)

    validate_query_llm_args(
        model=model,
        gpus=gpus,
        user_text=user_text,
    )

    job_id = str(uuid.uuid4())

    jobs[job_id] = Job(
        id=job_id,
        command=QUERY_LLM_COMMAND,
        model=model,
        gpus=gpus,
        user_text=user_text,
        status=JobStatus.queued,
    )

    return RedirectResponse("/", status_code=303)


# =============================================================================
# API submit endpoints
# =============================================================================

@app.post("/api/jobs/query-llm")
def submit_query_llm_job_api(
    model: str = Form(...),
    gpus: int = Form(1),
    user_text: str = Form("hello"),
    x_broker_token: Optional[str] = Header(default=None),
):
    verify_broker_token(x_broker_token)

    validate_query_llm_args(
        model=model,
        gpus=gpus,
        user_text=user_text,
    )

    job_id = str(uuid.uuid4())

    job = Job(
        id=job_id,
        command=QUERY_LLM_COMMAND,
        model=model,
        gpus=gpus,
        user_text=user_text,
        status=JobStatus.queued,
    )

    jobs[job_id] = job
    return job


# =============================================================================
# Worker polling and result posting
# =============================================================================

@app.get("/api/next-job")
def get_next_job(x_broker_token: Optional[str] = Header(default=None)):
    verify_broker_token(x_broker_token)

    for job in jobs.values():
        if job.status == JobStatus.queued:
            job.status = JobStatus.running

            return {
                "id": job.id,
                "command": job.command,
                "model": job.model,
                "gpus": job.gpus,
                "user_text": job.user_text,
            }

    return {"id": None}


@app.post("/api/jobs/{job_id}/result")
def post_result(
    job_id: str,
    result: str = Form(...),
    success: bool = Form(...),
    x_broker_token: Optional[str] = Header(default=None),
):
    verify_broker_token(x_broker_token)

    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")

    job = jobs[job_id]
    job.result = result
    job.status = JobStatus.done if success else JobStatus.failed

    return job


@app.get("/api/jobs")
def list_jobs(x_broker_token: Optional[str] = Header(default=None)):
    verify_broker_token(x_broker_token)
    return {"jobs": list(jobs.values())}


@app.get("/api/jobs/{job_id}")
def get_job(
    job_id: str,
    x_broker_token: Optional[str] = Header(default=None),
):
    verify_broker_token(x_broker_token)

    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")

    return jobs[job_id]


@app.post("/api/clear")
def clear_jobs(x_broker_token: Optional[str] = Header(default=None)):
    verify_broker_token(x_broker_token)

    jobs.clear()

    return {
        "status": "cleared",
        "jobs_count": len(jobs),
    }