|
import asyncio
|
|
import websockets
|
|
import json
|
|
import sys
|
|
|
|
|
|
class AsyncClient:
|
|
|
|
def __init__(self, url="ws://127.0.0.1:8470"):
|
|
|
|
self.url = url
|
|
self.ws = None
|
|
self.queue_in = asyncio.Queue()
|
|
self.queue_out = asyncio.Queue()
|
|
self.communication_task = None
|
|
|
|
async def ensure_connection():
|
|
|
|
if not self.ws:
|
|
self.ws = await websockets.connect(self.url)
|
|
|
|
return self.ws
|
|
|
|
async def ensure_communication(self):
|
|
|
|
if not self.communication_task:
|
|
self.communication_task = asyncio.create_task(self.communicate())
|
|
|
|
return self.communication_task
|
|
|
|
async def chat(self, message):
|
|
await self.ensure_communication()
|
|
await self.queue_out.put(message)
|
|
while True:
|
|
|
|
while True:
|
|
try:
|
|
response = await asyncio.wait_for(self.queue_in.get(), 0.1)
|
|
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
break
|
|
|
|
yield response
|
|
|
|
if response["done"]:
|
|
break
|
|
|
|
async def communicate(self):
|
|
loop = asyncio.get_event_loop()
|
|
async with websockets.connect(self.url) as websocket:
|
|
while True:
|
|
message_content = None
|
|
while not message_content:
|
|
try:
|
|
message_content = await asyncio.wait_for(
|
|
self.queue_out.get(), 0.1
|
|
)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
response = await websocket.send(json.dumps(message_content))
|
|
|
|
while True:
|
|
response = json.loads(await websocket.recv())
|
|
|
|
if response["done"]:
|
|
break
|
|
|
|
await self.queue_in.put(response)
|
|
|
|
await self.queue_in.put(response)
|
|
|
|
|
|
async def cli_client(url="ws://127.0.0.1:8470"):
|
|
|
|
loop = asyncio.get_event_loop()
|
|
async_client = AsyncClient(url)
|
|
|
|
while True:
|
|
|
|
sys.stdout.write("> ")
|
|
sys.stdout.flush()
|
|
message_content = await loop.run_in_executor(None, sys.stdin.readline)
|
|
|
|
async for response in async_client.chat(message_content):
|
|
|
|
print(response["content"], end="", flush=True)
|
|
|
|
if response["done"]:
|
|
break
|
|
|
|
print("")
|
|
|
|
|
|
def main():
|
|
url = "ws://127.0.0.1:8470"
|
|
try:
|
|
url = sys.argv[1]
|
|
except IndexError:
|
|
pass
|
|
asyncio.run(cli_client(url))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|