import asyncio
import websockets

class RimeClient:
    def __init__(self, speaker, api_key):
        self.url = f"wss://users-ws.rime.ai/ws?speaker={speaker}&modelId=mistv2&audioFormat=mp3"
        self.auth_headers = {
            "Authorization": f"Bearer {api_key}"
        }
        self.audio_data = b''

    async def send_tokens(self, websocket, message):
        for token in message:
            await websocket.send(token)

    async def handle_audio(self, websocket):
        while True:
            try:
                audio = await websocket.recv()
            except websockets.exceptions.ConnectionClosedOK:
                break
            self.audio_data += audio

    async def run(self, message):
        async with websockets.connect(self.url, additional_headers=self.auth_headers) as websocket:
            await asyncio.gather(
                self.send_tokens(websocket, message),
                self.handle_audio(websocket),
            )

    def save_audio(self, file_path):
        with open(file_path, 'wb') as f:
            f.write(self.audio_data)

message = [
    "This ",
    "is ",
    "a ",
    "test ",
    "<CLEAR>",
    "This ",
    "is ",
    "a ",
    "sentence, ",
    "that ",
    "will ",
    "produce ",
    "audio ",
    "across ",
    "two ",
    "messages.",
    "<EOS>",
]

client = RimeClient("cove", api_key="xxx")
asyncio.run(client.run(message))

client.save_audio("output.mp3")