import os
import queue
import asyncio
import concurrent.futures
import functools
import io
import sys
import random
from threading import Thread
import time
from dotenv import load_dotenv
import pyaudio
import speech_recognition as sr
import websockets
from aioconsole import ainput # for async input
from pydub import AudioSegment
from simpleaudio import WaveObject
load_dotenv()
executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)
web2_initial_message = True
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 44100
class AudioPlayer:
def __init__(self):
self.play_thread = None
self.stop_flag = False
self.queue = queue.Queue()
def play_audio(self):
while not self.stop_flag or not self.queue.empty():
try:
wav_data = self.queue.get_nowait()
except queue.Empty:
continue
wave_obj = WaveObject.from_wave_file(wav_data)
play_obj = wave_obj.play()
while play_obj.is_playing() and not self.stop_flag:
time.sleep(0.1)
if self.stop_flag:
play_obj.stop()
def start_playing(self, wav_data):
self.stop_flag = False
self.queue.put(wav_data)
if self.play_thread is None or not self.play_thread.is_alive():
self.play_thread = Thread(target=self.play_audio)
self.play_thread.start()
def stop_playing(self):
if self.play_thread and self.play_thread.is_alive():
self.stop_flag = True
self.play_thread.join()
self.play_thread = None
def add_to_queue(self, wav_data):
self.queue.put(wav_data)
audio_player = AudioPlayer()
def get_input_device_id():
p = pyaudio.PyAudio()
devices = [(i, p.get_device_info_by_index(i)['name'])
for i in range(p.get_device_count())
if p.get_device_info_by_index(i).get('maxInputChannels')]
print('Available devices:')
for id, name in devices:
print(f"Device id {id} - {name}")
return int(input('Please select device id: '))
async def handle_audio(websocket, device_id):
with sr.Microphone(device_index=device_id, sample_rate=RATE) as source:
recognizer = sr.Recognizer()
print('Source sample rate: ', source.SAMPLE_RATE)
print('Source width: ', source.SAMPLE_WIDTH)
print('Adjusting for ambient noise...Wait for 2 seconds')
recognizer.energy_threshold = 5000
recognizer.dynamic_energy_ratio = 6
recognizer.dynamic_energy_adjustment_damping = 0.85
recognizer.non_speaking_duration = 0.5
recognizer.pause_threshold = 0.8
recognizer.phrase_threshold = 0.5
recognizer.adjust_for_ambient_noise(source, duration=2)
listen_func = functools.partial(
recognizer.listen, source, phrase_time_limit=30)
print('Okay, start talking!')
while True:
print('[*]', end="") # indicate that we are listening
audio = await asyncio.get_event_loop().run_in_executor(executor, listen_func)
await websocket.send(audio.frame_data)
print('[-]', end="") # indicate that we are done listening
await asyncio.sleep(2)
async def handle_text(websocket):
print('You: ', end="", flush=False)
while True:
message = await ainput()
await websocket.send(message)
initial_message = True
async def receive_message(websocket, websocket2):
web1_init_message = await websocket.recv()
print('web1_init_message: ', web1_init_message)
web2_init_message = await websocket2.recv()
print('web1_init_message: ', web2_init_message)
message_to_websocket1 = "Suppose I'm Steve Jobs now. What question do you have for me?"
await websocket.send(message_to_websocket1)
web1_message = ''
while True:
try:
message = await websocket.recv()
print('here')
except websockets.exceptions.ConnectionClosedError as e:
print("Connection closed unexpectedly: ", e)
break
except Exception as e:
print("An error occurred: ", e)
break
if isinstance(message, str):
if message == '[end]\n':
if not web1_message:
continue
# remove everything before '> ' in the message
message_to_websocket2 = web1_message[web1_message.find('> ') + 2:]
# print('message_to_websocket2: ', message_to_websocket2)
await websocket2.send(message_to_websocket2)
web2_message = ''
j = 0
while True:
j += 1
try:
message = await websocket2.recv()
except websockets.exceptions.ConnectionClosedError as e:
print("Connection closed unexpectedly: ", e)
break
except Exception as e:
print("An error occurred: ", e)
break
if isinstance(message, str):
if message == '[end]\n':
# print('\nWebsocket2: ', end="", flush=False)
if not web2_message:
# print('skip')
continue
# remove everything before '> ' in the message
print(web2_message)
message_from_websocket2 = web2_message[web2_message.find('> ') + 2:]
await websocket.send(message_from_websocket2)
break
elif message.startswith('[+]'):
# stop playing audio
audio_player.stop_playing()
# indicate the transcription is done
# print(f"\nnWebsocket2: {message}", end="\n", flush=False)
elif message.startswith('[=]'):
# indicate the response is done
# print(f"nWebsocket2: {web2_message}", end="\n", flush=False)
pass
else:
# print('\nmessage++\n')
web2_message += message
elif isinstance(message, bytes):
global web2_initial_message
if web2_initial_message:
web2_initial_message = False
continue
audio_data = io.BytesIO(message)
audio = AudioSegment.from_mp3(audio_data)
wav_data = io.BytesIO()
audio.export(wav_data, format="wav")
# Start playing audio
audio_player.start_playing(wav_data)
elif message.startswith('[+]'):
# stop playing audio
audio_player.stop_playing()
# indicate the transcription is done
print(f"\n{message}", end="\n", flush=False)
elif message.startswith('[=]'):
# indicate the response is done
print(f"{message}", end="\n", flush=False)
else:
web1_message += message
print(f"{message}", end="", flush=False)
elif isinstance(message, bytes):
audio_data = io.BytesIO(message)
audio = AudioSegment.from_mp3(audio_data)
wav_data = io.BytesIO()
audio.export(wav_data, format="wav")
# Start playing audio
audio_player.start_playing(wav_data)
else:
print("Unexpected message")
break
def select_model():
llm_model_selection = input(
'1: gpt-3.5-turbo-16k \n'
'2: gpt-4 \n'
'3: claude-2 \n'
'Select llm model:')
if llm_model_selection == '1':
llm_model = 'gpt-3.5-turbo-16k'
elif llm_model_selection == '2':
llm_model = 'gpt-4'
elif llm_model_selection == '3':
llm_model = 'claude-2'
return llm_model
async def start_client(client_id, url):
api_key = os.getenv('AUTH_API_KEY')
llm_model = select_model()
uri = f"ws://{url}/ws/{client_id}?api_key={api_key}&llm_model={llm_model}"
async with websockets.connect(uri) as websocket:
uri2 = f"ws://{url}/ws/9999999?api_key={api_key}&llm_model={llm_model}"
# send client platform info
async with websockets.connect(uri2) as websocket2:
await websocket.send('terminal')
await websocket2.send('terminal')
print(f"Client #{client_id} connected to websocket1")
print(f"Client 9999999 connected to websocket2")
welcome_message = await websocket.recv()
welcome_message2 = await websocket2.recv()
print(f"{welcome_message}")
character = input('Select character: ')
await websocket.send(character)
await websocket2.send('6')
mode = input('Select mode (1: audio, 2: text): ')
if mode.lower() == '1':
device_id = get_input_device_id()
send_task = asyncio.create_task(handle_audio(websocket, device_id))
else:
send_task = asyncio.create_task(handle_text(websocket))
receive_task = asyncio.create_task(receive_message(websocket, websocket2))
await asyncio.gather(receive_task, send_task)
async def main(url):
client_id = random.randint(0, 1000000)
task = asyncio.create_task(start_client(client_id, url))
try:
await task
except KeyboardInterrupt:
task.cancel()
await asyncio.wait_for(task, timeout=None)
print("Client stopped by user")
if __name__ == "__main__":
url = sys.argv[1] if len(sys.argv) > 1 else 'localhost:8000'
asyncio.run(main(url))