Create app.py

app.py

@@ -0,0 +1,292 @@
+import gradio as gr
+import asyncio
+import websockets
+import json
+import uuid
+import argparse
+import urllib.parse
+from datetime import datetime
+import logging
+import sys
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger("chat-node")
+
+# Dictionary to store active connections
+active_connections = {}
+# Dictionary to store message history for each chat room
+chat_history = {}
+
+# Get node name from URL or command line
+def get_node_name():
+    parser = argparse.ArgumentParser(description='Start a chat node with a specific name')
+    parser.add_argument('--node-name', type=str, default=None, help='Name for this chat node')
+    parser.add_argument('--port', type=int, default=7860, help='Port to run the Gradio interface on')
+    
+    args = parser.parse_args()
+    node_name = args.node_name
+    port = args.port
+    
+    # If no node name specified, generate a random one
+    if not node_name:
+        node_name = f"node-{uuid.uuid4().hex[:8]}"
+        
+    return node_name, port
+
+async def websocket_handler(websocket, path):
+    """Handle WebSocket connections."""
+    try:
+        # Extract room_id from path if present
+        path_parts = path.strip('/').split('/')
+        room_id = path_parts[0] if path_parts else "default"
+        
+        # Register the new client
+        client_id = str(uuid.uuid4())
+        if room_id not in active_connections:
+            active_connections[room_id] = {}
+            chat_history[room_id] = []
+        
+        active_connections[room_id][client_id] = websocket
+        
+        # Send welcome message and chat history
+        welcome_msg = {
+            "type": "system",
+            "content": f"Welcome to room '{room_id}'! Connected from node '{NODE_NAME}'",
+            "timestamp": datetime.now().isoformat(),
+            "sender": "system",
+            "room_id": room_id
+        }
+        await websocket.send(json.dumps(welcome_msg))
+        
+        # Send chat history
+        for msg in chat_history[room_id]:
+            await websocket.send(json.dumps(msg))
+        
+        # Broadcast join notification
+        join_msg = {
+            "type": "system",
+            "content": f"User joined the room",
+            "timestamp": datetime.now().isoformat(),
+            "sender": "system",
+            "room_id": room_id
+        }
+        await broadcast_message(join_msg, room_id)
+        
+        logger.info(f"New client {client_id} connected to room {room_id}")
+        
+        # Handle messages from this client
+        async for message in websocket:
+            try:
+                data = json.loads(message)
+                
+                # Add metadata to the message
+                data["timestamp"] = datetime.now().isoformat()
+                data["sender_node"] = NODE_NAME
+                data["room_id"] = room_id
+                
+                # Store in history
+                chat_history[room_id].append(data)
+                if len(chat_history[room_id]) > 100:  # Limit history to 100 messages
+                    chat_history[room_id] = chat_history[room_id][-100:]
+                
+                # Broadcast to all clients in the room
+                await broadcast_message(data, room_id)
+                
+            except json.JSONDecodeError:
+                error_msg = {
+                    "type": "error",
+                    "content": "Invalid JSON format",
+                    "timestamp": datetime.now().isoformat(),
+                    "sender": "system",
+                    "room_id": room_id
+                }
+                await websocket.send(json.dumps(error_msg))
+    
+    except websockets.exceptions.ConnectionClosed:
+        logger.info(f"Client {client_id} disconnected from room {room_id}")
+    finally:
+        # Remove the client when disconnected
+        if room_id in active_connections and client_id in active_connections[room_id]:
+            del active_connections[room_id][client_id]
+            
+            # Broadcast leave notification
+            leave_msg = {
+                "type": "system",
+                "content": f"User left the room",
+                "timestamp": datetime.now().isoformat(),
+                "sender": "system",
+                "room_id": room_id
+            }
+            await broadcast_message(leave_msg, room_id)
+            
+            # Clean up empty rooms
+            if not active_connections[room_id]:
+                del active_connections[room_id]
+                # Optionally, you might want to keep the chat history
+
+async def broadcast_message(message, room_id):
+    """Broadcast a message to all clients in a room."""
+    if room_id in active_connections:
+        disconnected_clients = []
+        
+        for client_id, websocket in active_connections[room_id].items():
+            try:
+                await websocket.send(json.dumps(message))
+            except websockets.exceptions.ConnectionClosed:
+                disconnected_clients.append(client_id)
+        
+        # Clean up disconnected clients
+        for client_id in disconnected_clients:
+            del active_connections[room_id][client_id]
+
+async def start_websocket_server(host='0.0.0.0', port=8765):
+    """Start the WebSocket server."""
+    server = await websockets.serve(websocket_handler, host, port)
+    logger.info(f"WebSocket server started on ws://{host}:{port}")
+    return server
+
+def send_message(message, username, room_id):
+    """Function to send a message from the Gradio interface."""
+    if not message.strip():
+        return None
+    
+    loop = asyncio.get_event_loop()
+    msg_data = {
+        "type": "chat",
+        "content": message,
+        "username": username,
+        "room_id": room_id
+    }
+    
+    loop.create_task(broadcast_message(msg_data, room_id))
+    
+    # Format the message for display
+    formatted_msg = f"{username}: {message}"
+    return formatted_msg
+
+def join_room(room_id, chat_history_output):
+    """Join a specific chat room."""
+    if not room_id.strip():
+        return "Please enter a valid room ID", chat_history_output
+    
+    # Sanitize the room ID
+    room_id = urllib.parse.quote(room_id.strip())
+    
+    # Create the room if it doesn't exist
+    if room_id not in chat_history:
+        chat_history[room_id] = []
+    
+    # Format existing messages
+    formatted_history = []
+    for msg in chat_history[room_id]:
+        if msg.get("type") == "chat":
+            formatted_history.append(f"{msg.get('username', 'Anonymous')}: {msg.get('content', '')}")
+        elif msg.get("type") == "system":
+            formatted_history.append(f"System: {msg.get('content', '')}")
+    
+    return f"Joined room: {room_id}", formatted_history
+
+def create_gradio_interface():
+    """Create and return the Gradio interface."""
+    with gr.Blocks(title=f"Chat Node: {NODE_NAME}") as interface:
+        gr.Markdown(f"# Chat Node: {NODE_NAME}")
+        gr.Markdown("Join a room by entering a room ID below or create a new one.")
+        
+        with gr.Row():
+            room_id_input = gr.Textbox(label="Room ID", placeholder="Enter room ID")
+            join_button = gr.Button("Join Room")
+        
+        chat_history_output = gr.Textbox(label="Chat History", lines=15, max_lines=15)
+        
+        with gr.Row():
+            username_input = gr.Textbox(label="Username", placeholder="Enter your username", value="User")
+            message_input = gr.Textbox(label="Message", placeholder="Type your message here")
+            send_button = gr.Button("Send")
+        
+        # Current room display
+        current_room_display = gr.Textbox(label="Current Room", value="Not joined any room yet")
+        
+        # Event handlers
+        join_button.click(
+            join_room,
+            inputs=[room_id_input, chat_history_output],
+            outputs=[current_room_display, chat_history_output]
+        )
+        
+        def send_and_clear(message, username, room_id):
+            if not room_id.startswith("Joined room:"):
+                return "Please join a room first", message
+            
+            actual_room_id = room_id.replace("Joined room: ", "").strip()
+            formatted_msg = send_message(message, username, actual_room_id)
+            
+            if formatted_msg:
+                return "", formatted_msg
+            return message, None
+        
+        send_button.click(
+            send_and_clear,
+            inputs=[message_input, username_input, current_room_display],
+            outputs=[message_input, chat_history_output]
+        )
+        
+        # Enter key to send message
+        message_input.submit(
+            send_and_clear,
+            inputs=[message_input, username_input, current_room_display],
+            outputs=[message_input, chat_history_output]
+        )
+    
+    return interface
+
+async def main():
+    """Main function to start the application."""
+    global NODE_NAME
+    NODE_NAME, port = get_node_name()
+    
+    # Start WebSocket server
+    server = await start_websocket_server()
+    
+    # Create and launch Gradio interface
+    interface = create_gradio_interface()
+    
+    # Custom middleware to extract node name from URL query parameters
+    from starlette.middleware.base import BaseHTTPMiddleware
+    
+    class NodeNameMiddleware(BaseHTTPMiddleware):
+        async def dispatch(self, request, call_next):
+            global NODE_NAME
+            query_params = dict(request.query_params)
+            if "node_name" in query_params:
+                NODE_NAME = query_params["node_name"]
+                logger.info(f"Node name set to {NODE_NAME} from URL parameter")
+            
+            response = await call_next(request)
+            return response
+    
+    # Apply middleware
+    app = gr.routes.App.create_app(interface)
+    app.add_middleware(NodeNameMiddleware)
+    
+    # Launch with the modified app
+    gr.routes.mount_gradio_app(app, interface, path="/")
+    
+    # Run the FastAPI app with uvicorn
+    import uvicorn
+    config = uvicorn.Config(app, host="0.0.0.0", port=port)
+    server = uvicorn.Server(config)
+    
+    logger.info(f"Starting Gradio interface on http://0.0.0.0:{port} with node name '{NODE_NAME}'")
+    
+    await server.serve()
+    
+    # Keep the WebSocket server running
+    await asyncio.Future()
+
+if __name__ == "__main__":
+    asyncio.run(main())