ContentAgent / ui /contentagentui.py
yetessam's picture
Increase max wait time and messages to prep users for long waits when its a cold start
b7c4eb0 verified
raw
history blame
5.41 kB
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 ----------
@staticmethod
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)