Spaces:
Sleeping
Sleeping
| 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 = 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") | |
| 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 ---------- | |
| def _call_agent(text: str, agent) -> str: | |
| try: | |
| 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: {e}" | |
| # ---------- UI build ---------- | |
| 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(): | |
| lines = [] | |
| def push(s): | |
| lines.append(s) | |
| return ("\n".join(lines), gr.update(), gr.update(), None) | |
| yield push("Waking endpoint… (this can take several minutes for cold starts)") | |
| ok, err = wake_endpoint(self.endpoint_uri, max_wait=600, poll_every=5.0, log=lambda m: lines.append(m)) | |
| yield push(lines[-1] if lines else "…") # flush last log line | |
| if not ok: | |
| yield push(f"[Server message] {err or 'wake failed'}") | |
| return | |
| yield push("Endpoint awake ✅. Checking health…") | |
| healthy, msg = is_endpoint_healthy(self.endpoint_uri) | |
| if not healthy: | |
| yield push(f"[Server message] {msg}") | |
| return | |
| yield push("Initializing agent…") | |
| try: | |
| agent = self.agent_initializer(self.endpoint_uri) | |
| except Exception as e: | |
| yield push(f"Agent init failed: {e}") | |
| return | |
| 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) | |