garvitsachdeva/SpindleFlow-RL
0
1"""2Animated robot orchestrator widget for the SpindleFlow RL demo.3Exports one public function: render_orchestrator(state, height=600)4 5All HTML/CSS/JS is self-contained — no CDN, no external calls.6Safe for Hugging Face Spaces iframe sandbox.7"""8 9from __future__ import annotations10import json11import math12 13# ── Agent color and icon maps ─────────────────────────────────────────────────14 15SPEC_COLORS = {16 "frontend_react": "#00d4ff",17 "backend_api": "#7c3aed",18 "database_architect": "#f59e0b",19 "devops_engineer": "#10b981",20 "security_analyst": "#ef4444",21 "product_strategist": "#8b5cf6",22 "ux_designer": "#ec4899",23 "tech_writer": "#94a3b8",24}25 26SPEC_ICONS = {27 "frontend_react": "FE",28 "backend_api": "API",29 "database_architect": "DB",30 "devops_engineer": "OPS",31 "security_analyst": "SEC",32 "product_strategist": "PM",33 "ux_designer": "UX",34 "tech_writer": "DOC",35}36 37_SPAWNED_COLOR = "#fbbf24" # gold for auto-spawned agents38_FALLBACK_COLORS = [ # cycle through for multiple unknown agents39 "#fbbf24", "#f472b6", "#34d399", "#fb923c", "#a78bfa",40]41 42 43def _agent_color(agent_id: str, spawned_ids: set) -> str:44 if agent_id in SPEC_COLORS:45 return SPEC_COLORS[agent_id]46 if agent_id in spawned_ids:47 return _SPAWNED_COLOR48 # deterministic fallback based on hash49 return _FALLBACK_COLORS[hash(agent_id) % len(_FALLBACK_COLORS)]50 51 52def _agent_icon(agent_id: str, spawned_ids: set) -> str:53 if agent_id in SPEC_ICONS:54 return SPEC_ICONS[agent_id]55 if agent_id in spawned_ids:56 return "⚡"57 return agent_id[:3].upper()58 59 60# ── Layout ────────────────────────────────────────────────────────────────────61 62def _agent_positions(agent_ids: list,63 canvas_w: int = 780,64 canvas_h: int = 560) -> dict:65 """Return {agent_id: (x, y)} in a straight vertical column on the right."""66 col_x = canvas_w - 11567 n = len(agent_ids)68 if n == 0:69 return {}70 pad_top = 5071 pad_bot = 5072 usable = canvas_h - pad_top - pad_bot73 step = usable / n74 positions = {}75 for i, aid in enumerate(agent_ids):76 y = round(pad_top + step * i + step / 2)77 positions[aid] = (col_x, y)78 return positions79 80 81# ── SVG builders ──────────────────────────────────────────────────────────────82 83def _robot_svg() -> str:84 return """85 <g id="robot" transform="translate(160, 280)">86 87 <!-- Antenna -->88 <line x1="0" y1="-115" x2="0" y2="-95" stroke="#00d4ff" stroke-width="2"/>89 <circle cx="0" cy="-120" r="5" fill="#00d4ff" class="antenna-pulse"/>90 91 <!-- Head -->92 <rect x="-38" y="-95" width="76" height="62" rx="10"93 fill="#0d1117" stroke="#00d4ff" stroke-width="1.5"94 class="head-glow"/>95 96 <!-- Left Eye -->97 <circle cx="-14" cy="-68" r="10" fill="#001a2e"/>98 <circle cx="-14" cy="-68" r="6" fill="#00d4ff" class="eye-left"/>99 <circle cx="-11" cy="-71" r="2" fill="white" opacity="0.6"/>100 101 <!-- Right Eye -->102 <circle cx="14" cy="-68" r="10" fill="#001a2e"/>103 <circle cx="14" cy="-68" r="6" fill="#00d4ff" class="eye-right"/>104 <circle cx="17" cy="-71" r="2" fill="white" opacity="0.6"/>105 106 <!-- Mouth -->107 <path d="M -14 -46 Q 0 -38 14 -46"108 fill="none" stroke="#00d4ff" stroke-width="2"109 stroke-linecap="round" class="mouth"/>110 111 <!-- Neck -->112 <rect x="-8" y="-33" width="16" height="10" rx="3"113 fill="#0d1117" stroke="#1a2a3a" stroke-width="1"/>114 115 <!-- Body -->116 <rect x="-45" y="-23" width="90" height="80" rx="12"117 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1.5"/>118 119 <!-- Core (spinning hexagon) -->120 <g class="core-spin" transform="translate(0, 17)">121 <polygon points="0,-18 15.6,-9 15.6,9 0,18 -15.6,9 -15.6,-9"122 fill="none" stroke="#00d4ff" stroke-width="1.5" opacity="0.8"/>123 <polygon points="0,-11 9.5,-5.5 9.5,5.5 0,11 -9.5,5.5 -9.5,-5.5"124 fill="rgba(0,212,255,0.15)" stroke="#00d4ff" stroke-width="1"/>125 <circle cx="0" cy="0" r="4" fill="#00d4ff" class="core-pulse"/>126 </g>127 128 <!-- Left Arm -->129 <g id="arm-left">130 <rect x="-68" y="-18" width="24" height="12" rx="6"131 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1.5"/>132 <rect x="-72" y="-8" width="14" height="28" rx="7"133 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1.5"/>134 </g>135 136 <!-- Right Arm -->137 <g id="arm-right" class="arm-idle">138 <rect x="44" y="-18" width="24" height="12" rx="6"139 fill="#0a0f1a" stroke="#00d4ff" stroke-width="1.5"/>140 <rect x="58" y="-8" width="14" height="28" rx="7"141 fill="#0a0f1a" stroke="#00d4ff" stroke-width="1.5"/>142 <circle cx="65" cy="22" r="5" fill="#00d4ff" class="hand-glow"/>143 </g>144 145 <!-- Legs -->146 <rect x="-28" y="57" width="18" height="28" rx="6"147 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1.5"/>148 <rect x="10" y="57" width="18" height="28" rx="6"149 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1.5"/>150 151 <!-- Feet -->152 <ellipse cx="-19" cy="87" rx="16" ry="7"153 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1"/>154 <ellipse cx="19" cy="87" rx="16" ry="7"155 fill="#0a0f1a" stroke="#1a3a5a" stroke-width="1"/>156 157 <!-- Shadow -->158 <ellipse cx="0" cy="97" rx="50" ry="8"159 fill="rgba(0,212,255,0.05)"/>160 </g>161 """162 163 164def _agent_card_svg(agent_id: str, x: int, y: int,165 status: str, color: str,166 is_spawned: bool = False) -> str:167 """Returns SVG <g> for one agent card. status: idle | active | done."""168 icon = SPEC_ICONS.get(agent_id, ("⚡" if is_spawned else agent_id[:3].upper()))169 label = agent_id.replace("_", " ").title()170 label = label[:18] + ("…" if len(label) > 18 else "")171 172 status_class = {"idle": "agent-idle", "active": "agent-active",173 "done": "agent-done"}.get(status, "agent-idle")174 opacity = "1.0" if status != "idle" else "0.40"175 border = "#fbbf24" if is_spawned else color176 spawn_star = (177 f'<text x="26" y="-26" text-anchor="middle" font-size="10" fill="#fbbf24">⚡</text>'178 if is_spawned else ""179 )180 181 return f"""182 <g class="agent-card {status_class}" transform="translate({x},{y})"183 id="agent-{agent_id}" opacity="{opacity}">184 <circle cx="0" cy="0" r="36" fill="none"185 stroke="{border}" stroke-width="1.5"186 class="agent-ring" opacity="0.25"/>187 <rect x="-26" y="-26" width="52" height="52" rx="10"188 fill="#0a0f1a" stroke="{border}" stroke-width="1.5"189 opacity="0.95"/>190 <text x="0" y="5" text-anchor="middle" dominant-baseline="middle"191 fill="{color}" font-family="'JetBrains Mono', monospace"192 font-size="11" font-weight="700">{icon}</text>193 {spawn_star}194 <circle cx="20" cy="-20" r="5" fill="{color}" class="status-dot"/>195 <text x="0" y="40" text-anchor="middle"196 fill="#64748b" font-family="system-ui, sans-serif"197 font-size="8.5" letter-spacing="0.3">{label}</text>198 <g class="done-check" opacity="0">199 <circle cx="20" cy="-20" r="7" fill="#10b981"/>200 <text x="20" y="-16" text-anchor="middle" fill="white" font-size="9">✓</text>201 </g>202 </g>203 """204 205 206def _beam_svg(edges: list, agent_positions: dict) -> str:207 """Returns SVG beam lines for all current delegation edges."""208 robot_hand_x, robot_hand_y = 225, 302209 lines = []210 for caller, callee in edges:211 if callee not in agent_positions:212 continue213 tx, ty = agent_positions[callee]214 color = SPEC_COLORS.get(callee, _SPAWNED_COLOR)215 lines.append(f"""216 <line id="beam-{callee}"217 x1="{robot_hand_x}" y1="{robot_hand_y}" x2="{tx}" y2="{ty}"218 stroke="{color}" stroke-width="1.5" stroke-linecap="round"219 opacity="0.55" stroke-dasharray="6 4" class="beam-line beam-animate"/>220 <circle id="dot-{callee}" r="4" fill="{color}" opacity="0.9" class="beam-dot">221 <animateMotion dur="0.9s" repeatCount="indefinite"222 path="M {robot_hand_x},{robot_hand_y} L {tx},{ty}"/>223 </circle>224 <circle id="burst-{callee}" cx="{tx}" cy="{ty}" r="8"225 fill="none" stroke="{color}" stroke-width="2"226 opacity="0" class="burst-ring burst-animate"/>227 """)228 return "\n".join(lines)229 230 231# ── HTML template ─────────────────────────────────────────────────────────────232 233def _html_template(*, agents_svg, beams_svg, robot_svg, state_json,234 task_short, reward_html, step, phase, mode, mode_color) -> str:235 return f"""<!DOCTYPE html>236<html>237<head>238<meta charset="utf-8"/>239<style>240 * {{ box-sizing: border-box; margin: 0; padding: 0; }}241 body {{ background: transparent; font-family: 'JetBrains Mono', 'Fira Code', monospace; overflow: hidden; }}242 243 .canvas-wrap {{244 position: relative; width: 100%; height: 560px;245 background: radial-gradient(ellipse at 25% 50%, rgba(0,212,255,0.04) 0%, transparent 60%),246 radial-gradient(ellipse at 85% 50%, rgba(124,58,237,0.03) 0%, transparent 50%),247 #080d14;248 border-radius: 16px; border: 1px solid rgba(0,212,255,0.1); overflow: hidden;249 }}250 .canvas-wrap::before {{251 content: ''; position: absolute; inset: 0;252 background-image: linear-gradient(rgba(0,212,255,0.025) 1px, transparent 1px),253 linear-gradient(90deg, rgba(0,212,255,0.025) 1px, transparent 1px);254 background-size: 40px 40px; border-radius: 16px; pointer-events: none;255 }}256 svg.main-svg {{ position: absolute; top: 0; left: 0; width: 100%; height: 100%; }}257 258 .info-bar {{259 position: absolute; bottom: 0; left: 0; right: 0; height: 44px;260 background: rgba(0,0,0,0.5); border-top: 1px solid rgba(255,255,255,0.05);261 border-radius: 0 0 16px 16px; display: flex; align-items: center;262 padding: 0 20px; gap: 24px; font-size: 11px; color: #475569;263 }}264 .info-badge {{ display: flex; align-items: center; gap: 6px; }}265 .info-badge .label {{ font-size: 9px; text-transform: uppercase; letter-spacing: 1px; color: #334155; }}266 .info-badge .value {{ font-weight: 700; color: #94a3b8; }}267 .task-text {{ flex: 1; overflow: hidden; white-space: nowrap; text-overflow: ellipsis; color: #475569; font-size: 10px; }}268 269 .orch-label {{ position: absolute; top: 18px; left: 18px; font-size: 9px; font-weight: 700; text-transform: uppercase; letter-spacing: 2px; color: #00d4ff; opacity: 0.7; }}270 .agents-label {{ position: absolute; top: 18px; right: 18px; font-size: 9px; font-weight: 700; text-transform: uppercase; letter-spacing: 2px; color: #475569; opacity: 0.7; }}271 272 .divider-line {{273 position: absolute; left: 47%; top: 8%; height: 84%; width: 1px;274 background: linear-gradient(to bottom, transparent, rgba(0,212,255,0.12), transparent);275 }}276 277 /* Robot animations */278 @keyframes antenna-blink {{ 0%,90%,100% {{ opacity:1; }} 95% {{ opacity:0.2; }} }}279 .antenna-pulse {{ animation: antenna-blink 2.5s ease-in-out infinite; }}280 281 @keyframes core-rotation {{ from {{ transform: rotate(0deg); }} to {{ transform: rotate(360deg); }} }}282 .core-spin {{ transform-origin: 0px 17px; animation: core-rotation 4s linear infinite; }}283 284 @keyframes core-pulse {{ 0%,100% {{ opacity:0.8; r:4px; }} 50% {{ opacity:1; r:6px; fill:white; }} }}285 .core-pulse {{ animation: core-pulse 1.5s ease-in-out infinite; }}286 287 @keyframes eye-blink {{ 0%,92%,100% {{ ry:6px; }} 96% {{ ry:1px; }} }}288 .eye-left, .eye-right {{ animation: eye-blink 4s ease-in-out infinite; transform-box: fill-box; transform-origin: center; }}289 290 @keyframes hand-glow {{ 0%,100% {{ opacity:0.6; r:5px; }} 50% {{ opacity:1; r:8px; }} }}291 .hand-glow {{ animation: hand-glow 1.2s ease-in-out infinite; }}292 293 @keyframes head-glow-pulse {{ 0%,100% {{ filter: drop-shadow(0 0 4px rgba(0,212,255,0.3)); }} 50% {{ filter: drop-shadow(0 0 12px rgba(0,212,255,0.7)); }} }}294 .head-glow {{ animation: head-glow-pulse 2s ease-in-out infinite; }}295 296 @keyframes arm-extend {{ 0% {{ transform: rotate(0deg) translateX(0px); }} 100% {{ transform: rotate(-15deg) translateX(12px); }} }}297 .arm-delegating {{ transform-origin: 55px 0px; animation: arm-extend 0.4s ease-out forwards; }}298 299 /* Agent animations */300 @keyframes agent-active-pulse {{ 0%,100% {{ filter: drop-shadow(0 0 6px currentColor); }} 50% {{ filter: drop-shadow(0 0 18px currentColor); }} }}301 .agent-active {{ animation: agent-active-pulse 0.8s ease-in-out infinite; opacity: 1 !important; }}302 .agent-done {{ opacity: 1 !important; }}303 .agent-done .status-dot {{ fill: #10b981 !important; }}304 .agent-done .done-check {{ opacity: 1 !important; }}305 306 @keyframes ring-expand {{ from {{ r:28px; opacity:0.6; }} to {{ r:48px; opacity:0; }} }}307 .agent-active .agent-ring {{ animation: ring-expand 1s ease-out infinite; }}308 309 /* Beam animations */310 @keyframes beam-draw {{ from {{ stroke-dashoffset:200; opacity:0; }} to {{ stroke-dashoffset:0; opacity:0.55; }} }}311 .beam-animate {{ stroke-dasharray: 6 4; animation: beam-draw 0.4s ease-out forwards; }}312 313 @keyframes burst-expand {{ 0% {{ r:8px; opacity:0.9; stroke-width:3px; }} 100% {{ r:28px; opacity:0; stroke-width:1px; }} }}314 .burst-animate {{ animation: burst-expand 0.6s ease-out infinite; }}315 316 .robot-thinking .core-spin {{ animation-duration: 1.2s !important; }}317 .robot-thinking .antenna-pulse {{ animation: antenna-blink 0.6s ease-in-out infinite !important; }}318 319 /* Sequential reveal */320 @keyframes slide-in-right {{321 from {{ opacity: 0; transform: translateX(22px); }}322 to {{ opacity: 1; transform: translateX(0); }}323 }}324 325 #particles {{ position: absolute; top: 0; left: 0; width: 100%; height: 560px; pointer-events: none; }}326</style>327</head>328<body>329<div class="canvas-wrap" id="canvas-wrap">330 <canvas id="particles"></canvas>331 <div class="orch-label">Orchestrator</div>332 <div class="agents-label">Specialists</div>333 <div class="divider-line"></div>334 335 <svg class="main-svg" viewBox="0 0 780 560" xmlns="http://www.w3.org/2000/svg">336 <g id="beams-layer">{beams_svg}</g>337 <g id="agents-layer">{agents_svg}</g>338 <g id="robot-layer">{robot_svg}</g>339 </svg>340 341 <div class="info-bar">342 <div class="info-badge">343 <span class="label">Step</span>344 <span class="value">{step}</span>345 </div>346 <div class="info-badge">347 <span class="label">Phase</span>348 <span class="value">{phase}</span>349 </div>350 <div class="info-badge">351 <span class="label">Mode</span>352 <span class="value" style="color:{mode_color};">{mode}</span>353 </div>354 <div class="info-badge">355 <span class="label">Reward</span>356 <span class="value">{reward_html}</span>357 </div>358 <div class="task-text" title="{task_short}">{task_short}</div>359 </div>360</div>361 362<script>363const STATE = {state_json};364 365const robotLayer = document.getElementById('robot-layer');366const armRight = document.getElementById('arm-right');367 368if (STATE.robot_state === 'thinking' || STATE.robot_state === 'delegating') {{369 robotLayer.classList.add('robot-thinking');370}}371if (STATE.robot_state === 'delegating' && armRight) {{372 armRight.classList.remove('arm-idle');373 armRight.classList.add('arm-delegating');374}}375 376// Sequential reveal: agents appear one-by-one with staggered delays377if (STATE.mode === 'SEQUENTIAL' && !STATE.done && STATE.called.length > 0) {{378 STATE.called.forEach(function(agentId, idx) {{379 var el = document.getElementById('agent-' + agentId);380 if (!el) return;381 el.style.opacity = '0';382 (function(element, delay) {{383 setTimeout(function() {{384 element.style.transition = 'opacity 0.5s ease';385 element.style.opacity = '1';386 }}, delay);387 }})(el, 250 + idx * 650);388 }});389}}390 391function spawnParticles(x, y, color) {{392 const canvas = document.getElementById('particles');393 if (!canvas) return;394 const ctx = canvas.getContext('2d');395 canvas.width = canvas.offsetWidth;396 canvas.height = canvas.offsetHeight;397 const particles = [];398 for (let i = 0; i < 18; i++) {{399 const angle = (Math.PI * 2 * i) / 18;400 const speed = 1.5 + Math.random() * 2.5;401 particles.push({{ x, y, vx: Math.cos(angle)*speed, vy: Math.sin(angle)*speed, life: 1.0, r: 2+Math.random()*2, color }});402 }}403 function animate() {{404 ctx.clearRect(0, 0, canvas.width, canvas.height);405 let alive = false;406 particles.forEach(p => {{407 if (p.life <= 0) return;408 p.x += p.vx; p.y += p.vy; p.vx *= 0.92; p.vy *= 0.92; p.life -= 0.025; alive = true;409 ctx.beginPath(); ctx.arc(p.x, p.y, p.r, 0, Math.PI*2);410 ctx.fillStyle = color + Math.floor(p.life*255).toString(16).padStart(2,'0');411 ctx.fill();412 }});413 if (alive) requestAnimationFrame(animate);414 else ctx.clearRect(0, 0, canvas.width, canvas.height);415 }}416 animate();417}}418 419if (STATE.active) {{420 const activeEl = document.getElementById('agent-' + STATE.active);421 if (activeEl) {{422 const wrap = document.getElementById('canvas-wrap');423 const wRect = wrap.getBoundingClientRect();424 const ct = activeEl.getCTM();425 if (ct) {{426 const scaleX = wRect.width / 780;427 const scaleY = wRect.height / 560;428 const tx = ct.e * scaleX;429 const ty = ct.f * scaleY;430 const rect = activeEl.querySelector('rect');431 const agentColor = rect ? rect.getAttribute('stroke') : '#00d4ff';432 setTimeout(() => spawnParticles(tx, ty, agentColor), 300);433 }}434 }}435}}436 437(function breathe() {{438 const robot = document.getElementById('robot');439 if (!robot) return;440 let t = 0;441 function frame() {{442 t += 0.02;443 const dy = Math.sin(t) * 2.5;444 robot.setAttribute('transform', `translate(160, ${{280 + dy}})`);445 requestAnimationFrame(frame);446 }}447 frame();448}})();449</script>450</body>451</html>"""452 453 454# ── State assembler ───────────────────────────────────────────────────────────455 456def _build_html(state: dict) -> str:457 called = state.get("called", [])458 active = state.get("active", "")459 edges = state.get("edges", [])460 task = state.get("task", "")461 step = state.get("step", 0)462 mode = state.get("mode", "SEQUENTIAL")463 done = state.get("done", False)464 reward = state.get("reward", None)465 phase = state.get("phase", 1)466 spawned_ids = set(state.get("spawned", []))467 468 # Show only agents that were actually called (+ active if mid-step)469 all_agents = list(called)470 if active and active not in all_agents:471 all_agents.append(active)472 473 # Nothing delegated yet — robot is idle/thinking, no agent cards needed474 positions = _agent_positions(all_agents) if all_agents else {}475 476 def agent_status(aid):477 if aid == active: return "active"478 if aid in called: return "done"479 return "idle"480 481 agents_svg = "\n".join(482 _agent_card_svg(483 aid, *positions[aid],484 agent_status(aid),485 _agent_color(aid, spawned_ids),486 is_spawned=(aid in spawned_ids),487 )488 for aid in all_agents489 )490 beams_svg = _beam_svg(edges, positions)491 robot_svg = _robot_svg()492 493 robot_state = (494 "delegating" if active else495 "done" if done else496 "thinking" if step > 0 else497 "idle"498 )499 500 task_short = (task[:72] + "…") if len(task) > 72 else task501 502 if reward is not None:503 sign = "+" if reward >= 0 else ""504 reward_color = "#10b981" if reward >= 0 else "#ef4444"505 reward_html = f'<span style="color:{reward_color};font-weight:700;">{sign}{reward:.3f}</span>'506 else:507 reward_html = '<span style="color:#334155;">—</span>'508 509 mode_color = {510 "SEQUENTIAL": "#00d4ff",511 "PARALLEL": "#7c3aed",512 "FAN_OUT_REDUCE": "#f59e0b",513 "ITERATIVE": "#10b981",514 "STOP": "#ef4444",515 }.get(mode, "#64748b")516 517 state_json = json.dumps({518 "robot_state": robot_state,519 "active": active,520 "called": called,521 "step": step,522 "done": done,523 "mode": mode,524 "spawned": list(spawned_ids),525 })526 527 return _html_template(528 agents_svg = agents_svg,529 beams_svg = beams_svg,530 robot_svg = robot_svg,531 state_json = state_json,532 task_short = task_short,533 reward_html = reward_html,534 step = step,535 phase = phase,536 mode = mode,537 mode_color = mode_color,538 )539 540 541# ── Public API ────────────────────────────────────────────────────────────────542 543def render_orchestrator(state: dict, height: int = 600) -> None:544 """545 Render the animated robot orchestrator widget in a Streamlit page.546 547 state keys:548 called — list of specialist IDs called so far this episode549 active — specialist being called right now (or "")550 edges — list of [caller_id, callee_id] pairs551 task — task description string552 step — current step number553 mode — delegation mode name (e.g. "SEQUENTIAL")554 done — whether the episode is finished555 reward — cumulative reward float (or None)556 phase — curriculum phase int557 spawned — list of auto-spawned specialist IDs (shown in gold)558 """559 import streamlit.components.v1 as components560 components.html(_build_html(state), height=height, scrolling=False)561