feat: ui/ux improvements

app.py

@@ -1,6 +1,7 @@
 import gc
 import html
 import inspect
+import re
 import threading
 import time
 
@@ -70,7 +71,28 @@ PROMPT_TEMPLATE = (
     "<|assistant|>"
 )
 
+GENERAL_PROMPT_TEMPLATE = (
+    "<|user|>\n"
+    "You are Phi-3 Mini in a SQL generator demo. Reply naturally and briefly. "
+    "If the user asks for SQL, provide only the SQL query.\n\n"
+    "User: {message}<|end|>\n"
+    "<|assistant|>"
+)
+
 EMPTY_VALIDATOR = '<span class="validator-badge validator-empty">No SQL yet</span>'
+CHAT_VALIDATOR = '<span class="validator-badge validator-empty">Chat response</span>'
+EMPTY_CHAT_OUTPUT = ""
+LOAD_SCROLL_JS = """
+(selectedKey) => {
+  setTimeout(() => {
+    document.querySelector("#query-section")?.scrollIntoView({
+      behavior: "smooth",
+      block: "start"
+    });
+  }, 50);
+  return selectedKey;
+}
+"""
 
 _current_model_id = None
 _model = None
@@ -190,6 +212,68 @@ def clean_generation(text):
     return cleaned
 
 
+def is_sql_like(text):
+    text = (text or "").strip()
+    if not text:
+        return False
+    first_word = re.match(r"^\s*([A-Za-z]+)", text)
+    if not first_word:
+        return False
+    return first_word.group(1).upper() in {
+        "SELECT",
+        "WITH",
+        "INSERT",
+        "UPDATE",
+        "DELETE",
+        "CREATE",
+        "ALTER",
+        "DROP",
+    }
+
+
+def is_sql_intent(message, schema):
+    message = (message or "").strip().lower()
+    schema = (schema or "").strip()
+    if schema:
+        return True
+    if not message:
+        return False
+    sql_terms = {
+        "sql",
+        "query",
+        "select",
+        "table",
+        "schema",
+        "database",
+        "join",
+        "group by",
+        "order by",
+        "where",
+        "average",
+        "count",
+        "sum",
+        "rows",
+        "columns",
+    }
+    return any(term in message for term in sql_terms)
+
+
+def build_generation_prompt(schema, message):
+    schema = (schema or "").strip()
+    message = (message or "").strip()
+    if is_sql_intent(message, schema):
+        table_schema = schema or "No explicit schema provided. Infer the table and columns only if the request includes them."
+        return PROMPT_TEMPLATE.format(schema=table_schema, question=message)
+    return GENERAL_PROMPT_TEMPLATE.format(message=message)
+
+
+def format_generation_result(text):
+    cleaned = clean_generation(text)
+    if is_sql_like(cleaned):
+        return str(cleaned), EMPTY_CHAT_OUTPUT, validate_sql(cleaned)
+    return "", str(cleaned), CHAT_VALIDATOR
+
+
 def validate_sql(sql_text):
     sql_text = (sql_text or "").strip()
     if not sql_text:
@@ -224,12 +308,12 @@ def render_header():
     <section class="top-panel">
       <div>
         <h1>Phi-3 Mini SQL Generator</h1>
-        <p>QLoRA fine-tuned - b-mc2/sql-create-context - T4 GPU - 21 min</p>
+        <p>QLoRA fine-tuned - b-mc2/sql-create-context</p>
       </div>
       <div class="top-badges">
         <span class="badge badge-green">73.5% exact match</span>
         <span class="badge badge-cream">+71.5pp vs base</span>
-        <span class="badge badge-light">CPU - lazy load</span>
+        <span class="badge badge-light">CPU lazy load</span>
       </div>
     </section>
     """
@@ -238,8 +322,7 @@ def render_header():
 def render_step(number, title):
     return f"""
     <div class="step-title">
-      <span class="step-index">{number}</span>
-      <h2>{title}</h2>
+      <span>{number} &mdash; {title}</span>
     </div>
     """
 
@@ -248,8 +331,6 @@ def render_model_card(model_key, selected_key):
     model_def = model_by_key(model_key)
     selected = model_key == selected_key
     state_class = " selected" if selected else ""
-    action = "Selected" if selected else "Select"
-    ring = "selected-mark" if selected else "empty-mark"
     return f"""
     <article class="model-card{state_class}" role="button" tabindex="0">
       <div class="model-tag">{model_def["tag"]}</div>
@@ -260,8 +341,7 @@ def render_model_card(model_key, selected_key):
         <small>exact match</small>
       </div>
       <div class="model-card-footer">
-        <span>{action}</span>
-        <span class="{ring}"></span>
+        <span>{model_def["label"]}</span>
       </div>
     </article>
     """
@@ -292,21 +372,51 @@ def render_loading_overlay(model_key=None, visible=False):
     return (
         '<div class="loading-overlay">'
         '<div class="loading-card">'
-        f'<h2>Loading {model_def["short_label"]} model...</h2>'
+        f'<div class="loading-title">Loading {model_def["short_label"]} model</div>'
         '<div class="loading-line"><span></span></div>'
-        '<p>First load: ~3-5 min - model stays cached for this session</p>'
+        '<p>First load: ~3-5 min &mdash; cached for session</p>'
         "</div>"
         "</div>"
     )
 
 
 def model_metadata(model_key=None):
-    return model_by_key(model_key or DEFAULT_MODEL_KEY)["metadata"]
+    return """
+    <section class="stats-row">
+      <div class="stat-card"><strong>73.5%</strong><span>exact match</span></div>
+      <div class="stat-card"><strong>+71.5pp</strong><span>vs base</span></div>
+      <div class="stat-card"><strong>1,000</strong><span>examples</span></div>
+      <div class="stat-card"><strong>21 min</strong><span>T4 training</span></div>
+    </section>
+    """
+
+
+def schema_name_by_value(schema):
+    schema = (schema or "").strip()
+    for name, value in PRESETS.items():
+        if value == schema:
+            return name
+    return "custom"
+
+
+def render_schema_context(schema=""):
+    schema = (schema or "").strip()
+    if not schema:
+        return '<div class="schema-context empty"></div>'
+    label = schema_name_by_value(schema)
+    escaped_schema = html.escape(schema)
+    escaped_label = html.escape(label)
+    return (
+        '<div class="schema-context">'
+        f'<span>Context: {escaped_label}</span>'
+        f'<code>{escaped_schema}</code>'
+        "</div>"
+    )
 
 
 def query_control_updates(can_generate):
-    input_updates = [gr.update(interactive=True) for _ in range(7)]
-    return [*input_updates, gr.update(interactive=can_generate)]
+    context_updates = [gr.update(interactive=True) for _ in range(6)]
+    return [*context_updates, gr.update(interactive=True), gr.update(interactive=can_generate)]
 
 
 def render_message(message="", kind="error"):
@@ -339,11 +449,11 @@ def load_selected_model(selected_key):
         render_status(selected_key, None, state="loading"),
         render_loading_overlay(selected_key, visible=True),
         model_metadata(selected_key),
-        gr.update(interactive=False),
+        gr.update(interactive=False, visible=False),
         *query_control_updates(False),
         "",
         EMPTY_VALIDATOR,
-        gr.update(interactive=False),
+        gr.update(interactive=False, visible=False),
         render_message(),
         gr.update(visible=False),
     )
@@ -361,7 +471,7 @@ def load_selected_model(selected_key):
             *query_control_updates(False),
             "",
             EMPTY_VALIDATOR,
-            gr.update(interactive=False),
+            gr.update(interactive=False, visible=False),
             render_message(error),
             gr.update(visible=False),
         )
@@ -377,18 +487,28 @@ def load_selected_model(selected_key):
         *query_control_updates(True),
         "",
         EMPTY_VALIDATOR,
-        gr.update(interactive=False),
+        gr.update(interactive=False, visible=False),
         render_message(f"Loaded {model_def['model_id']} in {elapsed}s.", kind="ok"),
         gr.update(visible=False),
     )
 
 
 def set_preset(name):
-    return PRESETS[name]
+    schema = PRESETS[name]
+    return schema, render_schema_context(schema), gr.update(visible=True)
+
+
+def clear_schema_context():
+    return "", render_schema_context(""), gr.update(visible=False)
+
+
+def trim_chat_history(chat_history, max_exchanges=10):
+    history = list(chat_history or [])
+    return history[-max_exchanges * 2 :]
 
 
-def comparison_updates(saved_state, current_sql, loaded_key, schema, question):
-    if not saved_state:
+def comparison_updates(saved_state, current_sql, loaded_key):
+    if not saved_state or not (current_sql or "").strip():
         return gr.update(visible=False), "", "", "", ""
 
     loaded_def = model_by_key(loaded_key) if loaded_key else model_by_key(DEFAULT_MODEL_KEY)
@@ -402,39 +522,55 @@ def comparison_updates(saved_state, current_sql, loaded_key, schema, question):
 
 
 def render_compare_label(prefix, model_label, metric):
-    metric_html = f'<strong>{metric} match</strong>' if metric else ""
-    return f'<div class="compare-head"><span>{prefix} - {model_label}</span>{metric_html}</div>'
+    metric_html = f"<strong>{html.escape(metric)} match</strong>" if metric else ""
+    return (
+        f'<div class="compare-head"><span>{html.escape(prefix)} - '
+        f"{html.escape(model_label)}</span>{metric_html}</div>"
+    )
 
 
-def generate_sql(schema, question, loaded_key, saved_state):
-    schema = (schema or "").strip()
-    question = (question or "").strip()
+def generate_response(message, chat_history, active_schema, loaded_key, saved_state):
+    message = (message or "").strip()
+    active_schema = (active_schema or "").strip()
+    chat_history = list(chat_history or [])
     if not loaded_key or _model is None or _tokenizer is None:
-        compare = comparison_updates(saved_state, "", loaded_key, schema, question)
+        compare = comparison_updates(saved_state, "", loaded_key)
         return (
+            chat_history,
+            message,
+            active_schema,
+            "",
             "",
             EMPTY_VALIDATOR,
-            gr.update(interactive=False),
+            gr.update(interactive=False, visible=False),
             render_message("Load a model before generating SQL."),
             *compare,
         )
-    if not schema or not question:
-        compare = comparison_updates(saved_state, "", loaded_key, schema, question)
+    if not message:
+        compare = comparison_updates(saved_state, "", loaded_key)
         return (
+            chat_history,
+            "",
+            active_schema,
+            "",
             "",
             EMPTY_VALIDATOR,
-            gr.update(interactive=False),
-            render_message("Schema and question are required."),
+            gr.update(interactive=False, visible=False),
+            render_message("Type a message before sending."),
             *compare,
         )
 
     model_def = model_by_key(loaded_key)
     if _current_model_id != model_def["model_id"]:
-        compare = comparison_updates(saved_state, "", loaded_key, schema, question)
+        compare = comparison_updates(saved_state, "", loaded_key)
         return (
+            chat_history,
+            message,
+            active_schema,
+            "",
             "",
             EMPTY_VALIDATOR,
-            gr.update(interactive=False),
+            gr.update(interactive=False, visible=False),
             render_message("Loaded model state is inconsistent. Reload the selected model."),
             *compare,
         )
@@ -443,40 +579,58 @@ def generate_sql(schema, question, loaded_key, saved_state):
     try:
         torch, _, _, _ = import_model_runtime()
         with _model_lock:
-            prompt = PROMPT_TEMPLATE.format(schema=schema, question=question)
+            prompt = build_generation_prompt(active_schema, message)
             inputs = _tokenizer(prompt, return_tensors="pt")
             input_length = inputs["input_ids"].shape[-1]
             with torch.no_grad():
                 output_ids = _model.generate(
                     **inputs,
-                    max_new_tokens=150,
+                    max_new_tokens=80,
                     do_sample=False,
                     use_cache=False,
                 )
             generated_ids = output_ids[0][input_length:]
-            sql_text = clean_generation(_tokenizer.decode(generated_ids, skip_special_tokens=False))
+            generated_text = _tokenizer.decode(generated_ids, skip_special_tokens=False)
     except Exception as exc:
-        compare = comparison_updates(saved_state, "", loaded_key, schema, question)
+        compare = comparison_updates(saved_state, "", loaded_key)
         return (
+            chat_history,
+            message,
+            active_schema,
+            "",
             "",
             EMPTY_VALIDATOR,
-            gr.update(interactive=False),
+            gr.update(interactive=False, visible=False),
             render_message(f"Generation failed: {type(exc).__name__}: {exc}"),
             *compare,
         )
 
     elapsed = int(time.time() - started)
-    compare = comparison_updates(saved_state, sql_text, loaded_key, schema, question)
+    sql_text, chat_text, validator = format_generation_result(generated_text)
+    display_response = f"```sql\n{sql_text}\n```" if sql_text else chat_text
+    new_history = trim_chat_history(
+        [
+            *chat_history,
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": display_response},
+        ]
+    )
+    compare = comparison_updates(saved_state, sql_text, loaded_key)
+    response_kind = "SQL" if sql_text.strip() else "chat response"
     return (
+        new_history,
+        "",
+        active_schema,
+        message,
         str(sql_text),
-        validate_sql(sql_text),
-        gr.update(interactive=bool(sql_text.strip())),
-        render_message(f"Generated with {model_def['model_id']} in {elapsed}s.", kind="ok"),
+        validator,
+        gr.update(interactive=bool(sql_text.strip()), visible=bool(sql_text.strip())),
+        render_message(f"Generated {response_kind} with {model_def['model_id']} in {elapsed}s.", kind="ok"),
         *compare,
     )
 
 
-def save_for_comparison(sql_text, loaded_key, schema, question):
+def save_for_comparison(sql_text, loaded_key, active_schema, last_message):
     sql_text = (sql_text or "").strip()
     if not sql_text or not loaded_key:
         return (
@@ -486,7 +640,7 @@ def save_for_comparison(sql_text, loaded_key, schema, question):
             "",
             "",
             "",
-            gr.update(interactive=False),
+            gr.update(interactive=False, visible=False),
             render_message("Generate SQL before saving a comparison."),
         )
 
@@ -495,8 +649,8 @@ def save_for_comparison(sql_text, loaded_key, schema, question):
         "sql": sql_text,
         "model_label": model_def["short_label"],
         "match": model_def["exact_match"],
-        "schema": schema or "",
-        "question": question or "",
+        "schema_context": active_schema or "",
+        "user_message": last_message or "",
     }
     return (
         saved,
@@ -511,201 +665,188 @@ def save_for_comparison(sql_text, loaded_key, schema, question):
 
 
 CSS = """
-@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@500;600;700;800&family=Space+Mono:wght@400;700&display=swap');
+@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;500;700&display=swap');
 
 :root {
-  --bg: #000000;
-  --panel: #242422;
-  --panel-2: #2e2f2d;
-  --panel-3: #111110;
-  --line: #686866;
-  --line-bright: #f3f0ea;
-  --text: #f4f1ea;
-  --muted: #b8b4ad;
-  --mint: #19b991;
-  --mint-soft: #dff8ef;
-  --cream: #fff2dd;
-  --cream-text: #8b4f12;
-  --code-blue: #2878a9;
+  --bg-base: #0c0c0b;
+  --bg-surface: #1a1a18;
+  --bg-raised: #242422;
+  --border: rgba(255, 255, 255, 0.1);
+  --border-hi: rgba(255, 255, 255, 0.22);
+  --text-hi: #f2ede4;
+  --text-mid: #9e9a93;
+  --text-lo: #5f5d58;
+  --text-primary: var(--text-hi);
+  --text-secondary: var(--text-mid);
+  --text-muted: var(--text-lo);
+  --teal: #1D9E75;
+  --teal-soft: #dff8ef;
+  --teal-text: #0F6E56;
+  --amber-soft: #FAEEDA;
+  --amber-text: #854F0B;
 }
 
 * {
   box-sizing: border-box;
 }
 
-.gradio-container {
-  background: var(--bg) !important;
-  color: var(--text) !important;
-  font-family: Outfit, ui-sans-serif, system-ui, sans-serif !important;
-}
-
+.gradio-container,
 .gradio-container .main,
 .gradio-container .wrap,
 .gradio-container .contain {
-  background: var(--bg) !important;
+  background: var(--bg-base) !important;
+  color: var(--text-primary) !important;
+  font-family: Space Mono, ui-monospace, SFMono-Regular, Menlo, Consolas, monospace !important;
 }
 
 .app-shell {
-  max-width: 1210px;
-  margin: 28px auto 52px;
-  padding: 0 22px;
+  max-width: 1120px;
+  margin: 22px auto 44px;
+  padding: 0 20px;
 }
 
 .top-panel {
-  align-items: flex-start;
-  background: var(--panel);
-  border: 2px solid #454542;
-  border-radius: 12px;
+  align-items: center;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
   display: grid;
-  gap: 24px;
-  grid-template-columns: 1fr auto;
-  min-height: 118px;
-  padding: 18px 25px;
+  gap: 16px;
+  grid-template-columns: minmax(0, 1fr) auto;
+  padding: 14px 16px;
 }
 
 .top-panel h1 {
-  color: var(--text);
-  font-size: 28px;
-  font-weight: 600;
+  color: var(--text-primary);
+  font-size: 15px;
+  font-weight: 500;
   letter-spacing: 0;
-  line-height: 1.1;
-  margin: 0 0 10px;
+  line-height: 1.25;
+  margin: 0 0 4px;
 }
 
 .top-panel p {
-  color: var(--muted);
-  font-size: 19px;
-  font-weight: 800;
+  color: var(--text-secondary);
+  font-size: 13px;
+  font-weight: 400;
   letter-spacing: 0;
+  line-height: 1.35;
   margin: 0;
 }
 
 .top-badges {
+  align-items: center;
   display: flex;
   flex-wrap: wrap;
-  gap: 10px;
+  gap: 8px;
   justify-content: flex-end;
-  max-width: 560px;
 }
 
-.badge {
-  border-radius: 7px;
+.badge,
+.validator-badge,
+.model-tag {
+  border-radius: 5px;
   display: inline-flex;
-  font-size: 18px;
-  font-weight: 700;
+  font-size: 11px;
+  font-weight: 500;
+  letter-spacing: 0;
   line-height: 1;
-  padding: 10px 14px;
+  padding: 6px 8px;
 }
 
-.badge-green {
-  background: #dcf5ee;
-  color: #147963;
+.badge-green,
+.validator-ok {
+  background: var(--teal-soft);
+  color: var(--teal-text);
 }
 
-.badge-cream {
-  background: var(--cream);
-  color: var(--cream-text);
+.badge-cream,
+.validator-warn {
+  background: var(--amber-soft);
+  color: var(--amber-text);
 }
 
-.badge-light {
-  background: #f2f1ed;
-  color: #6f6b65;
+.badge-light,
+.validator-empty {
+  background: var(--bg-raised);
+  color: var(--text-secondary);
+  border: 0.5px solid var(--border);
 }
 
 .step-title {
-  align-items: center;
-  display: flex;
-  gap: 10px;
-  margin: 42px 0 14px;
+  color: var(--text-secondary);
+  font-size: 11px;
+  font-weight: 500;
+  letter-spacing: 0.08em;
+  line-height: 1;
+  margin: 32px 0 12px;
+  text-transform: uppercase;
 }
 
-.step-index {
-  align-items: center;
-  background: var(--mint-soft);
-  border: 2px solid #8ce8d0;
-  border-radius: 50%;
-  color: #11866e;
+.step-title span {
   display: inline-flex;
-  font-family: Space Mono, monospace;
-  font-size: 17px;
-  font-weight: 700;
-  height: 32px;
-  justify-content: center;
-  width: 32px;
 }
 
-.step-title h2 {
-  color: #cfcac1;
-  font-family: Space Mono, monospace;
-  font-size: 22px;
-  font-weight: 700;
-  letter-spacing: -0.03em;
-  line-height: 1;
-  margin: 0;
-  text-transform: uppercase;
-}
-
-.model-grid {
+.model-grid,
+.compare-grid,
+.stats-row {
   display: grid;
-  gap: 18px;
+  gap: 12px;
   grid-template-columns: repeat(2, minmax(0, 1fr));
 }
 
-.model-grid > div {
+.model-grid > div,
+.compare-grid > div,
+.stats-row > div {
   min-width: 0;
 }
 
 .model-card {
-  background: var(--panel-2);
-  border: 2px solid #454542;
-  border-radius: 10px;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
   cursor: pointer;
-  height: 100%;
-  min-height: 266px;
-  padding: 27px 25px 24px;
-  transition: border-color 180ms ease, transform 180ms ease;
+  min-height: 176px;
+  padding: 16px;
+  transition: border-color 160ms ease, background 160ms ease;
 }
 
-.model-card.selected {
-  border-color: var(--mint);
+.model-card:hover {
+  border-color: var(--border-hi);
 }
 
-.model-card:hover {
-  transform: translateY(-2px);
+.model-card.selected {
+  border: 1.5px solid var(--teal);
 }
 
 .model-tag {
-  background: #fff3df;
-  border-radius: 5px;
-  color: #8b5a18;
-  display: inline-flex;
-  font-size: 17px;
-  font-weight: 700;
-  line-height: 1;
-  margin-bottom: 16px;
-  padding: 7px 10px;
+  background: var(--amber-soft);
+  color: var(--amber-text);
+  margin-bottom: 18px;
 }
 
 .model-card.selected .model-tag {
-  background: #e3f7f0;
-  color: #16816a;
+  background: var(--teal-soft);
+  color: var(--teal-text);
 }
 
 .model-card h3 {
-  color: #f6f3ec;
-  font-size: 24px;
-  font-weight: 700;
-  letter-spacing: -0.02em;
+  color: var(--text-primary);
+  font-size: 15px;
+  font-weight: 500;
+  letter-spacing: 0;
+  line-height: 1.3;
   margin: 0 0 8px;
 }
 
 .model-card code {
-  color: #aaa7a0;
+  color: var(--text-secondary);
   display: block;
-  font-family: Space Mono, monospace;
-  font-size: 17px;
-  font-weight: 700;
-  margin-bottom: 17px;
+  font-family: inherit;
+  font-size: 12px;
+  font-weight: 400;
+  line-height: 1.35;
+  margin-bottom: 18px;
   overflow: hidden;
   text-overflow: ellipsis;
   white-space: nowrap;
@@ -715,185 +856,257 @@ CSS = """
   align-items: baseline;
   display: flex;
   gap: 8px;
-  margin-bottom: 24px;
+  margin-bottom: 16px;
 }
 
 .model-score span {
-  color: #f7f3eb;
-  font-size: 40px;
+  color: var(--text-primary);
+  font-size: 28px;
   font-weight: 500;
-  letter-spacing: -0.04em;
+  letter-spacing: 0;
   line-height: 1;
 }
 
 .model-card.selected .model-score span {
-  color: #06896d;
+  color: var(--teal);
 }
 
-.model-score small {
-  color: #cac5ba;
-  font-size: 18px;
-  font-weight: 700;
+.model-score small,
+.model-card-footer {
+  color: var(--text-secondary);
+  font-size: 13px;
+  font-weight: 400;
 }
 
 .model-card-footer {
-  align-items: center;
-  color: #aaa59c;
   display: flex;
-  font-size: 18px;
-  font-weight: 700;
-  justify-content: space-between;
 }
 
-.model-card.selected .model-card-footer {
-  color: #06896d;
+#load-button,
+#generate-button,
+#save-button {
+  width: 100% !important;
+}
+
+.gradio-container button {
+  border-radius: 6px !important;
+  font-family: Space Mono, ui-monospace, monospace !important;
+  font-size: 11px !important;
+  font-weight: 500 !important;
+  letter-spacing: 0 !important;
+  transition: background 160ms ease, border-color 160ms ease, color 160ms ease, opacity 160ms ease !important;
+}
+
+#load-button button,
+#generate-button button {
+  background: var(--bg-raised) !important;
+  border: 0.5px solid var(--border-hi) !important;
+  color: var(--text-primary) !important;
+  min-height: 40px !important;
+  width: 100% !important;
+}
+
+#load-button button:hover,
+#generate-button button:hover {
+  background: var(--text-primary) !important;
+  color: var(--bg-base) !important;
+}
+
+#save-button button {
+  background: transparent !important;
+  border: 0.5px solid var(--border-hi) !important;
+  color: var(--text-primary) !important;
+  min-height: 38px !important;
+  width: 100% !important;
+}
+
+#save-button button:hover {
+  border-color: var(--text-primary) !important;
+}
+
+#save-button button:disabled,
+#generate-button button:disabled {
+  opacity: 0.4 !important;
+}
+
+.status-pill {
+  align-items: center;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
+  color: var(--text-secondary);
+  display: inline-flex;
+  font-size: 13px;
+  font-weight: 400;
+  gap: 8px;
+  margin: 12px 0 0;
+  padding: 8px 10px;
 }
 
-.empty-mark,
-.selected-mark {
+.status-pill span {
+  background: var(--text-muted);
   border-radius: 50%;
   display: inline-flex;
-  height: 35px;
-  width: 35px;
+  height: 7px;
+  width: 7px;
 }
 
-.empty-mark {
-  border: 2px solid #686866;
+.status-ready span,
+.status-loading span {
+  background: var(--teal);
 }
 
-.selected-mark {
-  align-items: center;
-  background: #d9fbf1;
-  border: 2px solid #86e6cc;
-  justify-content: center;
+.stats-row {
+  grid-template-columns: repeat(4, minmax(0, 1fr));
+  margin-top: 12px;
 }
 
-.selected-mark::after {
-  border: 2px solid #178870;
-  border-radius: 2px;
-  content: "";
-  height: 9px;
-  width: 9px;
+.stat-card {
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
+  padding: 12px;
 }
 
-.model-actions {
-  margin-top: 28px;
+.stat-card strong {
+  color: var(--text-primary);
+  display: block;
+  font-size: 15px;
+  font-weight: 500;
+  line-height: 1.2;
+  margin-bottom: 4px;
 }
 
-.gradio-container button {
-  border-radius: 6px !important;
-  font-family: Space Mono, monospace !important;
-  font-weight: 700 !important;
-  transition: background 180ms ease, border-color 180ms ease, color 180ms ease, transform 180ms ease !important;
+.stat-card span {
+  color: var(--text-secondary);
+  display: block;
+  font-size: 11px;
+  font-weight: 400;
+  line-height: 1.25;
 }
 
-.gradio-container button.primary,
-.gradio-container .primary > button,
-#load-button button,
-#generate-button button,
-#save-button button {
-  background: #000 !important;
-  border: 2px solid var(--line-bright) !important;
-  color: #f7f3ec !important;
+.query-section {
+  scroll-margin-top: 24px;
 }
 
-#load-button,
-#generate-button,
-#save-button {
-  width: 100% !important;
+.chat-history {
+  background: var(--bg-surface) !important;
+  border: 0.5px solid var(--border) !important;
+  border-radius: 6px !important;
+  margin-bottom: 14px;
 }
 
-#load-button button,
-#generate-button button,
-#save-button button {
-  min-height: 62px !important;
-  width: 100% !important;
+.chat-history .bubble-wrap,
+.chat-history .message,
+.chat-history .prose {
+  font-family: Space Mono, ui-monospace, monospace !important;
+  font-size: 13px !important;
+  line-height: 1.45 !important;
 }
 
-#generate-button button {
-  font-size: 24px !important;
+.chat-history .message-row {
+  padding: 6px 8px !important;
 }
 
-#load-button button:hover,
-#generate-button button:hover,
-#save-button button:hover {
-  background: #f4f0e8 !important;
-  color: #050505 !important;
-  transform: translateY(-1px);
+.chat-history pre,
+.chat-history code {
+  font-family: Space Mono, ui-monospace, monospace !important;
+  font-size: 12px !important;
 }
 
-#save-button button {
-  font-size: 22px !important;
+.schema-context-row {
+  align-items: center;
+  gap: 8px !important;
+  margin: 8px 0 12px;
 }
 
-.status-pill {
+.schema-context {
   align-items: center;
-  border-radius: 18px;
-  display: inline-flex;
-  font-size: 21px;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
+  color: var(--text-secondary);
+  display: flex;
+  gap: 8px;
+  min-height: 30px;
+  padding: 7px 9px;
+}
+
+.schema-context.empty {
+  display: none;
+}
+
+.schema-context span {
+  color: var(--teal);
+  font-size: 11px;
   font-weight: 500;
-  gap: 12px;
-  margin: 44px 0 0;
-  padding: 9px 18px;
 }
 
-.status-pill span {
-  border-radius: 50%;
-  display: inline-flex;
-  height: 11px;
-  width: 11px;
+.schema-context code {
+  color: var(--text-secondary);
+  font-family: inherit;
+  font-size: 11px;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
 }
 
-.status-ready,
-.status-loading {
-  background: #dff8ef;
-  border: 2px solid #76e4c9;
-  color: #137a65;
+.field-label,
+.preset-label {
+  color: var(--text-secondary);
+  font-size: 11px;
+  font-weight: 500;
+  letter-spacing: 0.08em;
+  line-height: 1;
+  margin: 0 0 10px;
+  text-transform: uppercase;
 }
 
-.status-ready span,
-.status-loading span {
-  background: var(--mint);
+.composer-row {
+  align-items: stretch;
+  gap: 8px !important;
 }
 
-.status-empty {
-  background: #1f1f1d;
-  border: 2px solid #53534f;
-  color: #c4bfb6;
+#message-input {
+  flex: 1 1 auto;
 }
 
-.status-empty span {
-  background: #77736c;
+#message-input textarea {
+  min-height: 42px !important;
 }
 
-.section-divider {
-  border-top: 2px solid #e8e3db;
-  margin: 42px 0 0;
+#clear-schema-button button {
+  background: transparent !important;
+  border: 0.5px solid var(--border) !important;
+  color: var(--text-secondary) !important;
+  min-height: 30px !important;
+  min-width: 34px !important;
+  width: 34px !important;
 }
 
-.preset-label {
-  color: #f2eee6;
-  font-family: Space Mono, monospace;
-  font-size: 20px;
-  font-weight: 700;
-  margin: 0 0 10px;
+#clear-schema-button button:hover {
+  border-color: var(--border-hi) !important;
+  color: var(--text-primary) !important;
 }
 
 .preset-row-gradio {
-  gap: 8px !important;
+  gap: 6px !important;
+  margin-bottom: 10px;
 }
 
 .preset-row-gradio button {
-  background: #242422 !important;
-  border: 2px solid #62625f !important;
-  color: #d4d0c8 !important;
-  font-size: 16px !important;
-  min-height: 38px !important;
+  background: var(--bg-raised) !important;
+  border: 0.5px solid var(--border) !important;
+  border-radius: 999px !important;
+  color: var(--text-secondary) !important;
+  font-size: 11px !important;
+  min-height: 28px !important;
+  padding: 0 10px !important;
 }
 
 .preset-row-gradio button:hover {
-  border-color: #f3f0ea !important;
-  color: #f3f0ea !important;
+  border-color: var(--border-hi) !important;
+  color: var(--text-primary) !important;
 }
 
 .gradio-container .block,
@@ -905,153 +1118,140 @@ CSS = """
 
 .gradio-container label,
 .gradio-container .label-wrap span {
-  color: #f4f0e8 !important;
-  font-family: Space Mono, monospace !important;
-  font-size: 20px !important;
-  font-weight: 700 !important;
+  color: var(--text-secondary) !important;
+  font-family: Space Mono, ui-monospace, monospace !important;
+  font-size: 11px !important;
+  font-weight: 500 !important;
 }
 
 textarea,
 input,
 .cm-editor {
-  background: #242422 !important;
-  border: 2px solid #646461 !important;
-  border-radius: 10px !important;
-  color: #f2eee6 !important;
-  font-family: Space Mono, monospace !important;
-  font-size: 18px !important;
-  font-weight: 700 !important;
+  background: var(--bg-raised) !important;
+  border: 0.5px solid var(--border) !important;
+  border-radius: 6px !important;
+  color: var(--text-primary) !important;
+  font-family: Space Mono, ui-monospace, monospace !important;
+  font-size: 13px !important;
+  font-weight: 400 !important;
+  line-height: 1.45 !important;
+}
+
+textarea::placeholder,
+input::placeholder {
+  color: var(--text-muted) !important;
 }
 
 textarea {
-  min-height: 108px !important;
+  min-height: 132px !important;
+}
+
+.section-divider {
+  border-top: 0.5px solid var(--border);
+  margin: 28px 0 0;
 }
 
 .output-shell {
-  border: 2px solid var(--line-bright);
-  border-radius: 10px;
-  margin-top: 14px;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
+  margin-top: 12px;
   overflow: hidden;
 }
 
 .output-head {
   align-items: center;
-  background: #121210;
-  border-bottom: 1px solid #3f3f3c;
+  background: var(--bg-surface);
+  border-bottom: 0.5px solid var(--border);
   display: flex;
   justify-content: space-between;
-  min-height: 57px;
-  padding: 0 21px;
+  min-height: 34px;
+  padding: 0 12px;
 }
 
 .output-head span:first-child {
-  color: #c8c3bb;
-  font-family: Space Mono, monospace;
-  font-size: 17px;
-  font-weight: 700;
-}
-
-.validator-badge {
-  border-radius: 5px;
-  display: inline-flex;
-  font-size: 17px;
-  font-weight: 700;
-  line-height: 1;
-  padding: 8px 12px;
-}
-
-.validator-ok {
-  background: #edf9df;
-  color: #54822d;
-}
-
-.validator-warn {
-  background: #fff3df;
-  color: #8b5a18;
-}
-
-.validator-empty {
-  background: #272724;
-  color: #b9b3aa;
+  color: var(--text-secondary);
+  font-size: 11px;
+  font-weight: 500;
+  letter-spacing: 0.08em;
+  text-transform: uppercase;
 }
 
 .validator-detail {
-  color: var(--muted);
-  font-size: 12px;
+  color: var(--text-secondary);
+  font-size: 11px;
   margin-left: 8px;
 }
 
 .output-shell .cm-editor,
 .output-shell pre,
-.output-shell code {
+.output-shell code,
+.compare-card .cm-editor,
+.compare-card pre,
+.compare-card code {
   border: 0 !important;
+  font-size: 12px !important;
+  font-weight: 400 !important;
 }
 
 .message-box {
-  color: #cfc7bc;
-  font-size: 15px;
-  font-weight: 700;
-  min-height: 28px;
+  color: var(--text-secondary);
+  font-size: 11px;
+  font-weight: 400;
+  min-height: 24px;
   padding-top: 8px;
 }
 
 .message-error {
-  color: #f4b4a8;
+  color: var(--amber-text);
 }
 
 .message-ok {
-  color: #9ee6d1;
-}
-
-.metadata-wrap {
-  margin-top: 36px;
+  color: var(--teal);
 }
 
 .comparison-panel {
-  margin-top: 40px;
-}
-
-.compare-grid {
-  display: grid;
-  gap: 18px;
-  grid-template-columns: repeat(2, minmax(0, 1fr));
+  margin-top: 28px;
 }
 
 .compare-card {
-  border: 2px solid var(--line-bright);
-  border-radius: 10px;
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border);
+  border-radius: 6px;
   overflow: hidden;
 }
 
 .compare-card.current {
-  border-color: #6de4c8;
+  border-color: rgba(29, 158, 117, 0.45);
 }
 
 .compare-head {
   align-items: center;
-  background: #131311;
-  color: #c8c3bb;
+  background: var(--amber-soft);
+  color: var(--amber-text);
   display: flex;
-  font-size: 18px;
-  font-weight: 700;
-  gap: 28px;
-  min-height: 62px;
-  padding: 0 22px;
+  font-size: 11px;
+  font-weight: 500;
+  gap: 16px;
+  justify-content: space-between;
+  min-height: 34px;
+  padding: 0 12px;
 }
 
 .compare-card.current .compare-head,
 .current-compare-head .compare-head {
-  background: #dff8ef;
-  color: #147963;
+  background: var(--teal-soft);
+  color: var(--teal-text);
 }
 
 .compare-head strong {
   color: inherit;
+  font-weight: 500;
 }
 
 .loading-overlay {
   align-items: center;
-  background: rgba(0, 0, 0, 0.72);
+  background: rgba(0, 0, 0, 0.6);
   bottom: 0;
   display: flex;
   justify-content: center;
@@ -1067,56 +1267,63 @@ textarea {
 }
 
 .loading-card {
-  background: #242422;
-  border: 2px solid #75e3c9;
-  border-radius: 10px;
-  max-width: 620px;
-  padding: 28px;
-  width: min(92vw, 620px);
+  background: var(--bg-surface);
+  border: 0.5px solid var(--border-hi);
+  border-radius: 6px;
+  max-width: 480px;
+  padding: 18px;
+  width: min(90vw, 480px);
 }
 
-.loading-card h2 {
-  color: #f4f0e8;
-  font-size: 25px;
-  margin: 0 0 18px;
+.loading-title {
+  color: var(--text-primary);
+  font-size: 13px;
+  font-weight: 500;
+  line-height: 1.35;
+  margin-bottom: 12px;
 }
 
 .loading-line {
-  background: #111110;
+  background: var(--bg-raised);
   border-radius: 999px;
-  height: 16px;
+  height: 8px;
   overflow: hidden;
 }
 
 .loading-line span {
   animation: loadingPulse 1.3s infinite ease-in-out;
-  background: #19b991;
+  background: var(--teal);
   border-radius: inherit;
   display: block;
   height: 100%;
-  width: 52%;
+  width: 45%;
 }
 
 .loading-card p {
-  color: #d2ccc2;
-  font-size: 17px;
-  font-weight: 700;
-  margin: 14px 0 0;
+  color: var(--text-secondary);
+  font-size: 11px;
+  font-weight: 400;
+  line-height: 1.4;
+  margin: 12px 0 0;
 }
 
 @keyframes loadingPulse {
-  0% { transform: translateX(-35%); }
-  50% { transform: translateX(40%); }
-  100% { transform: translateX(115%); }
+  0% { transform: translateX(-70%); }
+  50% { transform: translateX(70%); }
+  100% { transform: translateX(220%); }
 }
 
-@media (max-width: 820px) {
+@media (max-width: 860px) {
   .top-panel,
   .model-grid,
   .compare-grid {
     grid-template-columns: 1fr;
   }
 
+  .stats-row {
+    grid-template-columns: repeat(2, minmax(0, 1fr));
+  }
+
   .top-badges {
     justify-content: flex-start;
   }
@@ -1127,59 +1334,66 @@ textarea {
 }
 """
 
-
-blocks_kwargs = {"title": "Phi-3 Mini SQL Generator"}
-
-with gr.Blocks(**blocks_kwargs) as demo:
+with gr.Blocks(css=CSS, title="Phi-3 Mini SQL Generator") as demo:
     selected_model_key = gr.State(value=DEFAULT_MODEL_KEY)
     loaded_key_state = gr.State(value=None)
     saved_output = gr.State(value=None)
+    active_schema = gr.State(value="")
+    last_user_message = gr.State(value="")
 
     with gr.Column(elem_classes=["app-shell"]):
         loading_overlay = gr.HTML(render_loading_overlay(visible=False))
         gr.HTML(render_header())
 
-        gr.HTML(render_step("1", "Select Model"))
+        gr.HTML(render_step("01", "Model"))
         with gr.Row(elem_classes=["model-grid"]):
             base_model_card = gr.HTML(render_model_card(BASE_MODEL_KEY, DEFAULT_MODEL_KEY))
             fine_tuned_model_card = gr.HTML(render_model_card(FINE_TUNED_MODEL_KEY, DEFAULT_MODEL_KEY))
-        with gr.Row(elem_classes=["model-actions"]):
-            base_select = gr.Button("Select base", size="md")
-            fine_tuned_select = gr.Button("Select fine-tuned", size="md")
         load_button = gr.Button("Load selected model", variant="primary", elem_id="load-button")
         model_status = gr.HTML(render_status(DEFAULT_MODEL_KEY, None))
-
-        gr.HTML(render_step("2", "Query Input"))
-        gr.HTML('<div class="preset-label">SQL table schema</div>')
-        with gr.Row(elem_classes=["preset-row-gradio"]):
-            employees_preset = gr.Button("employees", size="sm")
-            orders_preset = gr.Button("orders", size="sm")
-            students_preset = gr.Button("students", size="sm")
-            products_preset = gr.Button("products", size="sm")
-            sales_preset = gr.Button("sales", size="sm")
-        schema_input = gr.Textbox(
-            label="",
-            value=PRESETS["employees"],
-            lines=4,
-            max_lines=8,
-            interactive=True,
-            show_label=False,
-        )
-        question_input = gr.Textbox(
-            label="Question",
-            value="What is the average salary per department?",
-            lines=2,
-            interactive=True,
-        )
-        generate_button = gr.Button(
-            "Generate SQL",
-            variant="primary",
-            interactive=False,
-            elem_id="generate-button",
-        )
+        model_info = gr.HTML(model_metadata(DEFAULT_MODEL_KEY))
+
+        with gr.Column(elem_id="query-section", elem_classes=["query-section"]):
+            gr.HTML(render_step("02", "Chat"))
+            chatbot_kwargs = {
+                "label": "",
+                "height": 360,
+                "show_label": False,
+                "elem_classes": ["chat-history"],
+            }
+            if "type" in inspect.signature(gr.Chatbot).parameters:
+                chatbot_kwargs["type"] = "messages"
+            chatbot = gr.Chatbot(**chatbot_kwargs)
+            gr.HTML('<div class="field-label">Schema context</div>')
+            with gr.Row(elem_classes=["preset-row-gradio", "prompt-chip-row"]):
+                employees_preset = gr.Button("employees", size="sm")
+                orders_preset = gr.Button("orders", size="sm")
+                students_preset = gr.Button("students", size="sm")
+                products_preset = gr.Button("products", size="sm")
+                sales_preset = gr.Button("sales", size="sm")
+            with gr.Row(elem_classes=["schema-context-row"]):
+                active_schema_pill = gr.HTML(render_schema_context(""))
+                clear_schema_button = gr.Button("x", size="sm", visible=False, elem_id="clear-schema-button")
+            with gr.Row(elem_classes=["composer-row"]):
+                message_input = gr.Textbox(
+                    label="",
+                    value="",
+                    placeholder="Type a message...",
+                    lines=1,
+                    max_lines=5,
+                    interactive=True,
+                    show_label=False,
+                    elem_id="message-input",
+                )
+                send_button = gr.Button(
+                    "Send",
+                    variant="primary",
+                    interactive=False,
+                    elem_id="generate-button",
+                )
 
         gr.HTML('<div class="section-divider"></div>')
-        gr.HTML(render_step("3", "Output"))
+        gr.HTML(render_step("03", "Output"))
         with gr.Column(elem_classes=["output-shell"]):
             with gr.Row(elem_classes=["output-head"]):
                 gr.HTML("<span>SQL</span>")
@@ -1194,18 +1408,11 @@ with gr.Blocks(**blocks_kwargs) as demo:
         save_button = gr.Button(
             "Save for comparison",
             interactive=False,
+            visible=False,
             elem_id="save-button",
         )
         error_output = gr.HTML(render_message())
 
-        with gr.Column(elem_classes=["metadata-wrap"]):
-            model_info = gr.Textbox(
-                label="Model metadata",
-                value=model_metadata(DEFAULT_MODEL_KEY),
-                lines=5,
-                interactive=False,
-            )
-
         with gr.Column(visible=False, elem_classes=["comparison-panel"]) as comparison_panel:
             with gr.Row(elem_classes=["compare-grid"]):
                 with gr.Column(elem_classes=["compare-card"]):
@@ -1226,22 +1433,12 @@ with gr.Blocks(**blocks_kwargs) as demo:
         students_preset,
         products_preset,
         sales_preset,
-        schema_input,
-        question_input,
-        generate_button,
+        clear_schema_button,
+        message_input,
+        send_button,
         save_button,
         error_output,
     ]
-    base_select.click(
-        select_model,
-        inputs=[gr.State(BASE_MODEL_KEY), loaded_key_state],
-        outputs=model_state_outputs,
-    )
-    fine_tuned_select.click(
-        select_model,
-        inputs=[gr.State(FINE_TUNED_MODEL_KEY), loaded_key_state],
-        outputs=model_state_outputs,
-    )
     base_model_card.click(
         select_model,
         inputs=[gr.State(BASE_MODEL_KEY), loaded_key_state],
@@ -1267,41 +1464,54 @@ with gr.Blocks(**blocks_kwargs) as demo:
             students_preset,
             products_preset,
             sales_preset,
-            schema_input,
-            question_input,
-            generate_button,
+            clear_schema_button,
+            message_input,
+            send_button,
             sql_output,
             validator_output,
             save_button,
             error_output,
             comparison_panel,
         ],
+        js=LOAD_SCROLL_JS,
     )
 
-    employees_preset.click(set_preset, inputs=gr.State("employees"), outputs=schema_input)
-    orders_preset.click(set_preset, inputs=gr.State("orders"), outputs=schema_input)
-    students_preset.click(set_preset, inputs=gr.State("students"), outputs=schema_input)
-    products_preset.click(set_preset, inputs=gr.State("products"), outputs=schema_input)
-    sales_preset.click(set_preset, inputs=gr.State("sales"), outputs=schema_input)
-
-    generate_button.click(
-        generate_sql,
-        inputs=[schema_input, question_input, loaded_key_state, saved_output],
-        outputs=[
-            sql_output,
-            validator_output,
-            save_button,
-            error_output,
-            comparison_panel,
-            saved_model_label,
-            saved_sql,
-            current_model_label,
-            current_sql,
-        ],
+    schema_context_outputs = [active_schema, active_schema_pill, clear_schema_button]
+    employees_preset.click(set_preset, inputs=gr.State("employees"), outputs=schema_context_outputs)
+    orders_preset.click(set_preset, inputs=gr.State("orders"), outputs=schema_context_outputs)
+    students_preset.click(set_preset, inputs=gr.State("students"), outputs=schema_context_outputs)
+    products_preset.click(set_preset, inputs=gr.State("products"), outputs=schema_context_outputs)
+    sales_preset.click(set_preset, inputs=gr.State("sales"), outputs=schema_context_outputs)
+    clear_schema_button.click(clear_schema_context, outputs=schema_context_outputs)
+
+    chat_generation_outputs = [
+        chatbot,
+        message_input,
+        active_schema,
+        last_user_message,
+        sql_output,
+        validator_output,
+        save_button,
+        error_output,
+        comparison_panel,
+        saved_model_label,
+        saved_sql,
+        current_model_label,
+        current_sql,
+    ]
+    send_button.click(
+        generate_response,
+        inputs=[message_input, chatbot, active_schema, loaded_key_state, saved_output],
+        outputs=chat_generation_outputs,
+    )
+    message_input.submit(
+        generate_response,
+        inputs=[message_input, chatbot, active_schema, loaded_key_state, saved_output],
+        outputs=chat_generation_outputs,
     )
     save_button.click(
         save_for_comparison,
-        inputs=[sql_output, loaded_key_state, schema_input, question_input],
+        inputs=[sql_output, loaded_key_state, active_schema, last_user_message],
         outputs=[
             saved_output,
             comparison_panel,
@@ -1321,7 +1531,4 @@ demo.queue(**queue_kwargs)
 
 
 if __name__ == "__main__":
-    launch_kwargs = {}
-    if "css" in inspect.signature(demo.launch).parameters:
-        launch_kwargs["css"] = CSS
-    demo.launch(**launch_kwargs)
+    demo.launch()