From 3516e8b75ba09ceef9236caf7fbb1c494a16bcb9 Mon Sep 17 00:00:00 2001 From: "Kumar R." Date: Thu, 21 Nov 2024 07:47:04 +0530 Subject: [PATCH] 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 | 95 +++++++++++++++++++++++++++++----- 1 file changed, 81 insertions(+), 14 deletions(-) diff --git a/examples/gui_update_example.py b/examples/gui_update_example.py index d2c5d82b4..8f3e1a1d5 100644 --- a/examples/gui_update_example.py +++ b/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 = `
Thinking - 
State: ${data.state}
-Current Action: ${data.current_action || 'None'}
+ 
State: ${data.state} Current Action: ${data.current_action || 'None'}
`; break; @@ -122,6 +136,16 @@ async def get(): `; break; + case 'SIMPLE_UPDATE': + updateHtml = ` +
+ ${data.role || 'System'}: ${data.status} +
+ ${new Date(data.timestamp).toLocaleString()} +
+
`; + break; + case 'planning_and_acting': updateHtml = `
@@ -155,21 +179,46 @@ async def get(): .then(response => response.json()) .then(data => console.log('Team started:', data)); } + + $(document).ready(function() { + initializeWebSocket(); + }); """ 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"}