Fix deprecation warnings and improve error handling

- Replace deprecated load_in_8bit with BitsAndBytesConfig
- Fix dropdown value to dynamically use first model
- Increase GPU duration to 600s for model loading
- Add better error handling for GPU task aborted errors
- Add model_choice validation

app.py

@@ -1,77 +1,73 @@
+from __future__ import annotations
+
+import json
 import os
-import time
-import gc
-import sys
-import threading
-from itertools import islice
-from datetime import datetime
-import re  # for parsing <think> blocks
+from typing import Any, Dict, List, Tuple
+
 import gradio as gr
+import spaces
 import torch
-from transformers import pipeline, TextIteratorStreamer, StoppingCriteria
-from transformers import AutoTokenizer
-from ddgs import DDGS
-import spaces  # Import spaces early to enable ZeroGPU support
-from torch.utils._pytree import tree_map
-
-# Global event to signal cancellation from the UI thread to the generation thread
-cancel_event = threading.Event()
+from transformers import AutoTokenizer, pipeline, BitsAndBytesConfig
 
-access_token=os.environ['HF_TOKEN']
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN environment variable must be set for private router checkpoints.")
 
-# Optional: Disable GPU visibility if you wish to force CPU usage
-# os.environ["CUDA_VISIBLE_DEVICES"] = ""
+ROUTER_SYSTEM_PROMPT = """You are the Router Agent coordinating Math, Code, and General-Search specialists.\nEmit ONLY strict JSON with keys route_plan, route_rationale, expected_artifacts,\nthinking_outline, handoff_plan, todo_list, difficulty, tags, acceptance_criteria, metrics.\nEach route_plan entry must be a tool call (e.g., /math(...), /code(...), /general-search(...)).\nBe concise but precise. Do not include prose outside of the JSON object."""
 
-# ------------------------------
-# Torch-Compatible Model Definitions with Adjusted Descriptions
-# ------------------------------
 MODELS = {
     "Router-Qwen3-32B-8bit": {
         "repo_id": "Alovestocode/router-qwen3-32b-merged",
-        "description": "CourseGPT-Pro router checkpoint built on Qwen3 32B and quantized to 8-bit for ZeroGPU deployment. Optimised for orchestrating math, code, and general-search agents while staying within the ZeroGPU memory envelope.",
+        "description": "Router checkpoint on Qwen3 32B merged and quantized for 8-bit ZeroGPU inference.",
         "params_b": 32.0,
     },
     "Router-Gemma3-27B-8bit": {
         "repo_id": "Alovestocode/router-gemma3-merged",
-        "description": "CourseGPT-Pro router checkpoint built on Gemma3 27B and quantized to 8-bit. Provides the same JSON routing schema with a Gemma-flavoured inductive bias for math-heavy prompts.",
+        "description": "Router checkpoint on Gemma3 27B merged and quantized for 8-bit ZeroGPU inference.",
         "params_b": 27.0,
     },
 }
 
-# Global cache for pipelines to avoid re-loading.
-PIPELINES = {}
+REQUIRED_KEYS = [
+    "route_plan",
+    "route_rationale",
+    "expected_artifacts",
+    "thinking_outline",
+    "handoff_plan",
+    "todo_list",
+    "difficulty",
+    "tags",
+    "acceptance_criteria",
+    "metrics",
+]
 
-def load_pipeline(model_name):
-    """
-    Load and cache a transformers pipeline for text generation.
-    Prefers 8-bit loading (bitsandbytes) to stay within ZeroGPU limits,
-    falling back to bf16/fp16/fp32 if quantized loading is unavailable.
-    """
-    global PIPELINES
+PIPELINES: Dict[str, Any] = {}
+
+
+def load_pipeline(model_name: str):
     if model_name in PIPELINES:
         return PIPELINES[model_name]
 
     repo = MODELS[model_name]["repo_id"]
-    tokenizer = AutoTokenizer.from_pretrained(repo, token=access_token)
+    tokenizer = AutoTokenizer.from_pretrained(repo, token=HF_TOKEN)
 
-    # First try to load in 8-bit to minimise VRAM usage.
     try:
+        quantization_config = BitsAndBytesConfig(load_in_8bit=True)
         pipe = pipeline(
             task="text-generation",
             model=repo,
             tokenizer=tokenizer,
             trust_remote_code=True,
             device_map="auto",
-            model_kwargs={"load_in_8bit": True},
+            model_kwargs={"quantization_config": quantization_config},
             use_cache=True,
-            token=access_token,
+            token=HF_TOKEN,
         )
         PIPELINES[model_name] = pipe
         return pipe
     except Exception as exc:
         print(f"8-bit load failed for {repo}: {exc}. Falling back to higher precision.")
 
-    # Fallback ladder when 8-bit is not available.
     for dtype in (torch.bfloat16, torch.float16, torch.float32):
         try:
             pipe = pipeline(
@@ -79,17 +75,16 @@ def load_pipeline(model_name):
                 model=repo,
                 tokenizer=tokenizer,
                 trust_remote_code=True,
-                dtype=dtype,
                 device_map="auto",
+                dtype=dtype,
                 use_cache=True,
-                token=access_token,
+                token=HF_TOKEN,
             )
             PIPELINES[model_name] = pipe
             return pipe
         except Exception:
             continue
 
-    # Final fallback with framework defaults.
     pipe = pipeline(
         task="text-generation",
         model=repo,
@@ -97,523 +92,234 @@ def load_pipeline(model_name):
         trust_remote_code=True,
         device_map="auto",
         use_cache=True,
-        token=access_token,
+        token=HF_TOKEN,
     )
     PIPELINES[model_name] = pipe
     return pipe
 
 
-def retrieve_context(query, max_results=6, max_chars=50):
-    """
-    Retrieve search snippets from DuckDuckGo (runs in background).
-    Returns a list of result strings.
-    """
-    try:
-        with DDGS() as ddgs:
-            return [f"{i+1}. {r.get('title','No Title')} - {r.get('body','')[:max_chars]}"
-                    for i, r in enumerate(islice(ddgs.text(query, region="wt-wt", safesearch="off", timelimit="y"), max_results))]
-    except Exception:
-        return []
-
-def format_conversation(history, system_prompt, tokenizer):
-    if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
-        messages = [{"role": "system", "content": system_prompt.strip()}] + history
-        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
-    else:
-        # Fallback for base LMs without chat template
-        prompt = system_prompt.strip() + "\n"
-        for msg in history:
-            if msg['role'] == 'user':
-                prompt += "User: " + msg['content'].strip() + "\n"
-            elif msg['role'] == 'assistant':
-                prompt += "Assistant: " + msg['content'].strip() + "\n"
-        if not prompt.strip().endswith("Assistant:"):
-            prompt += "Assistant: "
-        return prompt
-
-def get_duration(user_msg, chat_history, system_prompt, enable_search, max_results, max_chars, model_name, max_tokens, temperature, top_k, top_p, repeat_penalty, search_timeout):
-    # Get model size from the MODELS dict (more reliable than string parsing)
-    model_size = MODELS[model_name].get("params_b", 4.0)  # Default to 4B if not found
-    
-    # Only use AOT for models >= 2B parameters
-    use_aot = model_size >= 2
-    
-    # Adjusted for H200 performance: faster inference, quicker compilation
-    base_duration = 20 if not use_aot else 40  # Reduced base times
-    token_duration = max_tokens * 0.005  # ~200 tokens/second average on H200
-    search_duration = 10 if enable_search else 0  # Reduced search time
-    aot_compilation_buffer = 20 if use_aot else 0  # Faster compilation on H200
-    
-    return base_duration + token_duration + search_duration + aot_compilation_buffer
-
-@spaces.GPU(duration=get_duration)
-def chat_response(user_msg, chat_history, system_prompt,
-                  enable_search, max_results, max_chars,
-                  model_name, max_tokens, temperature,
-                  top_k, top_p, repeat_penalty, search_timeout):
-    """
-    Generates streaming chat responses, optionally with background web search.
-    This version includes cancellation support.
-    """
-    # Clear the cancellation event at the start of a new generation
-    cancel_event.clear()
+def build_router_prompt(
+    user_task: str,
+    context: str,
+    acceptance: str,
+    extra_guidance: str,
+    difficulty: str,
+    tags: str,
+) -> str:
+    prompt_parts = [ROUTER_SYSTEM_PROMPT.strip(), "\n### Router Inputs\n"]
+    prompt_parts.append(f"Difficulty: {difficulty or 'intermediate'}")
+    prompt_parts.append(f"Tags: {tags or 'general'}")
+    if acceptance.strip():
+        prompt_parts.append(f"Acceptance criteria: {acceptance.strip()}")
+    if extra_guidance.strip():
+        prompt_parts.append(f"Additional guidance: {extra_guidance.strip()}")
+    if context.strip():
+        prompt_parts.append("\n### Supporting context\n" + context.strip())
+    prompt_parts.append("\n### User task\n" + user_task.strip())
+    prompt_parts.append("\nReturn only JSON.")
+    return "\n".join(prompt_parts)
+
+
+def extract_json_from_text(text: str) -> str:
+    start = text.find("{")
+    if start == -1:
+        raise ValueError("Router output did not contain a JSON object.")
+    depth = 0
+    in_string = False
+    escape = False
+    for idx in range(start, len(text)):
+        ch = text[idx]
+        if in_string:
+            if escape:
+                escape = False
+            elif ch == "\\":
+                escape = True
+            elif ch == '"':
+                in_string = False
+            continue
+        if ch == '"':
+            in_string = True
+            continue
+        if ch == '{':
+            depth += 1
+        elif ch == '}':
+            depth -= 1
+            if depth == 0:
+                return text[start : idx + 1]
+    raise ValueError("Router output JSON appears truncated.")
+
+
+def validate_router_plan(plan: Dict[str, Any]) -> Tuple[bool, List[str]]:
+    issues: List[str] = []
+    for key in REQUIRED_KEYS:
+        if key not in plan:
+            issues.append(f"Missing key: {key}")
+    route_plan = plan.get("route_plan")
+    if not isinstance(route_plan, list) or not route_plan:
+        issues.append("route_plan must be a non-empty list of tool calls")
+    metrics = plan.get("metrics")
+    if not isinstance(metrics, dict):
+        issues.append("metrics must be an object containing primary/secondary entries")
+    todo = plan.get("todo_list")
+    if not isinstance(todo, list) or not todo:
+        issues.append("todo_list must contain at least one checklist item")
+    return len(issues) == 0, issues
+
+
+def format_validation_message(ok: bool, issues: List[str]) -> str:
+    if ok:
+        return "✅ Router plan includes all required fields."
+    bullets = "\n".join(f"- {issue}" for issue in issues)
+    return f"❌ Issues detected:\n{bullets}"
+
+
+@spaces.GPU(duration=600)
+def generate_router_plan(
+    user_task: str,
+    context: str,
+    acceptance: str,
+    extra_guidance: str,
+    difficulty: str,
+    tags: str,
+    model_choice: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> Tuple[str, Dict[str, Any], str, str]:
+    if not user_task.strip():
+        raise gr.Error("User task is required.")
     
-    history = list(chat_history or [])
-    history.append({'role': 'user', 'content': user_msg})
-
-    # Launch web search if enabled
-    debug = ''
-    search_results = []
-    if enable_search:
-        debug = 'Search task started.'
-        thread_search = threading.Thread(
-            target=lambda: search_results.extend(
-                retrieve_context(user_msg, int(max_results), int(max_chars))
-            )
-        )
-        thread_search.daemon = True
-        thread_search.start()
-    else:
-        debug = 'Web search disabled.'
+    if model_choice not in MODELS:
+        raise gr.Error(f"Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}")
 
     try:
-        cur_date = datetime.now().strftime('%Y-%m-%d')
-        # merge any fetched search results into the system prompt
-        if search_results:
-            
-            enriched = system_prompt.strip() + \
-            f'''\n# The following contents are the search results related to the user's message:
-            {search_results}
-            In the search results I provide to you, each result is formatted as [webpage X begin]...[webpage X end], where X represents the numerical index of each article. Please cite the context at the end of the relevant sentence when appropriate. Use the citation format [citation:X] in the corresponding part of your answer. If a sentence is derived from multiple contexts, list all relevant citation numbers, such as [citation:3][citation:5]. Be sure not to cluster all citations at the end; instead, include them in the corresponding parts of the answer.
-            When responding, please keep the following points in mind:
-            - Today is {cur_date}.
-            - Not all content in the search results is closely related to the user's question. You need to evaluate and filter the search results based on the question.
-            - For listing-type questions (e.g., listing all flight information), try to limit the answer to 10 key points and inform the user that they can refer to the search sources for complete information. Prioritize providing the most complete and relevant items in the list. Avoid mentioning content not provided in the search results unless necessary.
-            - For creative tasks (e.g., writing an essay), ensure that references are cited within the body of the text, such as [citation:3][citation:5], rather than only at the end of the text. You need to interpret and summarize the user's requirements, choose an appropriate format, fully utilize the search results, extract key information, and generate an answer that is insightful, creative, and professional. Extend the length of your response as much as possible, addressing each point in detail and from multiple perspectives, ensuring the content is rich and thorough.
-            - If the response is lengthy, structure it well and summarize it in paragraphs. If a point-by-point format is needed, try to limit it to 5 points and merge related content.
-            - For objective Q&A, if the answer is very brief, you may add one or two related sentences to enrich the content.
-            - Choose an appropriate and visually appealing format for your response based on the user's requirements and the content of the answer, ensuring strong readability.
-            - Your answer should synthesize information from multiple relevant webpages and avoid repeatedly citing the same webpage.
-            - Unless the user requests otherwise, your response should be in the same language as the user's question.
-            # The user's message is:
-            '''
-        else:
-            enriched = system_prompt
-
-        # wait up to 1s for snippets, then replace debug with them
-        if enable_search:
-            thread_search.join(timeout=float(search_timeout))
-            if search_results:
-                debug = "### Search results merged into prompt\n\n" + "\n".join(
-                    f"- {r}" for r in search_results
-                )
-            else:
-                debug = "*No web search results found.*"
-
-        # merge fetched snippets into the system prompt
-        if search_results:
-            enriched = system_prompt.strip() + \
-            f'''\n# The following contents are the search results related to the user's message:
-            {search_results}
-            In the search results I provide to you, each result is formatted as [webpage X begin]...[webpage X end], where X represents the numerical index of each article. Please cite the context at the end of the relevant sentence when appropriate. Use the citation format [citation:X] in the corresponding part of your answer. If a sentence is derived from multiple contexts, list all relevant citation numbers, such as [citation:3][citation:5]. Be sure not to cluster all citations at the end; instead, include them in the corresponding parts of the answer.
-            When responding, please keep the following points in mind:
-            - Today is {cur_date}.
-            - Not all content in the search results is closely related to the user's question. You need to evaluate and filter the search results based on the question.
-            - For listing-type questions (e.g., listing all flight information), try to limit the answer to 10 key points and inform the user that they can refer to the search sources for complete information. Prioritize providing the most complete and relevant items in the list. Avoid mentioning content not provided in the search results unless necessary.
-            - For creative tasks (e.g., writing an essay), ensure that references are cited within the body of the text, such as [citation:3][citation:5], rather than only at the end of the text. You need to interpret and summarize the user's requirements, choose an appropriate format, fully utilize the search results, extract key information, and generate an answer that is insightful, creative, and professional. Extend the length of your response as much as possible, addressing each point in detail and from multiple perspectives, ensuring the content is rich and thorough.
-            - If the response is lengthy, structure it well and summarize it in paragraphs. If a point-by-point format is needed, try to limit it to 5 points and merge related content.
-            - For objective Q&A, if the answer is very brief, you may add one or two related sentences to enrich the content.
-            - Choose an appropriate and visually appealing format for your response based on the user's requirements and the content of the answer, ensuring strong readability.
-            - Your answer should synthesize information from multiple relevant webpages and avoid repeatedly citing the same webpage.
-            - Unless the user requests otherwise, your response should be in the same language as the user's question.
-            # The user's message is:
-            '''
-        else:
-            enriched = system_prompt
-
-        pipe = load_pipeline(model_name)
-
-        prompt = format_conversation(history, enriched, pipe.tokenizer)
-        prompt_debug = f"\n\n--- Prompt Preview ---\n```\n{prompt}\n```"
-        streamer = TextIteratorStreamer(pipe.tokenizer,
-                                        skip_prompt=True,
-                                        skip_special_tokens=True)
-        gen_thread = threading.Thread(
-            target=pipe,
-            args=(prompt,),
-            kwargs={
-                'max_new_tokens': max_tokens,
-                'temperature': temperature,
-                'top_k': top_k,
-                'top_p': top_p,
-                'repetition_penalty': repeat_penalty,
-                'streamer': streamer,
-                'return_full_text': False,
-            }
+        prompt = build_router_prompt(
+            user_task=user_task,
+            context=context,
+            acceptance=acceptance,
+            extra_guidance=extra_guidance,
+            difficulty=difficulty,
+            tags=tags,
         )
-        gen_thread.start()
-
-        # Buffers for thought vs answer
-        thought_buf = ''
-        answer_buf = ''
-        in_thought = False
-        assistant_message_started = False
-
-        # First yield contains the user message
-        yield history, debug
-
-        # Stream tokens
-        for chunk in streamer:
-            # Check for cancellation signal
-            if cancel_event.is_set():
-                if assistant_message_started and history and history[-1]['role'] == 'assistant':
-                    history[-1]['content'] += " [Generation Canceled]"
-                yield history, debug
-                break
-            
-            text = chunk
-
-            # Detect start of thinking
-            if not in_thought and '<think>' in text:
-                in_thought = True
-                history.append({'role': 'assistant', 'content': '', 'metadata': {'title': '💭 Thought'}})
-                assistant_message_started = True
-                after = text.split('<think>', 1)[1]
-                thought_buf += after
-                if '</think>' in thought_buf:
-                    before, after2 = thought_buf.split('</think>', 1)
-                    history[-1]['content'] = before.strip()
-                    in_thought = False
-                    answer_buf = after2
-                    history.append({'role': 'assistant', 'content': answer_buf})
-                else:
-                    history[-1]['content'] = thought_buf
-                yield history, debug
-                continue
-
-            if in_thought:
-                thought_buf += text
-                if '</think>' in thought_buf:
-                    before, after2 = thought_buf.split('</think>', 1)
-                    history[-1]['content'] = before.strip()
-                    in_thought = False
-                    answer_buf = after2
-                    history.append({'role': 'assistant', 'content': answer_buf})
-                else:
-                    history[-1]['content'] = thought_buf
-                yield history, debug
-                continue
-
-            # Stream answer
-            if not assistant_message_started:
-                history.append({'role': 'assistant', 'content': ''})
-                assistant_message_started = True
-
-            answer_buf += text
-            history[-1]['content'] = answer_buf.strip()
-            yield history, debug
-
-        gen_thread.join()
-        yield history, debug + prompt_debug
-    except GeneratorExit:
-        # Handle cancellation gracefully
-        print("Chat response cancelled.")
-        # Don't yield anything - let the cancellation propagate
-        return
-    except Exception as e:
-        history.append({'role': 'assistant', 'content': f"Error: {e}"})
-        yield history, debug
-    finally:
-        gc.collect()
-
-
-def update_default_prompt(enable_search):
-    return f"You are a helpful assistant."
-
-def update_duration_estimate(model_name, enable_search, max_results, max_chars, max_tokens, search_timeout):
-    """Calculate and format the estimated GPU duration for current settings."""
-    try:
-        dummy_msg, dummy_history, dummy_system_prompt = "", [], ""
-        duration = get_duration(dummy_msg, dummy_history, dummy_system_prompt, 
-                              enable_search, max_results, max_chars, model_name, 
-                              max_tokens, 0.7, 40, 0.9, 1.2, search_timeout)
-        model_size = MODELS[model_name].get("params_b", 4.0)
-        return (f"⏱️ **Estimated GPU Time: {duration:.1f} seconds**\n\n"
-                f"📊 **Model Size:** {model_size:.1f}B parameters\n"
-                f"🔍 **Web Search:** {'Enabled' if enable_search else 'Disabled'}")
-    except Exception as e:
-        return f"⚠️ Error calculating estimate: {e}"
-
-# ------------------------------
-# Gradio UI
-# ------------------------------
-with gr.Blocks(
-    title="LLM Inference with ZeroGPU",
-    theme=gr.themes.Soft(
-        primary_hue="indigo",
-        secondary_hue="purple",
-        neutral_hue="slate",
-        radius_size="lg",
-        font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
-    ),
-    css="""
-        .duration-estimate { background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%); border-left: 4px solid #667eea; padding: 12px; border-radius: 8px; margin: 16px 0; }
-        .chatbot { border-radius: 12px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); }
-        button.primary { font-weight: 600; }
-        .gradio-accordion { margin-bottom: 12px; }
-    """
-) as demo:
-    # Header
-    gr.Markdown("""
-    # 🧠 ZeroGPU LLM Inference
-    ### Powered by Hugging Face ZeroGPU with Web Search Integration
-    """)
-    
-    with gr.Row():
-        # Left Panel - Configuration
-        with gr.Column(scale=3):
-            # Core Settings (Always Visible)
-            with gr.Group():
-                gr.Markdown("### ⚙️ Core Settings")
-                model_dd = gr.Dropdown(
-                    label="🤖 Model",
-                    choices=list(MODELS.keys()),
-                    value="Qwen3-1.7B",
-                    info="Select the language model to use"
+
+        generator = load_pipeline(model_choice)
+        result = generator(
+            prompt,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+        )[0]["generated_text"]
+
+        completion = result[len(prompt) :].strip() if result.startswith(prompt) else result.strip()
+
+        try:
+            json_block = extract_json_from_text(completion)
+            plan = json.loads(json_block)
+            ok, issues = validate_router_plan(plan)
+            validation_msg = format_validation_message(ok, issues)
+        except Exception as exc:
+            plan = {}
+            validation_msg = f"❌ JSON parsing failed: {exc}"
+
+        return completion, plan, validation_msg, prompt
+    except Exception as exc:
+        error_msg = f"❌ Generation failed: {str(exc)}"
+        return "", {}, error_msg, ""
+
+
+def clear_outputs():
+    return "", {}, "Awaiting generation.", ""
+
+
+def build_ui():
+    description = "Use the CourseGPT-Pro router checkpoints (Gemma3/Qwen3) hosted on ZeroGPU to generate structured routing plans."
+    with gr.Blocks(theme=gr.themes.Soft(), css="""
+        textarea { font-family: 'JetBrains Mono', 'Fira Code', monospace; }
+        .status-ok { color: #0d9488; font-weight: 600; }
+        .status-bad { color: #dc2626; font-weight: 600; }
+    """) as demo:
+        gr.Markdown("# 🛰️ Router Control Room — ZeroGPU" )
+        gr.Markdown(description)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                user_task = gr.Textbox(
+                    label="User Task / Problem Statement",
+                    placeholder="Describe the homework-style query that needs routing...",
+                    lines=8,
+                    value="Explain how to solve a constrained optimization homework problem that mixes calculus and coding steps.",
                 )
-                search_chk = gr.Checkbox(
-                    label="🔍 Enable Web Search",
-                    value=False,
-                    info="Augment responses with real-time web data"
+                context = gr.Textbox(
+                    label="Supporting Context (optional)",
+                    placeholder="Paste any retrieved evidence, PDFs, or rubric notes.",
+                    lines=4,
                 )
-                sys_prompt = gr.Textbox(
-                    label="📝 System Prompt",
+                acceptance = gr.Textbox(
+                    label="Acceptance Criteria",
+                    placeholder="Bullet list of 'definition of done' checks.",
                     lines=3,
-                    value=update_default_prompt(search_chk.value),
-                    placeholder="Define the assistant's behavior and personality..."
-                )
-            
-            # Duration Estimate
-            duration_display = gr.Markdown(
-                value=update_duration_estimate("Qwen3-1.7B", False, 4, 50, 1024, 5.0),
-                elem_classes="duration-estimate"
-            )
-            
-            # Advanced Settings (Collapsible)
-            with gr.Accordion("🎛️ Advanced Generation Parameters", open=False):
-                max_tok = gr.Slider(
-                    64, 16384, value=1024, step=32,
-                    label="Max Tokens",
-                    info="Maximum length of generated response"
+                    value="- Provide citations for every claim.\n- Ensure /math verifies /code output.",
                 )
-                temp = gr.Slider(
-                    0.1, 2.0, value=0.7, step=0.1,
-                    label="Temperature",
-                    info="Higher = more creative, Lower = more focused"
-                )
-                with gr.Row():
-                    k = gr.Slider(
-                        1, 100, value=40, step=1,
-                        label="Top-K",
-                        info="Number of top tokens to consider"
-                    )
-                    p = gr.Slider(
-                        0.1, 1.0, value=0.9, step=0.05,
-                        label="Top-P",
-                        info="Nucleus sampling threshold"
-                    )
-                rp = gr.Slider(
-                    1.0, 2.0, value=1.2, step=0.1,
-                    label="Repetition Penalty",
-                    info="Penalize repeated tokens"
-                )
-            
-            # Web Search Settings (Collapsible)
-            with gr.Accordion("🌐 Web Search Settings", open=False, visible=False) as search_settings:
-                mr = gr.Number(
-                    value=4, precision=0,
-                    label="Max Results",
-                    info="Number of search results to retrieve"
+                extra_guidance = gr.Textbox(
+                    label="Additional Guidance",
+                    placeholder="Special constraints, tools to avoid, etc.",
+                    lines=3,
                 )
-                mc = gr.Number(
-                    value=50, precision=0,
-                    label="Max Chars/Result",
-                    info="Character limit per search result"
+            with gr.Column(scale=2):
+                model_choice = gr.Dropdown(
+                    label="Router Checkpoint",
+                    choices=list(MODELS.keys()),
+                    value=list(MODELS.keys())[0] if MODELS else None,
+                    allow_custom_value=False,
                 )
-                st = gr.Slider(
-                    minimum=0.0, maximum=30.0, step=0.5, value=5.0,
-                    label="Search Timeout (s)",
-                    info="Maximum time to wait for search results"
+                difficulty = gr.Radio(
+                    label="Difficulty Tier",
+                    choices=["introductory", "intermediate", "advanced"],
+                    value="advanced",
+                    interactive=True,
                 )
-            
-            # Actions
-            with gr.Row():
-                clr = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1)
-        
-        # Right Panel - Chat Interface
-        with gr.Column(scale=7):
-            chat = gr.Chatbot(
-                type="messages",
-                height=600,
-                label="💬 Conversation",
-                show_copy_button=True,
-                avatar_images=(None, "🤖"),
-                bubble_full_width=False
-            )
-            
-            # Input Area
-            with gr.Row():
-                txt = gr.Textbox(
-                    placeholder="💭 Type your message here... (Press Enter to send)",
-                    scale=9,
-                    container=False,
-                    show_label=False,
-                    lines=1,
-                    max_lines=5
+                tags = gr.Textbox(
+                    label="Tags",
+                    placeholder="Comma-separated e.g. calculus, optimization, python",
+                    value="calculus, optimization, python",
                 )
-                with gr.Column(scale=1, min_width=120):
-                    submit_btn = gr.Button("📤 Send", variant="primary", size="lg")
-                    cancel_btn = gr.Button("⏹️ Stop", variant="stop", visible=False, size="lg")
-            
-            # Example Prompts
-            gr.Examples(
-                examples=[
-                    ["Explain quantum computing in simple terms"],
-                    ["Write a Python function to calculate fibonacci numbers"],
-                    ["What are the latest developments in AI? (Enable web search)"],
-                    ["Tell me a creative story about a time traveler"],
-                    ["Help me debug this code: def add(a,b): return a+b+1"]
-                ],
-                inputs=txt,
-                label="💡 Example Prompts"
-            )
-            
-            # Debug/Status Info (Collapsible)
-            with gr.Accordion("🔍 Debug Info", open=False):
-                dbg = gr.Markdown()
-    
-    # Footer
-    gr.Markdown("""
-    ---
-    💡 **Tips:** 
-    - Use **Advanced Parameters** to fine-tune creativity and response length
-    - Enable **Web Search** for real-time, up-to-date information
-    - Try different **models** for various tasks (reasoning, coding, general chat)
-    - Click the **Copy** button on responses to save them to your clipboard
-    """, elem_classes="footer")
-
-    # --- Event Listeners ---
-
-    # Group all inputs for cleaner event handling
-    chat_inputs = [txt, chat, sys_prompt, search_chk, mr, mc, model_dd, max_tok, temp, k, p, rp, st]
-    # Group all UI components that can be updated.
-    ui_components = [chat, dbg, txt, submit_btn, cancel_btn]
-
-    def submit_and_manage_ui(user_msg, chat_history, *args):
-        """
-        Orchestrator function that manages UI state and calls the backend chat function.
-        It uses a try...finally block to ensure the UI is always reset.
-        """
-        if not user_msg.strip():
-            # If the message is empty, do nothing.
-            # We yield an empty dict to avoid any state changes.
-            yield {}
-            return
-
-        # 1. Update UI to "generating" state.
-        #    Crucially, we do NOT update the `chat` component here, as the backend
-        #    will provide the correctly formatted history in the first response chunk.
-        yield {
-            txt: gr.update(value="", interactive=False),
-            submit_btn: gr.update(interactive=False),
-            cancel_btn: gr.update(visible=True),
-        }
-
-        cancelled = False
-        try:
-            # 2. Call the backend and stream updates
-            backend_args = [user_msg, chat_history] + list(args)
-            for response_chunk in chat_response(*backend_args):
-                yield {
-                    chat: response_chunk[0],
-                    dbg: response_chunk[1],
-                }
-        except GeneratorExit:
-            # Mark as cancelled and re-raise to prevent "generator ignored GeneratorExit"
-            cancelled = True
-            print("Generation cancelled by user.")
-            raise
-        except Exception as e:
-            print(f"An error occurred during generation: {e}")
-            # If an error happens, add it to the chat history to inform the user.
-            error_history = (chat_history or []) + [
-                {'role': 'user', 'content': user_msg},
-                {'role': 'assistant', 'content': f"**An error occurred:** {str(e)}"}
-            ]
-            yield {chat: error_history}
-        finally:
-            # Only reset UI if not cancelled (to avoid "generator ignored GeneratorExit")
-            if not cancelled:
-                print("Resetting UI state.")
-                yield {
-                    txt: gr.update(interactive=True),
-                    submit_btn: gr.update(interactive=True),
-                    cancel_btn: gr.update(visible=False),
-                }
-
-    def set_cancel_flag():
-        """Called by the cancel button, sets the global event."""
-        cancel_event.set()
-        print("Cancellation signal sent.")
-    
-    def reset_ui_after_cancel():
-        """Reset UI components after cancellation."""
-        cancel_event.clear()  # Clear the flag for next generation
-        print("UI reset after cancellation.")
-        return {
-            txt: gr.update(interactive=True),
-            submit_btn: gr.update(interactive=True),
-            cancel_btn: gr.update(visible=False),
-        }
-
-    # Event for submitting text via Enter key or Submit button
-    submit_event = txt.submit(
-        fn=submit_and_manage_ui,
-        inputs=chat_inputs,
-        outputs=ui_components,
-    )
-    submit_btn.click(
-        fn=submit_and_manage_ui,
-        inputs=chat_inputs,
-        outputs=ui_components,
-    )
+                max_new_tokens = gr.Slider(256, 1024, value=640, step=32, label="Max New Tokens")
+                temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature")
+                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+
+        generate_btn = gr.Button("Generate Router Plan", variant="primary")
+        clear_btn = gr.Button("Clear", variant="secondary")
+
+        with gr.Row():
+            raw_output = gr.Textbox(label="Raw Model Output", lines=12)
+            plan_json = gr.JSON(label="Parsed Router Plan")
+        validation_msg = gr.Markdown("Awaiting generation.")
+        prompt_view = gr.Textbox(label="Full Prompt", lines=10)
+
+        generate_btn.click(
+            generate_router_plan,
+            inputs=[
+                user_task,
+                context,
+                acceptance,
+                extra_guidance,
+                difficulty,
+                tags,
+                model_choice,
+                max_new_tokens,
+                temperature,
+                top_p,
+            ],
+            outputs=[raw_output, plan_json, validation_msg, prompt_view],
+        )
 
-    # Event for the "Cancel" button.
-    # It sets the cancel flag, cancels the submit event, then resets the UI.
-    cancel_btn.click(
-        fn=set_cancel_flag,
-        cancels=[submit_event]
-    ).then(
-        fn=reset_ui_after_cancel,
-        outputs=ui_components
-    )
+        clear_btn.click(fn=clear_outputs, outputs=[raw_output, plan_json, validation_msg, prompt_view])
 
-    # Listeners for updating the duration estimate
-    duration_inputs = [model_dd, search_chk, mr, mc, max_tok, st]
-    for component in duration_inputs:
-        component.change(fn=update_duration_estimate, inputs=duration_inputs, outputs=duration_display)
+    return demo
 
-    # Toggle web search settings visibility
-    def toggle_search_settings(enabled):
-        return gr.update(visible=enabled)
-    
-    search_chk.change(
-        fn=lambda enabled: (update_default_prompt(enabled), gr.update(visible=enabled)),
-        inputs=search_chk,
-        outputs=[sys_prompt, search_settings]
-    )
-    
-    # Clear chat action
-    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
-    
-    demo.launch()
+
+demo = build_ui()
+
+if __name__ == "__main__":  # pragma: no cover
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))