""" Custom Gradio UI for the Incident Response Env. Replaces the default OpenEnv playground with one that: - Shows the legal Role choices directly in the label: "Role (attacker / defender)". - Binds each action to an integer 0..N-1 and shows the integer-to-command mapping inline (with the table re-rendering when the role flips). - Adds a "Run model" button that loads the GRPO checkpoint from HF and auto-generates the next action, with live status messages. - Embeds the GRPO training notebook below the playground via an iframe. """ from __future__ import annotations import asyncio import json from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import gradio as gr from .incident_response_env_environment import ATTACKS, DEFENSES try: from .model_runner import generate_action as _model_generate_action _MODEL_AVAILABLE = True except ImportError: _MODEL_AVAILABLE = False _MODEL_EXECUTOR = ThreadPoolExecutor(max_workers=1) ROLE_CHOICES: List[str] = ["attacker", "defender"] def _action_choices(role: str) -> List[Tuple[str, int]]: """Return (label, value) pairs for the Action ID dropdown.""" items = ATTACKS if role == "attacker" else DEFENSES prefix = "ATTACK" if role == "attacker" else "DEFEND" return [(f"{i}: {prefix} {name}", i) for i, name in enumerate(items)] def _command_for(role: str, action_id: int) -> str: """Translate a (role, integer) pair into the env's expected command string.""" if role == "attacker": return f"ATTACK: {ATTACKS[action_id]}" return f"DEFEND: {DEFENSES[action_id]}" def _mapping_md(role: str) -> str: """Markdown table showing the active integer-to-command mapping.""" items = ATTACKS if role == "attacker" else DEFENSES prefix = "ATTACK" if role == "attacker" else "DEFEND" rows = "\n".join(f"| **{i}** | `{prefix}: {name}` |" for i, name in enumerate(items)) return ( f"#### Action ID mapping for role `{role}`\n" f"_The integer you pick is sent as `{prefix}: `._\n\n" "| ID | Command sent |\n" "|---:|:-------------|\n" f"{rows}\n" ) def prepare_notebook_html(notebook_path: Path, out_path: Path) -> Path: """ Convert a Jupyter notebook into a self-contained HTML file. Tries `nbconvert` first (full Jupyter rendering, including outputs and syntax highlighting). Falls back to a minimal hand-rolled renderer when nbconvert is unavailable so the server still starts cleanly. """ out_path.parent.mkdir(parents=True, exist_ok=True) try: import nbformat from nbconvert import HTMLExporter nb = nbformat.read(str(notebook_path), as_version=4) exporter = HTMLExporter(template_name="lab") exporter.embed_images = True body, _ = exporter.from_notebook_node(nb) out_path.write_text(body, encoding="utf-8") return out_path except Exception as exc: # pragma: no cover - fallback path try: nb = json.loads(notebook_path.read_text(encoding="utf-8")) except Exception as read_exc: out_path.write_text( f"

Could not load notebook: {read_exc}

", encoding="utf-8", ) return out_path parts: List[str] = [ "" "GRPO Notebook (raw)" "", "

Rendered with the built-in fallback " f"(install nbconvert for full output). Reason: {exc}

", "

GRPO Training Notebook

", ] for cell in nb.get("cells", []): src = "".join(cell.get("source", [])) if cell.get("cell_type") == "markdown": parts.append(f"

{src}

") else: escaped = ( src.replace("&", "&") .replace("<", "<") .replace(">", ">") ) parts.append(f"

{escaped}

") parts.append("") out_path.write_text("\n".join(parts), encoding="utf-8") return out_path def _strip_prompt_from_response_for_display(data: Dict[str, Any]) -> Dict[str, Any]: """ Remove the (often very long) `observation.prompt` field from the JSON shown in the UI, while keeping the server response unchanged. """ if not isinstance(data, dict): return {} out = dict(data) obs = out.get("observation") if isinstance(obs, dict) and "prompt" in obs: obs2 = dict(obs) obs2.pop("prompt", None) out["observation"] = obs2 return out def build_custom_ui( web_manager: Any, metadata: Optional[Any], quick_start_md: Optional[str], notebook_iframe_url: str, ) -> gr.Blocks: """Build the full /web Gradio app: playground + embedded GRPO notebook.""" title = getattr(metadata, "name", None) or "Incident Response Env" readme_content = getattr(metadata, "readme_content", None) if metadata else None display_title = f"OpenEnv Agentic Environment: {title}" async def reset_env(): try: data = await web_manager.reset_environment() obs = data.get("observation", {}) or {} next_role = obs.get("next_role", "attacker") health = obs.get("health", 1.0) status_text = obs.get("status", "STABLE") role_value = next_role if next_role in ROLE_CHOICES else "attacker" new_choices = _action_choices(role_value) # User requested: never show the prompt in the UI. obs_md = "" return ( gr.update(value=role_value), gr.update( choices=new_choices, value=new_choices[0][1] if new_choices else None, label=f"Action ID (0..{len(new_choices) - 1})", ), _mapping_md(role_value), obs_md, json.dumps(_strip_prompt_from_response_for_display(data), indent=2), f"Reset OK | health={health:.0%} | status={status_text} | next_role={role_value}", ) except Exception as exc: return ( gr.update(), gr.update(), gr.update(), "", "", f"Error during reset: {exc}", ) async def step_env(role: str, action_id: Optional[int]): try: # If the user clicks Step before any Reset, the underlying env session # has no episode state yet. Auto-reset once to avoid None-state errors. try: if getattr(web_manager, "episode_state", None) is not None and ( web_manager.episode_state.current_observation is None ): await web_manager.reset_environment() except Exception: # If we can't read/reset episode_state for any reason, keep going # and let the environment raise a more specific error. pass if role not in ROLE_CHOICES: return (gr.update(), gr.update(), gr.update(), "", "", "Pick a role first.") if action_id is None: return (gr.update(), gr.update(), gr.update(), "", "", "Pick an Action ID first.") ai = int(action_id) universe = ATTACKS if role == "attacker" else DEFENSES if not 0 <= ai < len(universe): return ( gr.update(), gr.update(), gr.update(), "", "", f"Invalid Action ID for `{role}` (must be 0..{len(universe) - 1}).", ) command = _command_for(role, ai) data = await web_manager.step_environment({"command": command, "role": role}) obs = data.get("observation", {}) or {} next_role = obs.get("next_role", role) health = obs.get("health", 1.0) status_text = obs.get("status", "STABLE") done = bool(obs.get("done", False)) meta = obs.get("metadata", {}) or {} role_value = next_role if next_role in ROLE_CHOICES else role new_choices = _action_choices(role_value) done_md = " Episode complete — click Reset." if done else "" breached = meta.get("breached") breach_part = f" | breached={breached}" if breached is not None else "" # User requested: never show the prompt in the UI. obs_md = "" return ( gr.update(value=role_value), gr.update( choices=new_choices, value=new_choices[0][1] if new_choices else None, label=f"Action ID (0..{len(new_choices) - 1})", ), _mapping_md(role_value), obs_md, json.dumps(_strip_prompt_from_response_for_display(data), indent=2), f"Step OK | health={health:.0%} | status={status_text}{breach_part} | next_role={next_role}.{done_md}", ) except Exception as exc: return ( gr.update(), gr.update(), gr.update(), "", "", f"Error during step: {exc}", ) def get_state_sync(): try: return json.dumps(web_manager.get_state(), indent=2) except Exception as exc: return f"Error: {exc}" def on_role_change(role: str): new_choices = _action_choices(role) return ( gr.update( choices=new_choices, value=new_choices[0][1] if new_choices else None, label=f"Action ID (0..{len(new_choices) - 1})", ), _mapping_md(role), ) async def run_model(role: str): """ Load the GRPO checkpoint for the current role and auto-generate an action, then step the environment with the result. """ _blank = (gr.update(), gr.update(), gr.update(), "", "", "") if not _MODEL_AVAILABLE: return (*_blank[:5], "transformers/torch not installed — cannot run model.") if role not in ROLE_CHOICES: return (*_blank[:5], "Select a role first.") # Ensure an episode is active try: current_obs = ( web_manager.episode_state.current_observation if web_manager.episode_state and web_manager.episode_state.current_observation else None ) except Exception: current_obs = None if current_obs is None: await web_manager.reset_environment() try: current_obs = web_manager.episode_state.current_observation except Exception: current_obs = {} prompt = (current_obs or {}).get("prompt", "") if isinstance(current_obs, dict) else "" if not prompt: return (*_blank[:5], "No prompt yet — click Reset first.") status_msgs: List[str] = [] def _status_cb(msg: str): status_msgs.append(msg) loop = asyncio.get_event_loop() try: command = await loop.run_in_executor( _MODEL_EXECUTOR, lambda: _model_generate_action(role, prompt, _status_cb), ) except Exception as exc: return (*_blank[:5], f"Model error: {exc}") all_status = " → ".join(status_msgs) if not command: return (*_blank[:5], f"Model produced no valid action. {all_status}") try: data = await web_manager.step_environment({"command": command, "role": role}) except Exception as exc: return (*_blank[:5], f"Step error after model action: {exc}") obs = data.get("observation", {}) or {} next_role = obs.get("next_role", role) health = obs.get("health", 1.0) status_text = obs.get("status", "STABLE") done = bool(obs.get("done", False)) role_value = next_role if next_role in ROLE_CHOICES else role new_choices = _action_choices(role_value) done_note = " Episode complete — click Reset." if done else "" return ( gr.update(value=role_value), gr.update(choices=new_choices, value=new_choices[0][1] if new_choices else None, label=f"Action ID (0..{len(new_choices) - 1})"), _mapping_md(role_value), "", json.dumps(_strip_prompt_from_response_for_display(data), indent=2), f"Model → {command} | health={health:.0%} | {status_text} | next={next_role}.{done_note}", ) initial_role = "attacker" initial_choices = _action_choices(initial_role) with gr.Blocks(title=display_title) as demo: # Fixed left sidebar with in-page navigation. gr.HTML( """ """ ) gr.HTML('

') with gr.Row(): with gr.Column(scale=1, elem_classes="col-left"): # Replace the default Quick Start section with the integer binding table, # so users can always see how Action IDs map to commands. mapping_md = gr.Markdown(_mapping_md(initial_role)) if readme_content: with gr.Accordion("README", open=False): gr.Markdown(readme_content) with gr.Column(scale=2, elem_classes="col-right"): gr.Markdown(f"# Playground — {title}\nClick **Reset** to start a new episode.") role_dd = gr.Dropdown( choices=ROLE_CHOICES, value=initial_role, label="Role (attacker / defender)", info="Whose turn it is. Reset will set this automatically based on the env.", allow_custom_value=False, ) action_dd = gr.Dropdown( choices=initial_choices, value=initial_choices[0][1], label=f"Action ID (0..{len(initial_choices) - 1})", info="Integer index into the role's action universe (see mapping below).", allow_custom_value=False, ) gr.Markdown("*Action ID mapping is shown on the left.*") with gr.Row(): step_btn = gr.Button("Step", variant="primary") reset_btn = gr.Button("Reset", variant="secondary") state_btn = gr.Button("Get state", variant="secondary") model_btn = gr.Button("Run model", variant="stop") status_box = gr.Textbox(label="Status", interactive=False) # Prompt is intentionally hidden; keep a placeholder for spacing # and to keep the same output wiring for click handlers. obs_md = gr.Markdown(value="") raw_json = gr.Code( label="Raw JSON response", language="json", interactive=False, ) # Notebook embed directly under the interactive playground, but full width. gr.HTML( f'''

GRPO Training Notebook

''' ) role_dd.change(on_role_change, inputs=role_dd, outputs=[action_dd, mapping_md]) reset_btn.click( fn=reset_env, outputs=[role_dd, action_dd, mapping_md, obs_md, raw_json, status_box], ) step_btn.click( fn=step_env, inputs=[role_dd, action_dd], outputs=[role_dd, action_dd, mapping_md, obs_md, raw_json, status_box], ) model_btn.click( fn=run_model, inputs=[role_dd], outputs=[role_dd, action_dd, mapping_md, obs_md, raw_json, status_box], ) state_btn.click(fn=get_state_sync, outputs=[raw_json]) return demo