a little cleanup

This commit is contained in:
Moishe Lettvin
2024-01-16 19:58:11 -05:00
parent cfaccefe9c
commit 755059c358
14 changed files with 62 additions and 53 deletions

View File

@@ -269,7 +269,7 @@ class LLMResponse(OrchestratorResponse):
yield out.strip() yield out.strip()
def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]: def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]:
return [QueueFrame(FrameType.AUDIO_FRAME, audio_frame)] return [QueueFrame(FrameType.AUDIO, audio_frame)]
def get_frames_from_chunk(self, chunk) -> Generator[list[QueueFrame], Any, None]: def get_frames_from_chunk(self, chunk) -> Generator[list[QueueFrame], Any, None]:
for audio_frame in self.services.tts.run_tts(chunk): for audio_frame in self.services.tts.run_tts(chunk):

View File

@@ -375,7 +375,7 @@ class Orchestrator(EventHandler):
# if interrupted, we just pull frames off the queue and discard them # if interrupted, we just pull frames off the queue and discard them
if not self.is_interrupted.is_set(): if not self.is_interrupted.is_set():
if frame: if frame:
if frame.frame_type == FrameType.AUDIO_FRAME: if frame.frame_type == FrameType.AUDIO:
chunk = frame.frame_data chunk = frame.frame_data
all_audio_frames.extend(chunk) all_audio_frames.extend(chunk)
@@ -385,7 +385,7 @@ class Orchestrator(EventHandler):
if l: if l:
self.mic.write_frames(bytes(b[:l])) self.mic.write_frames(bytes(b[:l]))
b = b[l:] b = b[l:]
elif frame.frame_type == FrameType.IMAGE_FRAME: elif frame.frame_type == FrameType.IMAGE:
self.set_image(frame.frame_data) self.set_image(frame.frame_data)
elif len(b): elif len(b):
self.mic.write_frames(bytes(b)) self.mic.write_frames(bytes(b))

View File

@@ -4,13 +4,14 @@ from dataclasses import dataclass
class FrameType(Enum): class FrameType(Enum):
START_STREAM = 0 START_STREAM = 0
END_STREAM = 1 END_STREAM = 1
AUDIO_FRAME = 2 AUDIO = 2
IMAGE_FRAME = 3 IMAGE = 3
SENTENCE_FRAME = 4 SENTENCE = 4
TEXT_CHUNK_FRAME = 5 TEXT_CHUNK = 5
LLM_MESSAGE_FRAME = 6 LLM_MESSAGE = 6
APP_MESSAGE_FRAME = 7 APP_MESSAGE = 7
IMAGE_DESCRIPTION = 8 IMAGE_DESCRIPTION = 8
TRANSCRIPTION = 9
@dataclass(frozen=True) @dataclass(frozen=True)
class QueueFrame: class QueueFrame:

View File

@@ -73,13 +73,13 @@ class LLMService(AIService):
if not self.output_queue: if not self.output_queue:
raise Exception("Output queue must be set before using the run method.") raise Exception("Output queue must be set before using the run method.")
if frame.frame_type == FrameType.LLM_MESSAGE_FRAME: if frame.frame_type == FrameType.LLM_MESSAGE:
if type(frame.frame_data) != list: if type(frame.frame_data) != list:
raise Exception("LLM service requires a dict for the data field") raise Exception("LLM service requires a dict for the data field")
messages: list[dict[str, str]] = frame.frame_data messages: list[dict[str, str]] = frame.frame_data
async for message in self.run_llm_async_sentences(messages): async for message in self.run_llm_async_sentences(messages):
await self.output_queue.put(QueueFrame(FrameType.SENTENCE_FRAME, message)) await self.output_queue.put(QueueFrame(FrameType.SENTENCE, message))
class TTSService(AIService): class TTSService(AIService):
@@ -98,13 +98,13 @@ class TTSService(AIService):
if not self.output_queue: if not self.output_queue:
raise Exception("Output queue must be set before using the run method.") raise Exception("Output queue must be set before using the run method.")
if frame.frame_type == FrameType.SENTENCE_FRAME: if frame.frame_type == FrameType.SENTENCE:
if type(frame.frame_data) != str: if type(frame.frame_data) != str:
raise Exception("TTS service requires a string for the data field") raise Exception("TTS service requires a string for the data field")
text = frame.frame_data text = frame.frame_data
async for audio in self.run_tts(text): async for audio in self.run_tts(text):
await self.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio)) await self.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
class ImageGenService(AIService): class ImageGenService(AIService):

View File

@@ -35,7 +35,11 @@ class DailyTransportService(EventHandler):
self.duration: float = duration self.duration: float = duration
self.expiration = time.time() + duration * 60 self.expiration = time.time() + duration * 60
self.output_queue = Queue() # This queue is used to marshal frames from the async output queue to the thread that emits audio & video.
# We need this to maintain the asynchronous behavior of asyncio queues -- to give async functions
# a chance to run while waiting for queue items -- but also to maintain thread safety.
self.threadsafe_output_queue = Queue()
self.is_interrupted = Event() self.is_interrupted = Event()
self.stop_threads = Event() self.stop_threads = Event()
self.story_started = False self.story_started = False
@@ -45,6 +49,9 @@ class DailyTransportService(EventHandler):
self.camera_height = 768 self.camera_height = 768
self.camera_enabled = False self.camera_enabled = False
self.output_queue = asyncio.Queue()
self.media_queue = asyncio.Queue()
self.other_participant_has_joined = False self.other_participant_has_joined = False
self.camera_thread = None self.camera_thread = None
@@ -62,12 +69,6 @@ class DailyTransportService(EventHandler):
}, },
} }
# This queue is used to marshal frames from the async output queue to the sync output queue
# We need this to maintain the asynchronous behavior of asyncio queues -- to give async functions
# a chance to run while waiting for queue items -- but also to maintain thread safety for the
# primary output queue.
self.async_output_queue = asyncio.Queue()
self.logger: logging.Logger = logging.getLogger("dailyai") self.logger: logging.Logger = logging.getLogger("dailyai")
self.event_handlers = {} self.event_handlers = {}
@@ -182,24 +183,25 @@ class DailyTransportService(EventHandler):
) )
if self.token: if self.token:
self.transcription_queue = asyncio.Queue()
self.client.start_transcription(self.transcription_settings) self.client.start_transcription(self.transcription_settings)
self.my_participant_id = self.client.participants()["local"]["id"] self.my_participant_id = self.client.participants()["local"]["id"]
async def get_transcriptions(self): async def get_media_frames(self):
while True: while True:
transcript = await self.transcription_queue.get() frame = await self.media_queue.get()
yield transcript yield frame
if frame.frame_type == FrameType.END_STREAM:
break
def get_async_output_queue(self): def get_async_output_queue(self):
return self.async_output_queue return self.output_queue
async def marshal_frames(self): async def marshal_frames(self):
while True: while True:
frame = await self.async_output_queue.get() frame = await self.output_queue.get()
self.output_queue.put(frame) self.threadsafe_output_queue.put(frame)
self.async_output_queue.task_done() self.output_queue.task_done()
if frame.frame_type == FrameType.END_STREAM: if frame.frame_type == FrameType.END_STREAM:
break break
@@ -222,7 +224,8 @@ class DailyTransportService(EventHandler):
self.stop_threads.set() self.stop_threads.set()
await self.async_output_queue.put(QueueFrame(FrameType.END_STREAM, None)) await self.media_queue.put(QueueFrame(FrameType.END_STREAM, None))
await self.output_queue.put(QueueFrame(FrameType.END_STREAM, None))
await async_output_queue_marshal_task await async_output_queue_marshal_task
if self.camera_thread and self.camera_thread.is_alive(): if self.camera_thread and self.camera_thread.is_alive():
@@ -258,9 +261,10 @@ class DailyTransportService(EventHandler):
def on_app_message(self, message, sender): def on_app_message(self, message, sender):
pass pass
def on_transcription_message(self, message): def on_transcription_message(self, message:dict):
if self.loop: if self.loop:
asyncio.run_coroutine_threadsafe(self.transcription_queue.put(message), self.loop) frame = QueueFrame(FrameType.TRANSCRIPTION, message)
asyncio.run_coroutine_threadsafe(self.media_queue.put(frame), self.loop)
def on_transcription_stopped(self, stopped_by, stopped_by_error): def on_transcription_stopped(self, stopped_by, stopped_by_error):
pass pass
@@ -291,7 +295,7 @@ class DailyTransportService(EventHandler):
all_audio_frames = bytearray() all_audio_frames = bytearray()
while True: while True:
try: try:
frames_or_frame: QueueFrame | list[QueueFrame] = self.output_queue.get() frames_or_frame: QueueFrame | list[QueueFrame] = self.threadsafe_output_queue.get()
if type(frames_or_frame) == QueueFrame: if type(frames_or_frame) == QueueFrame:
frames: list[QueueFrame] = [frames_or_frame] frames: list[QueueFrame] = [frames_or_frame]
elif type(frames_or_frame) == list: elif type(frames_or_frame) == list:
@@ -302,13 +306,13 @@ class DailyTransportService(EventHandler):
for frame in frames: for frame in frames:
if frame.frame_type == FrameType.END_STREAM: if frame.frame_type == FrameType.END_STREAM:
self.logger.info("Stopping frame consumer thread") self.logger.info("Stopping frame consumer thread")
self.output_queue.task_done() self.threadsafe_output_queue.task_done()
return return
# if interrupted, we just pull frames off the queue and discard them # if interrupted, we just pull frames off the queue and discard them
if not self.is_interrupted.is_set(): if not self.is_interrupted.is_set():
if frame: if frame:
if frame.frame_type == FrameType.AUDIO_FRAME: if frame.frame_type == FrameType.AUDIO:
chunk = frame.frame_data chunk = frame.frame_data
all_audio_frames.extend(chunk) all_audio_frames.extend(chunk)
@@ -318,7 +322,7 @@ class DailyTransportService(EventHandler):
if l: if l:
self.mic.write_frames(bytes(b[:l])) self.mic.write_frames(bytes(b[:l]))
b = b[l:] b = b[l:]
elif frame.frame_type == FrameType.IMAGE_FRAME: elif frame.frame_type == FrameType.IMAGE:
self.set_image(frame.frame_data) self.set_image(frame.frame_data)
elif len(b): elif len(b):
self.mic.write_frames(bytes(b)) self.mic.write_frames(bytes(b))
@@ -333,7 +337,7 @@ class DailyTransportService(EventHandler):
if frame.frame_type == FrameType.START_STREAM: if frame.frame_type == FrameType.START_STREAM:
self.is_interrupted.clear() self.is_interrupted.clear()
self.output_queue.task_done() self.threadsafe_output_queue.task_done()
except Empty: except Empty:
try: try:
if len(b): if len(b):

View File

@@ -73,7 +73,7 @@ class TestResponse(unittest.TestCase):
while expected_words: while expected_words:
actual_word:QueueFrame = output_queue.get() actual_word:QueueFrame = output_queue.get()
word = expected_words.pop(0) word = expected_words.pop(0)
self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME) self.assertEqual(actual_word.frame_type, FrameType.AUDIO)
self.assertEqual(actual_word.frame_data, bytes(word, "utf-8")) self.assertEqual(actual_word.frame_data, bytes(word, "utf-8"))
output_queue.task_done() output_queue.task_done()
@@ -128,10 +128,10 @@ class TestResponse(unittest.TestCase):
while expected_words and not stop_processing_output_queue.is_set(): while expected_words and not stop_processing_output_queue.is_set():
try: try:
actual_word:QueueFrame = output_queue.get_nowait() actual_word:QueueFrame = output_queue.get_nowait()
if actual_word.frame_type == FrameType.AUDIO_FRAME: if actual_word.frame_type == FrameType.AUDIO:
time.sleep(0.1) time.sleep(0.1)
word = expected_words.pop(0) word = expected_words.pop(0)
self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME) self.assertEqual(actual_word.frame_type, FrameType.AUDIO)
self.assertEqual(actual_word.frame_data, bytes(word, "utf-8")) self.assertEqual(actual_word.frame_data, bytes(word, "utf-8"))
output_queue.task_done() output_queue.task_done()
except Empty: except Empty:

View File

@@ -40,7 +40,7 @@ class StaticSpriteResponse(OrchestratorResponse):
self.image_bytes = img.tobytes() self.image_bytes = img.tobytes()
def do_play(self) -> None: def do_play(self) -> None:
self.output_queue.put(QueueFrame(FrameType.IMAGE_FRAME, self.image_bytes)) self.output_queue.put(QueueFrame(FrameType.IMAGE, self.image_bytes))
class IntroSpriteResponse(StaticSpriteResponse): class IntroSpriteResponse(StaticSpriteResponse):
@@ -73,8 +73,8 @@ class AnimatedSpriteLLMResponse(LLMResponse):
def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]: def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]:
return [ return [
QueueFrame(FrameType.AUDIO_FRAME, audio_frame), QueueFrame(FrameType.AUDIO, audio_frame),
QueueFrame(FrameType.IMAGE_FRAME, random.choice(self.image_bytes)) QueueFrame(FrameType.IMAGE, random.choice(self.image_bytes))
] ]

View File

@@ -38,7 +38,7 @@ async def main(room_url):
return return
async for audio in audio_generator: async for audio in audio_generator:
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio)) transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
# wait for the output queue to be empty, then leave the meeting # wait for the output queue to be empty, then leave the meeting
transport.output_queue.join() transport.output_queue.join()

View File

@@ -41,7 +41,7 @@ async def main(room_url):
audio_generator: AsyncGenerator[bytes, None] = tts.run_tts(f"Hello there, {participant['info']['userName']}!") audio_generator: AsyncGenerator[bytes, None] = tts.run_tts(f"Hello there, {participant['info']['userName']}!")
async for audio in audio_generator: async for audio in audio_generator:
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio)) transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
print("setting up call state handler") print("setting up call state handler")
@transport.event_handler("on_call_state_updated") @transport.event_handler("on_call_state_updated")

View File

@@ -29,7 +29,7 @@ async def main(room_url):
"role": "system", "role": "system",
"content": "You are an LLM in a WebRTC session, and this is a 'hello world' demo. Say hello to the world." "content": "You are an LLM in a WebRTC session, and this is a 'hello world' demo. Say hello to the world."
}] }]
await text_to_llm_queue.put(QueueFrame(FrameType.LLM_MESSAGE_FRAME, messages)) await text_to_llm_queue.put(QueueFrame(FrameType.LLM_MESSAGE, messages))
await text_to_llm_queue.put(QueueFrame(FrameType.END_STREAM, None)) await text_to_llm_queue.put(QueueFrame(FrameType.END_STREAM, None))
llm_task = asyncio.create_task(llm.run()) llm_task = asyncio.create_task(llm.run())

View File

@@ -27,7 +27,7 @@ async def main(room_url):
@transport.event_handler("on_participant_joined") @transport.event_handler("on_participant_joined")
async def on_participant_joined(transport, participant): async def on_participant_joined(transport, participant):
(_, image_bytes) = await image_task (_, image_bytes) = await image_task
transport.output_queue.put(QueueFrame(FrameType.IMAGE_FRAME, image_bytes)) transport.output_queue.put(QueueFrame(FrameType.IMAGE, image_bytes))
await transport.run() await transport.run()

View File

@@ -41,11 +41,11 @@ async def main(room_url:str):
)) ))
async for audio_chunk in tts.run_tts("My friend the LLM is now going to tell a joke about llamas."): async for audio_chunk in tts.run_tts("My friend the LLM is now going to tell a joke about llamas."):
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio_chunk)) transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio_chunk))
llm_response = await llm_response_task llm_response = await llm_response_task
async for audio_chunk in tts.run_tts(llm_response): async for audio_chunk in tts.run_tts(llm_response):
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio_chunk)) transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio_chunk))
# wait for the output queue to be empty, then leave the meeting # wait for the output queue to be empty, then leave the meeting

View File

@@ -98,12 +98,12 @@ async def main(room_url):
data = await month_data_task data = await month_data_task
transport.output_queue.put( transport.output_queue.put(
[ [
QueueFrame(FrameType.IMAGE_FRAME, data["image"]), QueueFrame(FrameType.IMAGE, data["image"]),
QueueFrame(FrameType.AUDIO_FRAME, data["audio"][0]), QueueFrame(FrameType.AUDIO, data["audio"][0]),
] ]
) )
for audio in data["audio"][1:]: for audio in data["audio"][1:]:
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio)) transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
# wait for the output queue to be empty, then leave the meeting # wait for the output queue to be empty, then leave the meeting
transport.output_queue.join() transport.output_queue.join()

View File

@@ -32,7 +32,11 @@ async def main(room_url:str, token):
] ]
sentence = "" sentence = ""
async for message in transport.get_transcriptions(): async for frame in transport.get_media_frames():
if frame.frame_type != FrameType.TRANSCRIPTION:
continue
message = frame.frame_data
if message["session_id"] == transport.my_participant_id: if message["session_id"] == transport.my_participant_id:
continue continue
@@ -46,7 +50,7 @@ async def main(room_url:str, token):
async for response in llm.run_llm_async_sentences(messages): async for response in llm.run_llm_async_sentences(messages):
full_response += response full_response += response
async for audio in tts.run_tts(response): async for audio in tts.run_tts(response):
transport.output_queue.put(QueueFrame(FrameType.AUDIO_FRAME, audio)) await transport.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
messages.append({"role": "assistant", "content": full_response}) messages.append({"role": "assistant", "content": full_response})