Spaces:

Almaatla
/

WebSocketChat

Sleeping

App Files Files Community

Almaatla commited on Apr 14

Commit

050938e

verified ·

1 Parent(s): 0242cbc

Update app.py

Browse files

Files changed (1) hide show

app.py +67 -11

app.py CHANGED Viewed

@@ -1,20 +1,30 @@
 from fastapi import FastAPI, WebSocket
 from fastapi.responses import HTMLResponse
 import uvicorn
 app = FastAPI()
 class ConnectionManager:
     def __init__(self):
-        self.active_connections = []
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
-        self.active_connections.append(websocket)
-    async def broadcast(self, message: str):
-        for connection in self.active_connections:
-            await connection.send_text(message)
 manager = ConnectionManager()
@@ -23,18 +33,27 @@ async def get():
     return HTMLResponse("""
     <html>
     <body>
-        <h1>Simple Chat</h1>
         <div id="chat" style="height:300px;overflow-y:scroll"></div>
         <input id="msg" type="text">
         <button onclick="send()">Send</button>
         <script>
             const ws = new WebSocket('wss://' + window.location.host + '/ws');
             ws.onmessage = e => {
                 document.getElementById('chat').innerHTML +=
-                    `<div>${e.data}</div>`;
             };
             const send = () => {
-                ws.send(document.getElementById('msg').value);
                 document.getElementById('msg').value = '';
             };
         </script>
@@ -42,15 +61,52 @@ async def get():
     </html>
     """)
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await manager.connect(websocket)
     try:
         while True:
             message = await websocket.receive_text()
-            await manager.broadcast(f"User: {message}")
-    except:
-        manager.active_connections.remove(websocket)
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)

 from fastapi import FastAPI, WebSocket
 from fastapi.responses import HTMLResponse
 import uvicorn
+import json
 app = FastAPI()
 class ConnectionManager:
     def __init__(self):
+        self.active_connections = {}  # WebSocket: source
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
+        self.active_connections[websocket] = None
+    def set_source(self, websocket: WebSocket, source: str):
+        if websocket in self.active_connections:
+            self.active_connections[websocket] = source
+    async def send_to_destination(self, destination: str, message: str):
+        for ws, src in self.active_connections.items():
+            if src == destination:
+                await ws.send_text(message)
+    def remove(self, websocket: WebSocket):
+        if websocket in self.active_connections:
+            del self.active_connections[websocket]
 manager = ConnectionManager()
     return HTMLResponse("""
     <html>
     <body>
+        <h1>Chat Client</h1>
         <div id="chat" style="height:300px;overflow-y:scroll"></div>
         <input id="msg" type="text">
         <button onclick="send()">Send</button>
         <script>
             const ws = new WebSocket('wss://' + window.location.host + '/ws');
+            ws.onopen = () => {
+                ws.send(JSON.stringify({ source: 'user' }));
+            };
             ws.onmessage = e => {
+                const msg = JSON.parse(e.data);
                 document.getElementById('chat').innerHTML +=
+                    `<div>${msg.content}</div>`;
             };
             const send = () => {
+                const message = {
+                    content: document.getElementById('msg').value,
+                    source: 'user',
+                    destination: 'proxy'
+                };
+                ws.send(JSON.stringify(message));
                 document.getElementById('msg').value = '';
             };
         </script>
     </html>
     """)
+@app.get("/proxy")
+async def get_proxy():
+    return HTMLResponse("""
+    <html>
+    <body>
+        <h1>Proxy Client</h1>
+        <script>
+            const ws = new WebSocket('wss://' + window.location.host + '/ws');
+            ws.onopen = () => {
+                ws.send(JSON.stringify({ source: 'proxy' }));
+            };
+            ws.onmessage = e => {
+                const msg = JSON.parse(e.data);
+                if (msg.destination === 'proxy') {
+                    const response = {
+                        content: "Hello user!",
+                        source: 'proxy',
+                        destination: 'user'
+                    };
+                    ws.send(JSON.stringify(response));
+                }
+            };
+        </script>
+    </body>
+    </html>
+    """)
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await manager.connect(websocket)
     try:
+        # Handle initial source identification
+        data = await websocket.receive_text()
+        init_msg = json.loads(data)
+        if 'source' in init_msg:
+            manager.set_source(websocket, init_msg['source'])
+        # Handle messages
         while True:
             message = await websocket.receive_text()
+            msg_data = json.loads(message)
+            await manager.send_to_destination(msg_data['destination'], message)
+    except Exception as e:
+        manager.remove(websocket)
+        await websocket.close()
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)