Update gui_update_example.py

Included logic to use the gui update mechanism from a user program - much better logic, faster updates.

examples/gui_update_example.py

@@ -12,37 +12,52 @@
 from metagpt.team import Team
 from metagpt.roles import Role
 from dotenv import load_dotenv
+from typing import List
 
 load_dotenv()
 
 app = FastAPI()
 
-# Store active WebSocket connections
-active_connections = []
+# Initialize WebSocket connections list
+active_connections: List[WebSocket] = []
 
 # Export the callback for other modules to use
 websocket_gui_callback = None
 
+
 async def initialize_websocket_callback():
     """Initialize the WebSocket callback function"""
     async def _websocket_gui_callback(update_info):
         """Send updates to all connected WebSocket clients"""
-        for connection in active_connections[:]:  # Create a copy of the list to safely modify it
+        logger.info(f"Active connections: {len(active_connections)}")
+        if not active_connections:
+            logger.warning("No active WebSocket connections to send updates to.")
+            return
+
+        for connection in active_connections[:]:
             try:
-                # Add timestamp to updates
-                update_info["timestamp"] = datetime.datetime.now().isoformat()
+                # Add timestamp if not present
+                if 'timestamp' not in update_info:
+                    update_info["timestamp"] = datetime.datetime.now().isoformat()
+
+                # Send update and immediately flush
                 await connection.send_json(update_info)
-                logger.debug(f"Sent WebSocket update: {update_info['event']}")
+                await asyncio.sleep(0.1)  # Small delay between messages
+
+                logger.debug(f"Sent WebSocket update: {update_info}")
             except Exception as e:
                 logger.error(f"Error sending WebSocket update: {e}")
                 try:
                     active_connections.remove(connection)
                 except ValueError:
-                    pass  # Connection was already removed
+                    pass
 
+    # Set both the global callback and the WebSocketClient callback
     global websocket_gui_callback
     websocket_gui_callback = _websocket_gui_callback
-    return websocket_gui_callback
+    WebSocketClient.set_callback(_websocket_gui_callback)
+    logger.info("WebSocket callback initialized")
+
 
 @app.get("/")
 async def get():
@@ -109,8 +124,7 @@ async def get():
                             updateHtml = `
                                 <div class="update thinking">
                                     <strong>Thinking</strong>
-                                    <pre>State: ${data.state}
-Current Action: ${data.current_action || 'None'}</pre>
+                                    <pre>State: ${data.state} Current Action: ${data.current_action || 'None'}</pre>
                                 </div>`;
                             break;
 
@@ -122,6 +136,16 @@ async def get():
                                 </div>`;
                             break;
 
+                        case 'SIMPLE_UPDATE':
+                            updateHtml = `
+                                <div class="update">
+                                    <strong>${data.role || 'System'}</strong>: ${data.status}
+                                    <div class="timestamp" style="font-size: 0.8em; color: #666;">
+                                        ${new Date(data.timestamp).toLocaleString()}
+                                    </div>
+                                </div>`;
+                            break;
+
                         case 'planning_and_acting':
                             updateHtml = `
                                 <div class="update planning">
@@ -155,21 +179,46 @@ async def get():
                         .then(response => response.json())
                         .then(data => console.log('Team started:', data));
                 }
+
+                $(document).ready(function() {
+                    initializeWebSocket();
+                });
             </script>
         </body>
     </html>
     """
     return HTMLResponse(content=html_content)
 
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
+    logger.info("WebSocket connection established")
     active_connections.append(websocket)
     try:
         while True:
-            await websocket.receive_text()
-    except:
-        active_connections.remove(websocket)
+            try:
+                # Set WebSocket to text mode for lower latency
+                data = await websocket.receive_text()
+                # Immediately acknowledge receipt
+                await websocket.send_json({"status": "received"})
+            except WebSocketDisconnect:
+                break
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}")
+    finally:
+        if websocket in active_connections:
+            active_connections.remove(websocket)
+        logger.info("WebSocket connection closed")
+
+
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Running setup_callbacks at startup.")
+    await initialize_websocket_callback()
+    logger.info("WebSocket callback initialized.")
+    print("Ready.")
+
 
 @app.get("/kickoff")
 async def kickoff():
@@ -180,6 +229,17 @@ async def kickoff():
     gpt4 = Config.default()
     gpt4.llm.model = "gpt-4o-mini"
 
+    logger.info("UPDATING GUI inside /kickoff...")
+    update_info = {
+        "event": "SIMPLE_UPDATE",
+        "status": "Inside kickoff, starting exec...",
+        "current_task": None,
+        "progress": 0,
+        "role": "Manager",
+        "profile": "Profile"
+    }
+    await gui_update_callback(update_info)
+
     action1 = Action(config=gpt4, name="AlexSay", instruction="Express your opinion with emotion and don't repeat it")
     action2 = Action(config=gpt4, name="BobSay", instruction="Express your opinion with emotion and don't repeat it")
     alex = Role(name="Alex", profile="Democratic candidate", goal="Win the election", actions=[action1], watch=[action2])
@@ -188,6 +248,13 @@ async def kickoff():
     WebSocketClient.set_callback(websocket_gui_callback)
     team = Team(investment=10.0, env=env, roles=[alex, bob], progress_callback=gui_update_callback)
 
-    asyncio.create_task(team.run(idea="Topic: climate change. Under 80 words per message.", send_to="Alex", n_round=5, auto_archive=False))
+    asyncio.create_task(
+        team.run(
+            idea="Topic: climate change. Under 80 words per message.",
+            send_to="Alex",
+            n_round=5,
+            auto_archive=False
+        )
+    )
 
     return {"status": "started"}