Add model, GPU and prompt options

app.py

@@ -9,15 +9,60 @@ 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")
+    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")
+    raise RuntimeError(
+        "Missing UI_TOKEN. Add it in Hugging Face Space "
+        "Settings → Variables and secrets → New secret."
+    )
+
+
+# =============================================================================
+# Safe, predefined options exposed in the UI
+# =============================================================================
+
+MODEL_OPTIONS = {
+    "granite-4.0-micro": "ibm-granite/granite-4.0-micro",
+    "granite-4.1-8b": "ibm-granite/granite-4.1-8b",
+    "granite-4.1-30b": "ibm-granite/granite-4.1-30b",
+    "qwen2.5-coder-32b": "Qwen/Qwen2.5-Coder-32B-Instruct",
+}
+
+# Replace these with your real second-parameter options.
+# This is intentionally a fixed allowlist.
+RUN_MODE_OPTIONS = {
+    "scan": "scan",
+    "summarize": "summarize",
+    "full": "full",
+}
+
+# Simple predefined operational commands.
+# These do not accept user free text.
+BASIC_COMMANDS = {
+    "hostname": "Run hostname",
+    "whoami": "Run whoami",
+    "pwd": "Show current directory",
+    "disk": "Show disk usage",
+    "date": "Show current date",
+    "list_home": "List home directory",
+}
+
 
+# =============================================================================
+# Data model
+# =============================================================================
 
 class JobStatus(str, Enum):
     queued = "queued"
@@ -32,32 +77,78 @@ class Job(BaseModel):
     status: JobStatus = JobStatus.queued
     result: Optional[str] = None
 
+    # Parameters for run_paper_reader
+    model: Optional[str] = None
+    run_mode: 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] = {}
 
 
-ALLOWED_COMMANDS = {
-    "hostname": "hostname",
-    "whoami": "whoami",
-    "pwd": "pwd",
-    "disk": "df -h",
-    "date": "date",
-    "list_home": "ls -lah ~ | head -50",
-}
-
+# =============================================================================
+# Security helpers
+# =============================================================================
 
 def verify_broker_token(x_broker_token: Optional[str]) -> None:
+    """
+    Used by the Fury worker.
+
+    The worker sends:
+        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 the browser UI.
+
+    The browser form sends the UI token as a form field.
+    """
     if ui_token != UI_TOKEN:
         raise HTTPException(status_code=401, detail="Invalid UI token")
 
 
+def validate_basic_command(command: str) -> None:
+    if command not in BASIC_COMMANDS:
+        raise HTTPException(status_code=400, detail=f"Command is not allowed: {command}")
+
+
+def validate_paper_reader_args(
+    model: str,
+    run_mode: 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 run_mode not in RUN_MODE_OPTIONS:
+        raise HTTPException(status_code=400, detail=f"Run mode is not allowed: {run_mode}")
+
+    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="User text is required")
+
+    if len(user_text.strip()) == 0:
+        raise HTTPException(status_code=400, detail="User text cannot be empty")
+
+    if len(user_text) > 10_000:
+        raise HTTPException(status_code=400, detail="User text is too long; max 10,000 characters")
+
+
+# =============================================================================
+# Health and API inspection
+# =============================================================================
+
 @app.get("/health")
 def health():
     return {
@@ -67,25 +158,77 @@ def health():
     }
 
 
+@app.get("/api/commands")
+def list_commands(x_broker_token: Optional[str] = Header(default=None)):
+    verify_broker_token(x_broker_token)
+
+    return {
+        "basic_commands": BASIC_COMMANDS,
+        "model_options": MODEL_OPTIONS,
+        "run_mode_options": RUN_MODE_OPTIONS,
+        "gpu_range": [1, 16],
+    }
+
+
+# =============================================================================
+# Browser UI
+# =============================================================================
+
 @app.get("/", response_class=HTMLResponse)
 def home():
+    basic_command_options_html = "\n".join(
+        f'<option value="{escape(name)}">{escape(name)} — {escape(label)}</option>'
+        for name, label in BASIC_COMMANDS.items()
+    )
+
+    model_options_html = "\n".join(
+        f'<option value="{escape(key)}">{escape(value)}</option>'
+        for key, value in MODEL_OPTIONS.items()
+    )
+
+    run_mode_options_html = "\n".join(
+        f'<option value="{escape(key)}">{escape(value)}</option>'
+        for key, value in RUN_MODE_OPTIONS.items()
+    )
+
+    gpu_options_html = "\n".join(
+        f'<option value="{i}">{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.run_mode:
+            details.append(f"run_mode={job.run_mode}")
+
+        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"user_text={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>{escape(job.result or "")}</pre></td>
+            <td><pre>{safe_details}</pre></td>
+            <td><pre>{safe_result}</pre></td>
         </tr>
         """
 
-    options = "\n".join(
-        f'<option value="{escape(name)}">{escape(name)} → {escape(cmd)}</option>'
-        for name, cmd in ALLOWED_COMMANDS.items()
-    )
-
     return f"""
     <!doctype html>
     <html>
@@ -95,9 +238,16 @@ def home():
                 body {{
                     font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
                     margin: 40px;
-                    max-width: 1200px;
+                    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;
@@ -105,10 +255,14 @@ def home():
                     border-radius: 8px;
                     background: #fafafa;
                 }}
-                input, select, button {{
+                input, select, textarea, button {{
                     padding: 6px;
                     margin: 4px;
                 }}
+                textarea {{
+                    width: 95%;
+                    font-family: monospace;
+                }}
                 table {{
                     border-collapse: collapse;
                     width: 100%;
@@ -125,23 +279,34 @@ def home():
                 }}
                 pre {{
                     white-space: pre-wrap;
-                    max-width: 800px;
-                    max-height: 400px;
+                    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>
 
-            <p>
-                Submit an approved command from Hugging Face Spaces.
-                The worker running on Fury will pick it up, execute it on Fury,
-                and return the result here.
-            </p>
+            <div class="warning">
+                This UI does not allow arbitrary shell commands. It only submits predefined
+                command names and validated parameters. The Fury worker validates again locally.
+            </div>
+
+            <form method="post" action="/submit-basic">
+                <h2>Basic Fury Command</h2>
 
-            <form method="post" action="/submit">
                 <div>
                     <label><strong>UI token:</strong></label>
                     <input type="password" name="ui_token" placeholder="Enter UI_TOKEN" required>
@@ -150,10 +315,53 @@ def home():
                 <div>
                     <label><strong>Command:</strong></label>
                     <select name="command">
-                        {options}
+                        {basic_command_options_html}
                     </select>
-                    <button type="submit">Submit job</button>
                 </div>
+
+                <button type="submit">Submit basic job</button>
+            </form>
+
+            <form method="post" action="/submit-paper-reader">
+                <h2>Run Paper Reader on Fury</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>Run mode:</strong></label>
+                    <select name="run_mode">
+                        {run_mode_options_html}
+                    </select>
+                </div>
+
+                <div>
+                    <label><strong>Number of GPUs:</strong></label>
+                    <select name="gpus">
+                        {gpu_options_html}
+                    </select>
+                </div>
+
+                <div>
+                    <label><strong>Free text to send to the LLM:</strong></label><br>
+                    <textarea
+                        name="user_text"
+                        rows="10"
+                        placeholder="Enter the text/question/instructions to send to the model..."
+                        required
+                    ></textarea>
+                </div>
+
+                <button type="submit">Submit paper reader job</button>
             </form>
 
             <p>
@@ -167,6 +375,7 @@ def home():
                     <th>ID</th>
                     <th>Command</th>
                     <th>Status</th>
+                    <th>Parameters</th>
                     <th>Result</th>
                 </tr>
                 {rows}
@@ -176,15 +385,17 @@ def home():
     """
 
 
-@app.post("/submit")
-def submit_from_ui(
+# =============================================================================
+# Browser submit endpoints
+# =============================================================================
+
+@app.post("/submit-basic")
+def submit_basic_from_ui(
     command: str = Form(...),
     ui_token: str = Form(...),
 ):
     verify_ui_token(ui_token)
-
-    if command not in ALLOWED_COMMANDS:
-        raise HTTPException(status_code=400, detail="Command is not allowed")
+    validate_basic_command(command)
 
     job_id = str(uuid.uuid4())
 
@@ -197,15 +408,49 @@ def submit_from_ui(
     return RedirectResponse("/", status_code=303)
 
 
-@app.post("/api/jobs")
-def submit_job(
+@app.post("/submit-paper-reader")
+def submit_paper_reader_from_ui(
+    model: str = Form(...),
+    run_mode: str = Form(...),
+    gpus: int = Form(...),
+    user_text: str = Form(...),
+    ui_token: str = Form(...),
+):
+    verify_ui_token(ui_token)
+
+    validate_paper_reader_args(
+        model=model,
+        run_mode=run_mode,
+        gpus=gpus,
+        user_text=user_text,
+    )
+
+    job_id = str(uuid.uuid4())
+
+    jobs[job_id] = Job(
+        id=job_id,
+        command="run_paper_reader",
+        model=model,
+        run_mode=run_mode,
+        gpus=gpus,
+        user_text=user_text,
+        status=JobStatus.queued,
+    )
+
+    return RedirectResponse("/", status_code=303)
+
+
+# =============================================================================
+# API submit endpoints
+# =============================================================================
+
+@app.post("/api/jobs/basic")
+def submit_basic_job_api(
     command: str = Form(...),
     x_broker_token: Optional[str] = Header(default=None),
 ):
     verify_broker_token(x_broker_token)
-
-    if command not in ALLOWED_COMMANDS:
-        raise HTTPException(status_code=400, detail="Command is not allowed")
+    validate_basic_command(command)
 
     job_id = str(uuid.uuid4())
 
@@ -216,10 +461,67 @@ def submit_job(
     )
 
     jobs[job_id] = job
+    return job
+
+
+@app.post("/api/jobs/paper-reader")
+def submit_paper_reader_job_api(
+    model: str = Form(...),
+    run_mode: str = Form(...),
+    gpus: int = Form(...),
+    user_text: str = Form(...),
+    x_broker_token: Optional[str] = Header(default=None),
+):
+    verify_broker_token(x_broker_token)
+
+    validate_paper_reader_args(
+        model=model,
+        run_mode=run_mode,
+        gpus=gpus,
+        user_text=user_text,
+    )
+
+    job_id = str(uuid.uuid4())
+
+    job = Job(
+        id=job_id,
+        command="run_paper_reader",
+        model=model,
+        run_mode=run_mode,
+        gpus=gpus,
+        user_text=user_text,
+        status=JobStatus.queued,
+    )
+
+    jobs[job_id] = job
+    return job
+
+
+# Backward-compatible endpoint for your previous add_job.sh.
+@app.post("/api/jobs")
+def submit_job_legacy_api(
+    command: str = Form(...),
+    x_broker_token: Optional[str] = Header(default=None),
+):
+    verify_broker_token(x_broker_token)
+    validate_basic_command(command)
 
+    job_id = str(uuid.uuid4())
+
+    job = Job(
+        id=job_id,
+        command=command,
+        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)
@@ -231,7 +533,10 @@ def get_next_job(x_broker_token: Optional[str] = Header(default=None)):
             return {
                 "id": job.id,
                 "command": job.command,
-                "shell_command": ALLOWED_COMMANDS[job.command],
+                "model": job.model,
+                "run_mode": job.run_mode,
+                "gpus": job.gpus,
+                "user_text": job.user_text,
             }
 
     return {"id": None}
@@ -262,13 +567,26 @@ def list_jobs(x_broker_token: Optional[str] = Header(default=None)):
     return {"jobs": list(jobs.values())}
 
 
-@app.get("/api/commands")
-def list_commands(x_broker_token: Optional[str] = Header(default=None)):
+@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 {
-        "commands": [
-            {"name": name, "shell_command": cmd}
-            for name, cmd in ALLOWED_COMMANDS.items()
-        ]
+        "status": "cleared",
+        "jobs_count": len(jobs),
     }