@ -4,19 +4,22 @@ import json
import time
import queue
import os
import traceback
from queue import Queue
from threading import Thread
import threading
import uvicorn
import re
from fastapi import FastAPI
from threading import Thread
from starlette . websockets import WebSocket
from stt import stt
from stt import stt _bytes
from tts import tts
from pathlib import Path
import asyncio
from i import configure_interpreter
import urllib . parse
from utils . put_kernel_messages_into_queue import put_kernel_messages_into_queue
from i import configure_interpreter
from interpreter import interpreter
app = FastAPI ( )
@ -30,10 +33,10 @@ def is_full_sentence(text):
def split_into_sentences ( text ) :
return re . split ( r ' (?<=[.!?]) \ s+ ' , text )
# Global q ueues
receive_queue = queue . Queue ( )
send_queue = queue . Queue ( )
recieve_computer_queue = queue . Queue ( ) # Just for computer messages from the device
# Q ueues
from_computer = queue . Queue ( ) # Just for computer messages from the device. Sync queue because interpreter.run is synchronous
from_user = asyncio . Queue ( ) # Just for user messages from the device.
to_device = asyncio . Queue ( ) # For messages we send.
# Switch code executor to device if that's set
@ -56,14 +59,14 @@ if os.getenv('CODE_RUNNER') == "device":
# Unless it was just sent to the device, send it wrapped in flags
if not ( interpreter . messages and interpreter . messages [ - 1 ] == message ) :
send_queu e. put ( { " role " : " assistant " , " type " : " code " , " format " : " python " , " start " : True } )
send_queu e. put ( message )
send_queu e. put ( { " role " : " assistant " , " type " : " code " , " format " : " python " , " end " : True } )
to_devic e. put ( { " role " : " assistant " , " type " : " code " , " format " : " python " , " start " : True } )
to_devic e. put ( message )
to_devic e. put ( { " role " : " assistant " , " type " : " code " , " format " : " python " , " end " : True } )
# Stream the response
print ( " Waiting for the device to respond... " )
while True :
chunk = recieve_computer_queue . get ( )
chunk = from_computer . get ( )
print ( " Server recieved from device: " , chunk )
if " end " in chunk :
break
@ -87,47 +90,52 @@ async def websocket_endpoint(websocket: WebSocket):
await websocket . accept ( )
receive_task = asyncio . create_task ( receive_messages ( websocket ) )
send_task = asyncio . create_task ( send_messages ( websocket ) )
try :
await asyncio . gather ( receive_task , send_task )
except Exception as e :
traceback . print_exc ( )
print ( f " Connection lost. Error: { e } " )
async def receive_messages ( websocket : WebSocket ) :
while True :
data = await websocket . receive_text ( )
if type ( data ) == dict and data [ " role " ] == " computer " :
recieve_computer_queue . put ( data ) # To be handled by interpreter.computer.run
data = await websocket . receive_json ( )
if data [ " role " ] == " computer " :
from_computer . put ( data ) # To be handled by interpreter.computer.run
elif data [ " role " ] == " user " :
await from_user . put ( data )
else :
receive_queue . put ( data )
raise ( " Unknown role: " , data )
async def send_messages ( websocket : WebSocket ) :
while True :
message = await asyncio. get_event_loop ( ) . run_in_executor ( None , send_queue . get )
message = await to_device. get ( )
print ( message )
await websocket . send_json ( message )
def queue _listener( ) :
audio_ file = bytearray ( )
async def user _listener( ) :
audio_ bytes = bytearray ( )
while True :
# Check 10x a second for new messages
while receive_queue . empty ( ) :
time . sleep ( 0.1 )
message = receive_queue . get ( )
message = json . loads ( message )
message = await from_user . get ( )
# Hold the audio in a buffer. If it's ready (we got end flag, stt it)
if message [ " type " ] == " audio " :
if " content " in message :
audio_ file . extend ( bytes ( ast . literal_eval ( message [ " content " ] ) ) )
audio_bytes . extend ( bytes ( ast . literal_eval ( message [ " content " ] ) ) )
if " end " in message :
content = stt ( audio_file , message [ " format " ] )
content = stt _bytes( audio_bytes , message [ " format " ] )
if content == None : # If it was nothing / silence
continue
audio_ file = bytearray ( )
audio_ bytes = bytearray ( )
message = { " role " : " user " , " type " : " message " , " content " : content }
else :
continue
# Ignore flags, we only needed them for audio ^
if " content " not in message :
continue
# Custom stop message will halt us
if message . get ( " content " ) and message . get ( " content " ) . lower ( ) . strip ( " .,! " ) == " stop " :
if message [ " content " ] . lower ( ) . strip ( " .,! " ) == " stop " :
continue
# Load, append, and save conversation history
@ -142,53 +150,59 @@ def queue_listener():
for chunk in interpreter . chat ( messages , stream = True ) :
# Send it to the user
send_queu . put ( chunk )
await to_devic . put ( chunk )
# Speak full sentences out loud
if chunk [ " role " ] == " assistant " and " content " in chunk :
print ( " Chunk role is assistant and content is present in chunk. " )
accumulated_text + = chunk [ " content " ]
print ( " Accumulated text: " , accumulated_text )
sentences = split_into_sentences ( accumulated_text )
print ( " Sentences after splitting: " , sentences )
if is_full_sentence ( sentences [ - 1 ] ) :
print ( " Last sentence is a full sentence. " )
for sentence in sentences :
print ( " Streaming sentence: " , sentence )
stream_tts_to_user ( sentence )
await stream_or_play_tts ( sentence )
accumulated_text = " "
print ( " Reset accumulated text. " )
else :
print ( " Last sentence is not a full sentence. " )
for sentence in sentences [ : - 1 ] :
print ( " Streaming sentence: " , sentence )
stream_tts_to_user ( sentence )
await stream_or_play_tts ( sentence )
accumulated_text = sentences [ - 1 ]
print ( " Accumulated text is now the last sentence: " , accumulated_text )
# If we have a new message, save our progress and go back to the top
if not receive_queue . empty ( ) :
if not from_user . empty ( ) :
with open ( conversation_history_path , ' w ' ) as file :
json . dump ( interpreter . messages , file )
break
def stream_tts_to_user ( sentence ) :
send_queue . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " start " : True } )
audio_bytes = tts ( sentence )
send_queue . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " content " : str ( audio_bytes ) } )
send_queue . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " end " : True } )
async def stream_or_play_tts ( sentence ) :
if os . getenv ( ' TTS_RUNNER ' ) == " server " :
tts ( sentence , play_audio = True )
else :
await to_device . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " start " : True } )
audio_bytes = tts ( sentence , play_audio = False )
await to_device . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " content " : str ( audio_bytes ) } )
await to_device . put ( { " role " : " assistant " , " type " : " audio " , " format " : " audio/mp3 " , " end " : True } )
# Create a thread for the queue listener
queue_thread = Thread ( target = queue_listener )
# Start the queue listener thread
queue_thread . start ( )
from uvicorn import Config , Server
# Run the FastAPI app
if __name__ == " __main__ " :
async def main ( ) :
# Start listening to the user
asyncio . create_task ( user_listener ( ) )
# Start watching the kernel if it's your job to do that
if os . getenv ( ' CODE_RUNNER ' ) == " server " :
asyncio . create_task ( put_kernel_messages_into_queue ( from_computer ) )
server_url = os . getenv ( ' SERVER_URL ' )
if not server_url :
raise ValueError ( " The environment variable SERVER_URL is not set. Please set it to proceed. " )
parsed_url = urllib . parse . urlparse ( server_url )
print ( " Starting `server.py`... " )
uvicorn . run ( app , host = parsed_url . hostname , port = parsed_url . port )
config = Config ( app , host = parsed_url . hostname , port = parsed_url . port , lifespan = ' on ' )
server = Server ( config )
await server . serve ( )
asyncio . run ( main ( ) )
killian
11 months ago