import asyncio
import json
import logging

import websockets
from jsonrpcserver.aio import methods

@methods.add
async def echo(message):
  return message

@methods.add
async def add(a, b):
  return a + b

async def handler(websocket, path):
  log.debug('accept: %s', websocket.remote_address)
  while True:
    try:
      request = await websocket.recv()
      response = await methods.dispatch(request)
      log.debug('response: %s', response)
      await websocket.send(str(response))
    except websockets.exceptions.ConnectionClosed:
      log.debug('connection closed')
      break

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger('api')
start_server = websockets.serve(handler, '0.0.0.0', 3001)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
