Update server.py

server.py

@@ -1,9 +1,14 @@
 from fastapi import FastAPI
 import gradio as gr
 
+from env import EmailTriageEnv
+from app import smart_agent_logic
+
 app = FastAPI()
 
-# --- API endpoints (required) ---
+# -------------------------
+# REQUIRED API ENDPOINTS (DO NOT REMOVE)
+# -------------------------
 @app.post("/reset")
 async def reset():
     return {"status": "ok"}
@@ -13,16 +18,63 @@ async def status():
     return {"status": "online"}
 
 
-# --- Gradio UI ---
+# -------------------------
+# GRADIO FUNCTION (THIS SHOWS REWARD + SCORE)
+# -------------------------
 def demo_fn(task):
-    return f"Running task: {task} (UI restored 🚀)"
+    env = EmailTriageEnv(task=task)
+    state = env.reset()
+
+    results = []
+    total_reward = 0.0
+    steps = 0
+
+    while True:
+        if state.get("done"):
+            break
+
+        desc = state["description"]
+        action = smart_agent_logic(desc)
+
+        state, reward, done, _, _ = env.step(action)
+
+        total_reward += reward
+        steps += 1
+
+        results.append(
+            f"### Step {steps}\n"
+            f"- 📧 Email: {desc}\n"
+            f"- 🤖 Action: {action}\n"
+            f"- ⭐ Reward: {reward:.2f}\n"
+            f"---\n"
+        )
+
+        if done:
+            break
+
+    score = total_reward / steps if steps > 0 else 0.0
+
+    return f"""
+# 📊 Email Triage Results
+
+{''.join(results)}
+
+## 🏁 Final Score: **{score:.3f} / 1.000**
+"""
+
 
+# -------------------------
+# GRADIO UI
+# -------------------------
 demo = gr.Interface(
     fn=demo_fn,
-    inputs=gr.Dropdown(["easy", "medium", "hard"]),
-    outputs="text",
-    title="Email Gatekeeper"
+    inputs=gr.Dropdown(["easy", "medium", "hard"], label="Select Difficulty"),
+    outputs=gr.Markdown(label="Results"),
+    title="📧 Email Gatekeeper",
+    description="AI Agent for Email Triage using Reinforcement Learning"
 )
 
-# ✅ Mount UI at root
+# -------------------------
+# MOUNT UI
+# -------------------------
 app = gr.mount_gradio_app(app, demo, path="/")