Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/image.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -9,12 +9,19 @@ license: mit
 short_description: RL training env — natural language to constraint AST
 base_path: /web
 ---
+
 [Space UI & Interface](https://huggingface.co/spaces/DecentSanage/constraint-env) 
 
 # Constraint Environment
 
 This is the environment for training LLMs to learn a specific DSL made for time table scheduling. Model can then directly output constraints from natural language. Why this is needed? Usually time table generation is an NP hard problem, for humans it could take weeks to generate a conflict free time table. To solve this problem, tools are created to generate them in reasonable time. One example of those tools is CP SAT. Users can write the hardcoded constraints and the solver will generate a time table based on those constraints. Well, what happens when you want to add new constraints? Yes, you have to directly change the code. What if there is a way to directly define constraints in natural language and the solver understands that automatically? That's what we have tried to do with this project. LLM might not be good at scheduling time tables which have dozens of constraints but what it is good at is understanding natural language. For the specific purpose of defining constraints for university time tables a DSL was created whose specification is as follows: 
 
+### Reviewer Quick Links
+* **Hugging Face Space:** https://huggingface.co/spaces/DecentSanage/constraint-env
+* **Playground UI:** https://decentSanage-constraint-env.hf.space/web
+* **Health Endpoint:** https://decentSanage-constraint-env.hf.space/health
+* **OpenAPI Schema:** https://decentSanage-constraint-env.hf.space/openapi.json
+
 ```
 program         ::= { constraint }
 
@@ -64,9 +71,14 @@ identifier      ::= letter { letter | digit | "" }
 number          ::= digit { digit }
 
 ```
-
 The model outputs a json which follows the above format which can directly be converted into CP-SAT constraints. We have also included `generator.py` which implements the DSL compiler engine directly into the timetable matrix grid natively!
 
+## System Workflow: The Interactive Compiler Loop
+
+This environment operates as a multi-step interactive compiler. The agent submits an AST, receives deterministic feedback from the OpenEnv server, and iteratively debugs its logic to maximize its reward.
+
+<img width="1440" height="3008" alt="image" src="https://github.com/user-attachments/assets/3e469652-08f6-440a-96db-0555067a0af0" />
+
 ## Action and Observation
 The dataset for the training operates using dynamic deep AST JSON nodes: 
 
@@ -199,23 +211,43 @@ INFO:     10.16.33.124:5187 - "GET /web HTTP/1.1" 307 Temporary Redirect
 INFO:     10.16.24.44:32462 - "GET /web/ HTTP/1.1" 200 OK
 ```
 
+## Custom Web UI
+
+When `ENABLE_WEB_INTERFACE=true` the server mounts a **tabbed Gradio interface** at `/web`:
+
+| Tab | Description |
+|-----|-------------|
+| **Playground** | Default OpenEnv UI with Reset / Step / Get State controls |
+| **Constraint Compiler** | Our custom tab — task selector, full AST code editor, sample loaders, node-structure reference, and a real-time compiler chatbot |
+
+### Trying the Compiler tab
+1. Select a difficulty (`easy / medium / hard`) and click **Reset / Load Task**.
+2. The prompt appears and the editor is pre-filled with a correct sample AST.
+3. Edit the AST and click **▶ Submit to Compiler** — the chatbot shows the reward, error code, and exact compiler message.
+4. Use **📚 Load Sample ASTs** to reload any of the 3 canonical examples instantly.
+
+> The custom tab is implemented in `server/gradio_ui.py` via the `gradio_builder` extension point provided by OpenEnv core.
+
 ## Project Structure
 
 ```
 constraint_env/
-├── .dockerignore         # Docker build exclusions
-├── __init__.py            # Module exports
-├── README.md              # This file
-├── openenv.yaml           # OpenEnv manifest
-├── pyproject.toml         # Project metadata and dependencies
-├── uv.lock                # Locked dependencies (generated)
-├── client.py              # ConstraintEnv client
-├── generator.py           # Native Time Table generator for target compiling
-├── models.py              # Action and Observation models
+├── .dockerignore              # Docker build exclusions
+├── __init__.py                # Module exports
+├── README.md                  # This file
+├── openenv.yaml               # OpenEnv manifest
+├── pyproject.toml             # Project metadata and dependencies
+├── uv.lock                    # Locked dependencies (generated)
+├── client.py                  # ConstraintEnv WebSocket client
+├── generator.py               # DSL → timetable matrix compiler
+├── models.py                  # Pydantic Action / Observation / State models
+├── inference.py               # Baseline evaluation loop (OpenEnv compatible)
+├── dataset_example.py         # Training data — 3 difficulty tiers
 └─��� server/
-    ├── __init__.py        # Server module exports
-    ├── graders.py         # Advanced threshold limits scaling for evaluations
-    ├── constraint_env_environment.py  # Core environment logic
-    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
-    └── Dockerfile         # Container image definition
+    ├── __init__.py            # Server module exports
+    ├── app.py                 # FastAPI app (HTTP + WebSocket + Gradio)
+    ├── gradio_ui.py           # Custom Gradio "Constraint Compiler" tab
+    ├── graders.py             # Reward calculation & pass/fail thresholds
+    ├── constraint_env_environment.py  # Core step/reset/validate logic
+    └── Dockerfile             # Container image definition
 ```

client.py

@@ -30,12 +30,6 @@ class ConstraintEnv(
     def _step_payload(self, action: ConstraintAction) -> Dict:
         """
         Convert ConstraintAction to JSON payload for step message.
-
-        Args:
-            action: ConstraintAction instance
-
-        Returns:
-            Dictionary representation suitable for JSON encoding
         """
         return {
             "ast_output": action.ast_output,
@@ -44,37 +38,38 @@ class ConstraintEnv(
     def _parse_result(self, payload: Dict) -> StepResult[ConstraintObservation]:
         """
         Parse server response into StepResult[ConstraintObservation].
-
-        Args:
-            payload: JSON response data from server
-
-        Returns:
-            StepResult with ConstraintObservation
         """
+        
+        if isinstance(payload, str):
+            raise ValueError(f"Server returned an error string instead of JSON: {payload}")
+            
         obs_data = payload.get("observation", {})
+        if isinstance(obs_data, str):
+            obs_data = {}
+            
         observation = ConstraintObservation(
             prompt=obs_data.get("prompt", ""),
-            info=obs_data.get("info", 0),
+            info=obs_data.get("info", {}),           # FIX: Changed from 0 to {}
             done=payload.get("done", False),
-            reward=payload.get("reward"),
+            reward=payload.get("reward", 0.01),
+            messages=obs_data.get("messages", [])    # FIX: Added the missing messages array
         )
 
         return StepResult(
             observation=observation,
-            reward=payload.get("reward"),
+            reward=payload.get("reward", 0.01),
             done=payload.get("done", False),
         )
 
     def _parse_state(self, payload: Dict) -> ConstraintState:
         """
         Parse server response into State object.
-
-        Args:
-            payload: JSON response from state request
-
-        Returns:
-            State object with episode_id and step_count
         """
+        if isinstance(payload, str):
+            raise ValueError(f"Server returned an error string instead of JSON: {payload}")
+            
         return ConstraintState(
             episode_id=payload.get("episode_id"),
-        )
+            step_count=payload.get("step_count", 0),  # FIX: Added missing step tracking
+            max_steps=payload.get("max_steps", 5)     # FIX: Added missing max step bounds
+        )

models.py

@@ -11,7 +11,9 @@ The constraint_env environment is a simple test environment that echoes back mes
 """
 
 from openenv.core.env_server.types import Action, Observation, State
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Union
+import json as _json
+from pydantic import field_validator
 
 
 class ConstraintAction(Action):
@@ -19,6 +21,16 @@ class ConstraintAction(Action):
 
     ast_output: str
 
+    @field_validator("ast_output", mode="before")
+    @classmethod
+    def _normalise(cls, v: Any) -> str:
+        """Accept dict (from Gradio UI) or str (from LLM/API), always store as JSON string."""
+        if isinstance(v, dict):
+            return _json.dumps(v)
+        if not isinstance(v, str):
+            return _json.dumps(v)
+        return v
+
 
 class ConstraintObservation(Observation):
     """Observation from the environment, user prompt and rewards"""

server/app.py

@@ -57,13 +57,22 @@ def _make_env():
     return ConstraintEnvironment(dataset=_DATASET)
 
 
+try:
+    from .gradio_ui import build_constraint_gradio_ui as _gradio_builder
+except ImportError:
+    try:
+        from constraint_env.server.gradio_ui import build_constraint_gradio_ui as _gradio_builder
+    except ImportError:
+        _gradio_builder = None
+
 # Create the app – pass the factory so create_app calls _make_env() per session.
 app = create_app(
     _make_env,
     ConstraintAction,
     ConstraintObservation,
-    env_name="constraint_env",
+    env_name="Timetable Constraint Environment (NL-to-AST)",
     max_concurrent_envs=1,
+    gradio_builder=_gradio_builder,
 )
 
 
@@ -71,15 +80,22 @@ app = create_app(
 # PWA manifest – browsers request this at root level for the web UI
 # ---------------------------------------------------------------------------
 
-from fastapi.responses import JSONResponse  # noqa: E402
+import os
+from pathlib import Path
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import JSONResponse
+
+_ASSETS_DIR = Path(__file__).parent.parent / "assets"
+if _ASSETS_DIR.exists():
+    app.mount("/assets", StaticFiles(directory=str(_ASSETS_DIR)), name="assets")
 
 
 @app.get("/manifest.json", include_in_schema=False)
 async def web_manifest():
     return JSONResponse(
         content={
-            "name": "Constraint Environment",
-            "short_name": "ConstraintEnv",
+            "name": "Timetable Constraint Environment (NL-to-AST)",
+            "short_name": "NL-to-AST",
             "description": "RL training environment: natural-language → constraint AST",
             "start_url": "/web/",
             "display": "standalone",

server/constraint_env_environment.py

@@ -107,9 +107,10 @@ class ConstraintEnvironment(Environment):
             self._difficulty = random.choice(["easy", "medium", "hard"])
 
         pool = self._dataset[self._difficulty]
-        idx  = self._indexes[self._difficulty]
+        # idx  = self._indexes[self._difficulty]
+        idx = 0
         self._current_sample = pool[idx]
-        self._indexes[self._difficulty] = (idx + 1) % len(pool)
+        # self._indexes[self._difficulty] = (idx + 1) % len(pool)
         self._state = ConstraintState(
             episode_id=str(uuid4()),
             step_count=0,
@@ -128,6 +129,8 @@ class ConstraintEnvironment(Environment):
         """
         Evaluate the agent's AST output and return a scored observation.
         """
+        if self._current_sample is None:
+            self.reset()
         self._state.step_count += 1
         info: Dict[str, Any] = {"difficulty": self._difficulty}
         messages: List[str] = []
@@ -140,20 +143,28 @@ class ConstraintEnvironment(Environment):
 
         # ── 1. Parse JSON ────────────────────────────────────────────
         try:
-            ast = json.loads(action.ast_output)
-            if isinstance(ast, str):
-                ast = json.loads(ast)
+            raw = action.ast_output
+            # The Gradio/WebUI may pass the action already parsed as a dict
+            if isinstance(raw, dict):
+                ast = raw
+            else:
+                ast = json.loads(raw)
+                # Handle double-encoded JSON (string that contains JSON)
+                if isinstance(ast, str):
+                    ast = json.loads(ast)
+            if not isinstance(ast, dict):
+                raise TypeError(f"Expected a JSON object, got {type(ast).__name__}")
             is_valid_json = True
-        except (json.JSONDecodeError, TypeError):
+        except (json.JSONDecodeError, TypeError) as exc:
             info["error"] = "invalid_json"
             messages.extend([
                 "Your last submitted AST:",
-                action.ast_output,
-                "Compiler Error: Syntax Error. Invalid JSON."
+                str(action.ast_output),
+                f"Compiler Error: Syntax Error. Invalid JSON — {exc}"
             ])
 
         # ── 2. Logic match (ignores "name") ──────────────────────────
-        if is_valid_json and "target_ast" in self._current_sample:
+        if is_valid_json and isinstance(ast, dict) and "target_ast" in self._current_sample:
             if self._logic_match(ast, self._current_sample["target_ast"]):
                 is_exact_match = True
                 info["exact_match"] = True
@@ -203,7 +214,7 @@ class ConstraintEnvironment(Environment):
     # ------------------------------------------------------------------
 
     @staticmethod
-    def _logic_match(ast: Dict[str, Any], target: Dict[str, Any]) -> bool:
+    def _logic_match(ast: Any, target: Dict[str, Any]) -> bool:
         """
         Compare two ASTs on every logically meaningful field, ignoring "name".
 
@@ -219,6 +230,10 @@ class ConstraintEnvironment(Environment):
         """
         _LOGIC_KEYS = {"type", "forall", "where", "assert", "minimize"}
 
+        # Guard: both sides must be dicts
+        if not isinstance(ast, dict) or not isinstance(target, dict):
+            return False
+
         # Collect only logic keys present in either dict
         all_keys = (set(ast.keys()) | set(target.keys())) & _LOGIC_KEYS
 

server/graders.py

@@ -56,4 +56,4 @@ def check_passed(difficulty: str, score: float) -> bool:
     }
     
     req_threshold = thresholds.get(difficulty, 0.8)
-    return score > req_threshold
+    return score >= req_threshold

server/gradio_ui.py

@@ -0,0 +1,289 @@
+"""
+Custom Gradio UI for the Constraint Environment.
+
+Provides a rich "Constraint Compiler" tab alongside the default OpenEnv Playground.
+Registered via the gradio_builder extension point in create_app().
+"""
+
+import json
+import gradio as gr
+
+# ---------------------------------------------------------------------------
+# Sample AST pre-loaded in the editor for judges to try immediately
+# ---------------------------------------------------------------------------
+
+_SAMPLE_EASY = json.dumps({
+    "type": "hard",
+    "name": "cs_department_meeting",
+    "forall": [
+        {"b": "branches"},
+        {"sub": {"subjects": "b"}},
+        {"d": "days"},
+        {"s": "slots"}
+    ],
+    "where": {
+        "operator": "AND",
+        "left": {
+            "operator": "==",
+            "left": {"name": "b"},
+            "right": "CS"
+        },
+        "right": {
+            "operator": "AND",
+            "left": {"operator": "==", "left": "d", "right": 2},
+            "right": {"operator": "==", "left": "s", "right": 3}
+        }
+    },
+    "assert": {
+        "operator": "==",
+        "left": {
+            "target": "schedule",
+            "args": [{"name": "b"}, {"name": "sub"}, "d", "s"]
+        },
+        "right": 0
+    }
+}, indent=2)
+
+_SAMPLE_MEDIUM = json.dumps({
+    "type": "hard",
+    "name": "subject_weekly_frequency",
+    "forall": [
+        {"b": "branches"},
+        {"sub": {"subjects": "b"}}
+    ],
+    "assert": {
+        "operator": "==",
+        "left": {
+            "operator": "sum",
+            "over": [{"d": "days"}, {"s": "slots"}],
+            "expression": {
+                "target": "schedule",
+                "args": [{"name": "b"}, {"name": "sub"}, "d", "s"]
+            }
+        },
+        "right": {"frequency": "sub"}
+    }
+}, indent=2)
+
+_SAMPLE_HARD = json.dumps({
+    "type": "hard",
+    "name": "no_classes_on_saturday",
+    "forall": [
+        {"b": "branches"},
+        {"sub": {"subjects": "b"}},
+        {"d": "days"},
+        {"s": "slots"}
+    ],
+    "where": {
+        "operator": "AND",
+        "left": {"operator": "==", "left": "d", "right": 5},
+        "right": {
+            "operator": "!=",
+            "left": {"type": "sub"},
+            "right": "online"
+        }
+    },
+    "assert": {
+        "operator": "==",
+        "left": {
+            "target": "schedule",
+            "args": [{"name": "b"}, {"name": "sub"}, "d", "s"]
+        },
+        "right": 0
+    }
+}, indent=2)
+
+_INDEX_MAPPINGS = """\
+**Day Index:** `Mon=0  Tue=1  Wed=2  Thu=3  Fri=4  Sat=5`
+**Slot Index:** `9:00=0  10:00=1  11:00=2  12:00=3  BREAK=4  2:00=5  3:00=6  4:00=7  5:00=8`
+"""
+
+_TASK_INFO = {
+    "easy":   ("🟢 EASY",   "The branch CS must not have classes on Wednesday and on 12:00.",  _SAMPLE_EASY),
+    "medium": ("🟡 MEDIUM", "Subjects must be equal to their defined frequency.",               _SAMPLE_MEDIUM),
+    "hard":   ("🔴 HARD",   "No classes should be scheduled on Saturday, except for online classes.", _SAMPLE_HARD),
+}
+
+_CSS = """
+#compiler-panel { border-left: 3px solid #7c3aed; padding-left: 12px; }
+.reward-good { color: #22c55e !important; font-weight: bold; }
+.reward-bad  { color: #ef4444 !important; font-weight: bold; }
+.task-badge  { font-size: 1.1em; font-weight: bold; }
+"""
+
+
+def build_constraint_gradio_ui(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md):
+    """Custom Gradio builder — returns a gr.Blocks shown in the 'Custom' tab."""
+
+    def _run_async(coro):
+        """Run an async coroutine from a sync Gradio handler."""
+        import asyncio, concurrent.futures
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        if loop and loop.is_running():
+            # We are inside an async context — run in a fresh thread with its own loop
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+                future = pool.submit(asyncio.run, coro)
+                return future.result()
+        else:
+            return asyncio.run(coro)
+
+    def _do_reset(task_id):
+        label, prompt, sample_ast = _TASK_INFO.get(task_id, _TASK_INFO["easy"])
+        try:
+            result = _run_async(web_manager.reset_environment({"task_id": task_id}))
+            obs = result if isinstance(result, dict) else {}
+            prompt_out = obs.get("observation", {}).get("prompt", prompt)
+        except Exception as e:
+            prompt_out = f"⚠️ Reset error: {e}\n\nDefault prompt: {prompt}"
+        return (
+            f"**Task:** {label}\n\n**Prompt:** {prompt_out}",
+            sample_ast,
+            "—",
+            "—",
+            "",
+            ""   # history is a plain string log
+        )
+
+    def _do_step(ast_text, history_log):
+        if not ast_text.strip():
+            return "—", "—", "⚠️ Please enter an AST JSON.", history_log
+
+        sep = "\n" + "─" * 60 + "\n"
+
+        try:
+            ast_obj = json.loads(ast_text)
+        except json.JSONDecodeError as e:
+            entry = f"[YOU]\n{ast_text[:300]}\n\n[COMPILER]\n❌ Invalid JSON: {e}"
+            return "0.01", "❌ invalid_json", str(e), history_log + sep + entry
+
+        try:
+            result = _run_async(web_manager.step_environment({"ast_output": json.dumps(ast_obj)}))
+            payload = result if isinstance(result, dict) else {}
+            reward = payload.get("reward", 0.01)
+            done = payload.get("done", False)
+            obs = payload.get("observation", {})
+            msgs = obs.get("messages", [])
+            info = obs.get("info", {})
+            error = info.get("error", "null")
+            exact = info.get("exact_match", False)
+        except Exception as e:
+            entry = f"[YOU]\n{ast_text[:300]}\n\n[COMPILER]\n⚠️ Server error: {e}"
+            return "—", "—", f"⚠️ {e}", history_log + sep + entry
+
+        reward_str = f"{reward:.3f}"
+        if exact:
+            status = "✅ exact_match"
+        elif error == "logic_mismatch":
+            status = "⚠️ logic_mismatch"
+        elif error == "bad_structure":
+            status = "❌ bad_structure"
+        else:
+            status = f"ℹ️ {error}"
+
+        compiler_msg = "\n".join(msgs) if msgs else ("✅ Correct! Episode complete." if exact else "No compiler feedback.")
+        done_badge = "  🏁 Episode Done" if done else ""
+        entry = (
+            f"[YOU]\n{ast_text[:400]}\n\n"
+            f"[COMPILER] Reward={reward_str} | {status}{done_badge}\n{compiler_msg}"
+        )
+        return reward_str, status, compiler_msg, history_log + sep + entry
+
+    with gr.Blocks(title="Constraint Compiler") as demo:
+        gr.HTML(f"<style>{_CSS}</style>")
+        gr.HTML(
+            '<div style="display:flex;align-items:center;gap:20px;margin-bottom:8px;">'
+            '<img src="/assets/image.png" style="height:110px;border-radius:10px;"/>'
+            '<div>'
+            '<h1 style="margin:0;font-size:1.6em;color:#e2e8f0;">Timetable Constraint Environment</h1>'
+            'Convert natural-language scheduling rules into a <strong>JSON AST</strong> and get instant compiler feedback.</p>'
+            '</div></div>'
+        )
+        gr.Markdown(_INDEX_MAPPINGS)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### ⚙️ Task Control")
+                task_dd = gr.Dropdown(
+                    choices=["easy", "medium", "hard"],
+                    value="easy",
+                    label="Select Difficulty",
+                    interactive=True
+                )
+                reset_btn = gr.Button("🔄 Reset / Load Task", variant="primary", size="lg")
+
+                gr.Markdown("### 📋 Active Prompt")
+                prompt_box = gr.Markdown("*Click Reset to load a task…*")
+
+                gr.Markdown("### 📊 Last Step Result")
+                with gr.Row():
+                    reward_box = gr.Textbox(label="Reward", value="—", interactive=False, scale=1)
+                    status_box = gr.Textbox(label="Status", value="—", interactive=False, scale=2)
+                compiler_feedback = gr.Textbox(
+                    label="Compiler Feedback",
+                    lines=5,
+                    interactive=False,
+                    placeholder="Compiler messages appear here…"
+                )
+
+                gr.Markdown("### 🗺️ Node Structure Reference")
+                gr.Markdown(
+                    "**Operator:** `{\"operator\": \"==\", \"left\": …, \"right\": …}`\n\n"
+                    "**Function:** `{\"target\": \"schedule\", \"args\": [{\"name\":\"b\"}, …]}`\n\n"
+                    "**Property:** `{\"frequency\": \"sub\"}` or `{\"type\": \"sub\"}`\n\n"
+                    "**Sum:** `{\"operator\": \"sum\", \"over\": [{\"d\":\"days\"}], \"expression\": …}`\n\n"
+                    "**Variable Ref:** `{\"name\": \"b\"}`"
+                )
+
+            with gr.Column(scale=2, elem_id="compiler-panel"):
+                gr.Markdown("### ✏️ AST Editor")
+                ast_editor = gr.Code(
+                    value=_SAMPLE_EASY,
+                    language="json",
+                    label="Your AST JSON",
+                    lines=30,
+                    interactive=True
+                )
+                with gr.Row():
+                    step_btn = gr.Button("▶ Submit to Compiler", variant="primary", size="lg")
+                    clear_btn = gr.Button("🗑 Clear History", size="lg")
+
+                gr.Markdown("### 💬 Compiler Conversation Log")
+                chat_box = gr.Textbox(
+                    label="Interaction Log",
+                    lines=18,
+                    interactive=False,
+                    placeholder="Submit to the compiler to see step-by-step feedback here…",
+                )
+
+        # ── Sample loaders ─────────────────────────────────────────��────
+        with gr.Accordion("📚 Load Sample ASTs", open=False):
+            with gr.Row():
+                sample_easy_btn   = gr.Button("🟢 Easy Sample",   size="sm")
+                sample_medium_btn = gr.Button("🟡 Medium Sample", size="sm")
+                sample_hard_btn   = gr.Button("🔴 Hard Sample",   size="sm")
+
+        history_state = gr.State("")
+
+        # ── Events ──────────────────────────────────────────────────────
+        reset_btn.click(
+            fn=_do_reset,
+            inputs=[task_dd],
+            outputs=[prompt_box, ast_editor, reward_box, status_box, compiler_feedback, history_state]
+        ).then(lambda h: h, inputs=[history_state], outputs=[chat_box])
+
+        step_btn.click(
+            fn=_do_step,
+            inputs=[ast_editor, history_state],
+            outputs=[reward_box, status_box, compiler_feedback, history_state]
+        ).then(lambda h: h, inputs=[history_state], outputs=[chat_box])
+
+        clear_btn.click(fn=lambda: ("", ""), outputs=[history_state, chat_box])
+
+        sample_easy_btn.click(fn=lambda: _SAMPLE_EASY,   outputs=[ast_editor])
+        sample_medium_btn.click(fn=lambda: _SAMPLE_MEDIUM, outputs=[ast_editor])
+        sample_hard_btn.click(fn=lambda: _SAMPLE_HARD,   outputs=[ast_editor])
+
+    return demo