Update ui/contentagentui.py

ui/contentagentui.py

@@ -1,117 +1,134 @@
-import gradio as gr
 import os
+import gradio as gr
 
+from status_check import is_endpoint_healthy      # your tiny checker
+from endpoint_utils import wake_endpoint          # your wake+poll helper
 
 
 class ContentAgentUI:
+    """
+    Gradio UI that:
+      - shows a minimal control panel first (status + Start button),
+      - wakes & inits the agent when requested,
+      - reveals the main chat panel only after the agent is ready.
+    """
 
     def __init__(self, endpoint_uri: str, is_healthy: bool, health_message: str, agent_initializer):
         self.endpoint_uri = endpoint_uri
-        self.is_healthy = is_healthy
-        self.health_message = health_message
-        self.agent_initializer = agent_initializer
-        self.agent = None
-
-        # Optional: Adjust path if needed for hosted environments
+        self.is_healthy = bool(is_healthy)
+        self.health_message = health_message or ""
+        self.agent_initializer = agent_initializer  # callable: (uri) -> CodeAgent
+
+        # set at build()
+        self.app: gr.Blocks | None = None
+        self.status_box = None
+        self.control_panel = None
+        self.main_panel = None
+        self.prompt = None
+        self.reply = None
+        self.agent_state = None  # where we store the agent object
+        # NOTE: we don't keep a separate self.agent; we read from agent_state in callbacks
+
+    # ---------- helpers ----------
+
+    def _read_css(self) -> str | None:
         css_path = os.path.join(os.getcwd(), "ui", "styles.css")
-
-        self.ca_gui = gr.Blocks(css=css_path if os.path.exists(css_path) else None)
-
-        self.sections = [
-            self.create_header,
-            self.create_user_guidance,
-            self.create_main,
-            self.create_examples,
-            self.create_footer,
-        ]
-
-        for section in self.sections:
-            section()
-
-        self.ca_gui.launch()
-
-    def call_agent(self, input_text):
+        if os.path.exists(css_path):
+            try:
+                with open(css_path, "r", encoding="utf-8") as f:
+                    return f.read()
+            except Exception:
+                return None
+        return None
+
+    def _initial_status_text(self) -> str:
+        if not self.endpoint_uri:
+            return "No endpoint URI configured."
+        if self.is_healthy:
+            return f"Endpoint healthy ✅ — {self.health_message or 'OK'}.\nClick 'Start Agent' to initialize."
+        return (
+            f"Endpoint not healthy: {self.health_message or 'unknown'}.\n"
+            "Click 'Start Agent' to wake the endpoint and initialize."
+        )
+
+    # ---------- agent call ----------
+
+    @staticmethod
+    def _call_agent(text: str, agent) -> str:
         try:
-            if self.agent:
-                return self.agent.get_response(input_text)
-            else:
-                return "Agent not loaded."
+            if agent is None:
+                return "Agent not initialized yet. Click 'Start Agent'."
+            return str(agent.run(text))  # smolagents.CodeAgent API
         except Exception as e:
-            return f"Error: {str(e)}"
-
-    def create_header(self):
-        with self.ca_gui:
-            gr.Markdown("# Content Agent")
-
-    def create_user_guidance(self):
-        with self.ca_gui:
-            gr.Markdown("""
-                Please enter text below to get started. The AI Agent will try to determine whether the language is polite and uses the following classification:
-
-                - `polite`
-                - `somewhat polite`
-                - `neutral`
-                - `impolite`
-
-                Technology:
+            return f"Error: {e}"
 
-                - App is running `deepseek-ai/DeepSeek-R1-Distill-Qwen-32B` text generation model.
-                - Agent uses Intel's Polite Guard NLP library tool
-                - Compute on GCP · Nvidia L4 · 4x GPUs · 96 GB
-            """)
+    # ---------- UI build ----------
 
-    def create_main(self):
-        with self.ca_gui:
-            with gr.Row():
-                with gr.Column():
-                    self.user_input = gr.Textbox(label="Your Input", placeholder="Enter something here...")
-                    self.submit_button = gr.Button("Submit")
-                    self.output = gr.Textbox(label="Content feedback", interactive=False, lines=10, max_lines=20)
-
-                    # Use bound method with `self`
-                    self.submit_button.click(self.call_agent, inputs=self.user_input, outputs=self.output)
-                    self.user_input.submit(self.call_agent, inputs=self.user_input, outputs=self.output)
-
-    def get_example(self):
-        example_root = os.path.join(os.path.dirname(__file__), "examples")
-        examples = []
-
-        if os.path.exists(example_root):
-            example_files = [os.path.join(example_root, f) for f in os.listdir(example_root) if f.endswith(".txt")]
-
-            for file_path in example_files:
-                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
-                    examples.append(f.read())
-
-        return examples
-
-    def create_examples(self):
-        examples = self.get_example()
-
-        with self.ca_gui:
-            if examples:
-                example_radio = gr.Radio(choices=examples, label="Try one of these examples:")
-                example_radio.change(fn=lambda ex: ex, inputs=example_radio, outputs=self.user_input)
-            else:
-                gr.Markdown("*No examples found.*")
-
-    def create_footer(self):
-        with self.ca_gui:
-            gr.Markdown("<div id='footer'>Thanks for trying it out!</div>")
-
-
-    
-    def launch(self):
-        # build gradio Blocks:
-        # - status textbox prefilled with health_message
-        # - “Wake Endpoint” button → calls wake logic, updates status
-        # - “Initialize Agent” button → if healthy, self.agent = agent_initializer(endpoint_uri); reveal main panel
-        ...
-
-
-    def pass_through_agent(self, agent):
-        # Assign the agent for future use in `call_agent`
-        self.agent = agent
+    def build(self) -> gr.Blocks:
+        if self.app is not None:
+            return self.app
 
+        css = self._read_css()
+        with gr.Blocks(css=css) as demo:
+            gr.Markdown("# Content Agent")
 
-        
+            # Control panel (shown first)
+            with gr.Group(visible=True) as self.control_panel:
+                self.status_box = gr.Textbox(
+                    label="Status",
+                    value=self._initial_status_text(),
+                    lines=8,
+                    interactive=False,
+                )
+                start_btn = gr.Button("Start Agent (wake if needed)")
+
+            # Main panel (hidden until agent is initialized)
+            with gr.Group(visible=False) as self.main_panel:
+                self.agent_state = gr.State(None)
+                self.prompt = gr.Textbox(label="Your Input", placeholder="Enter something here…")
+                self.reply = gr.Textbox(label="Content feedback", interactive=False, lines=12)
+                submit_btn = gr.Button("Submit")
+
+                submit_btn.click(self._call_agent, inputs=[self.prompt, self.agent_state], outputs=self.reply)
+                self.prompt.submit(self._call_agent, inputs=[self.prompt, self.agent_state], outputs=self.reply)
+
+            # Event: Start Agent (wake → health-check → init → reveal main panel)
+            def on_start():
+                # 1) show we're waking (stream status)
+                yield ("Waking endpoint…", gr.update(), gr.update(), None)
+
+                ok, err = wake_endpoint(self.endpoint_uri, max_wait=180, poll_every=5.0)
+                if not ok:
+                    # wake failed → keep control panel visible
+                    yield (f"Wake failed: {err or 'unknown error'}", gr.update(), gr.update(), None)
+                    return
+
+                # 2) check health
+                yield ("Endpoint awake ✅. Checking health…", gr.update(), gr.update(), None)
+                healthy, msg = is_endpoint_healthy(self.endpoint_uri)
+                if not healthy:
+                    yield (f"Endpoint not healthy: {msg}", gr.update(), gr.update(), None)
+                    return
+
+                # 3) initialize agent
+                yield ("Initializing agent…", gr.update(), gr.update(), None)
+                try:
+                    agent = self.agent_initializer(self.endpoint_uri)
+                except Exception as e:
+                    yield (f"Agent init failed: {e}", gr.update(), gr.update(), None)
+                    return
+
+                # 4) success → hide control panel, show main
+                yield ("Agent initialized ✅", gr.update(visible=False), gr.update(visible=True), agent)
+
+            # wire the button: outputs are (status_text, control_panel_update, main_panel_update, agent_state)
+            start_btn.click(on_start, inputs=None, outputs=[self.status_box, self.control_panel, self.main_panel, self.agent_state])
+
+            self.app = demo
+
+        return self.app
+
+    # ---------- public API ----------
+
+    def launch(self, **kwargs):
+        return self.build().launch(**kwargs)