demo.py

demo.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import aiohttp
import sys
import logging
import json
import time
from contextlib import asynccontextmanager, contextmanager, AsyncExitStack
from typing import AsyncGenerator, Generator

logger = logging.getLogger('websockets')
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())

@contextmanager
def _pyaudio() -> Generator[None, None, None]:
    yield

@contextmanager
def _pyaudio_open_stream(*args, **kwargs) -> Generator[None, None, None]:
    yield

@asynccontextmanager
async def _polite_websocket(ws: aiohttp.ClientWebSocketResponse) -> AsyncGenerator[aiohttp.ClientWebSocketResponse, None]:
    try:
        yield ws
    finally:
        print('Terminating connection')
        await ws.send_str('{"eof" : 1}')
        msg = await ws.receive()
        handle_received(msg)

def handle_received(msg):
    if msg.type == aiohttp.WSMsgType.TEXT:
        message = json.loads(msg.data)
        if "partial" in message:
            print(f'partial: {message["partial"]}', end="\r")
        if "text" in message:
            last_time = time.time()
            print()
            logger.info(f'final: {message}')
            return False
    elif msg.type == aiohttp.WSMsgType.CLOSE:
        return True
    elif msg.type == aiohttp.WSMsgType.ERROR:
        return True

async def hello(uri, token, input_file):
    headers = {}
    if token:
        headers['token'] = token
    async with aiohttp.ClientSession(headers=headers) as session:
        async with AsyncExitStack() as stack:
            ws = await stack.enter_async_context(session.ws_connect(uri))
            print(f'Connected to {uri}')
            print('Type Ctrl-C to exit')
            ws = await stack.enter_async_context(_polite_websocket(ws))
            stop = False
            with open(input_file, 'rb') as audio_file:
                chunk_size = 1600
                while not stop:
                    data = audio_file.read(chunk_size)
                    if len(data) == 0:
                        break
                    await ws.send_bytes(data)
                    msg = await ws.receive()
                    stop = handle_received(msg)

if len(sys.argv) != 4:
    print("Usage: python script.py <server_url> <jwt_token> <input_audio_file>")
    exit(1)

server = sys.argv[1]
jwt = sys.argv[2]
input_file = sys.argv[3]

try:
    loop = asyncio.get_event_loop()
    loop.run_until_complete(hello(server, jwt, input_file))
except (Exception, KeyboardInterrupt) as e:
    loop.stop()
    loop.run_until_complete(loop.shutdown_asyncgens())
    if isinstance(e, KeyboardInterrupt):
        print('Bye')
        exit(0)
    else:
        print(f'Oops! {e}')
        exit(1)