Update app.py

app.py

@@ -1,10 +1,5 @@
 """
-SchoolSpirit AI – Granite‑3.3‑2B chatbot Space
-----------------------------------------------
-• IBM Granite‑3.3‑2B‑Instruct (Apache‑2), runs in HF CPU Space.
-• Keeps last MAX_TURNS exchanges to fit context.
-• “Clear Chat” button resets conversation.
-• Extensive error‑handling: model‑load, inference, bad input.
+SchoolSpirit AI – Granite‑3.3‑2B chatbot Space (messages API)
 """
 
 import re
@@ -16,14 +11,14 @@ from transformers import (
 )
 from transformers.utils import logging as hf_logging
 
-# ────────── Configuration ───────────────────────────────────────────────────
+# ─── Config ────────────────────────────────────────────────────────────────
 hf_logging.set_verbosity_error()
 LOG = hf_logging.get_logger("SchoolSpirit")
 
 MODEL_ID     = "ibm-granite/granite-3.3-2b-instruct"
-MAX_TURNS    = 6          # history turns to keep
-MAX_TOKENS   = 200        # response length
-MAX_INPUT_CH = 400        # user message length guard
+MAX_TURNS    = 6          # keep last N user/assistant pairs
+MAX_TOKENS   = 200
+MAX_INPUT_CH = 400
 
 SYSTEM_MSG = (
     "You are SchoolSpirit AI, the upbeat digital mascot for a company that "
@@ -32,13 +27,11 @@ SYSTEM_MSG = (
     "say so and suggest contacting a human. Never request personal data."
 )
 
-# ────────── Model loading with fail‑safe ────────────────────────────────────
+# ─── Model load with fail‑safe ─────────────────────────────────────────────
 try:
     tok = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        device_map="auto",
-        torch_dtype="auto",
+        MODEL_ID, device_map="auto", torch_dtype="auto"
     )
     generator = pipeline(
         "text-generation",
@@ -54,77 +47,77 @@ except Exception as exc:  # noqa: BLE001
     generator = None
     LOG.error(MODEL_ERR)
 
-# ────────── Helper utilities ────────────────────────────────────────────────
-def truncate(hist):
-    """Return last MAX_TURNS (user,bot) tuples."""
-    return hist[-MAX_TURNS:] if len(hist) > MAX_TURNS else hist
-
-
+# ─── Helpers ───────────────────────────────────────────────────────────────
 def clean(text: str) -> str:
-    """Normalize whitespace and guarantee non‑empty."""
     return re.sub(r"\s+", " ", text.strip()) or "…"
 
 
-def safe_generate(prompt: str) -> str:
-    """Call model.generate, catch & log any error, always return a string."""
-    try:
-        completion = generator(prompt)[0]["generated_text"]
-        reply = clean(completion.split("AI:", 1)[-1])
-    except Exception as err:  # noqa: BLE001
-        LOG.error(f"Inference error: {err}")
-        reply = (
-            "Sorry—I'm having trouble right now. "
-            "Please try again in a moment."
-        )
-    return reply
+def trim(messages):
+    """Keep system + last MAX_TURNS*2 messages."""
+    if len(messages) <= 1 + MAX_TURNS * 2:
+        return messages
+    # messages[0] is system; keep last N pairs
+    return [messages[0]] + messages[-MAX_TURNS * 2:]
 
-# ────────── Chat callback ───────────────────────────────────────────────────
-def chat(history, user_msg):
-    history = list(history)  # guaranteed list of tuples
+# ─── Chat callback (messages API) ──────────────────────────────────────────
+def chat_fn(messages):
+    """
+    messages: list[dict] with keys 'role' ('system'/'user'/'assistant')
+    returns the updated messages list
+    """
 
-    # Fatal start‑up failure
+    # Startup failure?
     if MODEL_ERR:
-        history.append((user_msg, MODEL_ERR))
-        return history, ""
+        messages.append(
+            {"role": "assistant", "content": MODEL_ERR}
+        )
+        return messages
 
-    user_msg = clean(user_msg or "")
+    # Validate user input
+    user_msg = messages[-1]["content"]
+    user_msg = clean(user_msg)
     if not user_msg:
-        history.append(("", "Please enter a message."))
-        return history, ""
+        messages[-1]["content"] = "Please enter a message."
+        return messages
     if len(user_msg) > MAX_INPUT_CH:
-        history.append(
-            (user_msg, f"Message too long (>{MAX_INPUT_CH} chars).")
+        messages[-1]["content"] = (
+            f"Message too long (> {MAX_INPUT_CH} chars)."
         )
-        return history, ""
+        return messages
 
-    history = truncate(history)
+    messages = trim(messages)
 
     # Build prompt
-    prompt_lines = [SYSTEM_MSG]
-    for u, a in history:
-        prompt_lines += [f"User: {u}", f"AI: {a}"]
-    prompt_lines += [f"User: {user_msg}", "AI:"]
-    prompt = "\n".join(prompt_lines)
-
-    reply = safe_generate(prompt)
-
-    history.append((user_msg, reply))
-    return history, ""
-
-# ────────── Clear chat callback ─────────────────────────────────────────────
-def clear_chat():
-    return [], ""
-
-# ────────── UI definition ───────────────────────────────────────────────────
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
-    gr.Markdown("# SchoolSpirit AI Chat")
-    chatbot   = gr.Chatbot(type="tuples")
-    msg_box   = gr.Textbox(placeholder="Ask me anything about SchoolSpirit AI…")
-    send_btn  = gr.Button("Send")
-    clear_btn = gr.Button("Clear Chat", variant="secondary")
-
-    send_btn.click(chat, [chatbot, msg_box], [chatbot, msg_box])
-    msg_box.submit(chat, [chatbot, msg_box], [chatbot, msg_box])
-    clear_btn.click(clear_chat, outputs=[chatbot, msg_box])
-
-demo.launch()
+    prompt = ""
+    for m in messages:
+        if m["role"] == "system":
+            prompt += m["content"] + "\n"
+        elif m["role"] == "user":
+            prompt += f"User: {m['content']}\n"
+        else:  # assistant
+            prompt += f"AI: {m['content']}\n"
+    prompt += "AI:"
+
+    # Generate
+    try:
+        out = generator(prompt)[0]["generated_text"]
+        reply = clean(out.split("AI:", 1)[-1])
+    except Exception as err:  # noqa: BLE001
+        LOG.error(f"Inference error: {err}")
+        reply = (
+            "Sorry—I'm having trouble right now. "
+            "Please try again shortly."
+        )
+
+    messages.append({"role": "assistant", "content": reply})
+    return messages
+
+# ─── Gradio UI (ChatInterface handles Send & Clear) ────────────────────────
+gr.ChatInterface(
+    fn=chat_fn,
+    title="SchoolSpirit AI Chat",
+    theme=gr.themes.Soft(primary_hue="blue"),
+    system_prompt=SYSTEM_MSG,
+    chatbot=gr.Chatbot(height=480),
+    type="messages",  # modern, future‑proof format
+).launch()