zero_gpu_test / chat.py
vericudebuget's picture
Upload 7 files
9c72a28 verified
Raw
History Blame Contribute Delete
2.51 kB
import asyncio, json, time, sys
try:
import websockets
except ImportError:
print("pip install websockets")
sys.exit(1)
messages = []
clients = set()
async def server_handler(ws):
clients.add(ws)
try:
for m in messages:
await ws.send(json.dumps(m))
async for raw in ws:
data = json.loads(raw)
data['server_time'] = int(time.time() * 1000)
messages.append(data)
websockets.broadcast(clients, json.dumps(data))
finally:
clients.discard(ws)
async def run_server(host='127.0.0.1', port=7860):
async with websockets.serve(server_handler, host, port):
print(f"server on ws://{host}:{port}")
await asyncio.Future()
async def client(uri, username):
async with websockets.connect(uri) as ws:
pending, next_id = {}, 0
async def reader():
async for raw in ws:
data = json.loads(raw)
now = int(time.time() * 1000)
user, msg = data['user'], data['msg']
line = f"{user}: {msg}"
if user == username and data.get('id') in pending:
rtt = now - pending.pop(data['id'])
line += f" (rtt: {rtt}ms)"
elif data.get('server_time'):
sp = now - data['server_time']
line += f" (sping: {sp}ms)"
print(f"\r{'':<80}\r{line}\n{username}> ", end='', flush=True)
async def writer():
nonlocal next_id
loop = asyncio.get_running_loop()
print(f"{username}> ", end='', flush=True)
while True:
line = await loop.run_in_executor(None, input)
if not line:
print(f"{username}> ", end='', flush=True)
continue
nid = next_id
next_id += 1
pending[nid] = int(time.time() * 1000)
await ws.send(json.dumps({"user": username, "msg": line, "id": nid}))
await asyncio.gather(reader(), writer())
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == 'server':
port = int(sys.argv[2]) if len(sys.argv) > 2 else 7860
asyncio.run(run_server(port=port))
else:
uri = sys.argv[1] if len(sys.argv) > 1 else 'wss://vericudebuget-online-group.hf.space/ws'
name = sys.argv[2] if len(sys.argv) > 2 else input("Username: ")
asyncio.run(client(uri, name))