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()
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]:
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 not self.is_interrupted.is_set():
if frame:
if frame.frame_type == FrameType.AUDIO_FRAME:
if frame.frame_type == FrameType.AUDIO:
chunk = frame.frame_data
all_audio_frames.extend(chunk)
@@ -385,7 +385,7 @@ class Orchestrator(EventHandler):
if l:
self.mic.write_frames(bytes(b[:l]))
b = b[l:]
elif frame.frame_type == FrameType.IMAGE_FRAME:
elif frame.frame_type == FrameType.IMAGE:
self.set_image(frame.frame_data)
elif len(b):
self.mic.write_frames(bytes(b))

View File

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

View File

@@ -73,13 +73,13 @@ class LLMService(AIService):
if not self.output_queue:
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:
raise Exception("LLM service requires a dict for the data field")
messages: list[dict[str, str]] = frame.frame_data
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):
@@ -98,13 +98,13 @@ class TTSService(AIService):
if not self.output_queue:
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:
raise Exception("TTS service requires a string for the data field")
text = frame.frame_data
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):

View File

@@ -35,7 +35,11 @@ class DailyTransportService(EventHandler):
self.duration: float = duration
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.stop_threads = Event()
self.story_started = False
@@ -45,6 +49,9 @@ class DailyTransportService(EventHandler):
self.camera_height = 768
self.camera_enabled = False
self.output_queue = asyncio.Queue()
self.media_queue = asyncio.Queue()
self.other_participant_has_joined = False
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.event_handlers = {}
@@ -182,24 +183,25 @@ class DailyTransportService(EventHandler):
)
if self.token:
self.transcription_queue = asyncio.Queue()
self.client.start_transcription(self.transcription_settings)
self.my_participant_id = self.client.participants()["local"]["id"]
async def get_transcriptions(self):
async def get_media_frames(self):
while True:
transcript = await self.transcription_queue.get()
yield transcript
frame = await self.media_queue.get()
yield frame
if frame.frame_type == FrameType.END_STREAM:
break
def get_async_output_queue(self):
return self.async_output_queue
return self.output_queue
async def marshal_frames(self):
while True:
frame = await self.async_output_queue.get()
self.output_queue.put(frame)
self.async_output_queue.task_done()
frame = await self.output_queue.get()
self.threadsafe_output_queue.put(frame)
self.output_queue.task_done()
if frame.frame_type == FrameType.END_STREAM:
break
@@ -222,7 +224,8 @@ class DailyTransportService(EventHandler):
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
if self.camera_thread and self.camera_thread.is_alive():
@@ -258,9 +261,10 @@ class DailyTransportService(EventHandler):
def on_app_message(self, message, sender):
pass
def on_transcription_message(self, message):
def on_transcription_message(self, message:dict):
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):
pass
@@ -291,7 +295,7 @@ class DailyTransportService(EventHandler):
all_audio_frames = bytearray()
while True:
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:
frames: list[QueueFrame] = [frames_or_frame]
elif type(frames_or_frame) == list:
@@ -302,13 +306,13 @@ class DailyTransportService(EventHandler):
for frame in frames:
if frame.frame_type == FrameType.END_STREAM:
self.logger.info("Stopping frame consumer thread")
self.output_queue.task_done()
self.threadsafe_output_queue.task_done()
return
# if interrupted, we just pull frames off the queue and discard them
if not self.is_interrupted.is_set():
if frame:
if frame.frame_type == FrameType.AUDIO_FRAME:
if frame.frame_type == FrameType.AUDIO:
chunk = frame.frame_data
all_audio_frames.extend(chunk)
@@ -318,7 +322,7 @@ class DailyTransportService(EventHandler):
if l:
self.mic.write_frames(bytes(b[:l]))
b = b[l:]
elif frame.frame_type == FrameType.IMAGE_FRAME:
elif frame.frame_type == FrameType.IMAGE:
self.set_image(frame.frame_data)
elif len(b):
self.mic.write_frames(bytes(b))
@@ -333,7 +337,7 @@ class DailyTransportService(EventHandler):
if frame.frame_type == FrameType.START_STREAM:
self.is_interrupted.clear()
self.output_queue.task_done()
self.threadsafe_output_queue.task_done()
except Empty:
try:
if len(b):

View File

@@ -73,7 +73,7 @@ class TestResponse(unittest.TestCase):
while expected_words:
actual_word:QueueFrame = output_queue.get()
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"))
output_queue.task_done()
@@ -128,10 +128,10 @@ class TestResponse(unittest.TestCase):
while expected_words and not stop_processing_output_queue.is_set():
try:
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)
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"))
output_queue.task_done()
except Empty:

View File

@@ -40,7 +40,7 @@ class StaticSpriteResponse(OrchestratorResponse):
self.image_bytes = img.tobytes()
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):
@@ -73,8 +73,8 @@ class AnimatedSpriteLLMResponse(LLMResponse):
def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]:
return [
QueueFrame(FrameType.AUDIO_FRAME, audio_frame),
QueueFrame(FrameType.IMAGE_FRAME, random.choice(self.image_bytes))
QueueFrame(FrameType.AUDIO, audio_frame),
QueueFrame(FrameType.IMAGE, random.choice(self.image_bytes))
]

View File

@@ -38,7 +38,7 @@ async def main(room_url):
return
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
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']}!")
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")
@transport.event_handler("on_call_state_updated")

View File

@@ -29,7 +29,7 @@ async def main(room_url):
"role": "system",
"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))
llm_task = asyncio.create_task(llm.run())

View File

@@ -27,7 +27,7 @@ async def main(room_url):
@transport.event_handler("on_participant_joined")
async def on_participant_joined(transport, participant):
(_, 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()

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."):
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
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

View File

@@ -98,12 +98,12 @@ async def main(room_url):
data = await month_data_task
transport.output_queue.put(
[
QueueFrame(FrameType.IMAGE_FRAME, data["image"]),
QueueFrame(FrameType.AUDIO_FRAME, data["audio"][0]),
QueueFrame(FrameType.IMAGE, data["image"]),
QueueFrame(FrameType.AUDIO, data["audio"][0]),
]
)
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
transport.output_queue.join()

View File

@@ -32,7 +32,11 @@ async def main(room_url:str, token):
]
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:
continue
@@ -46,7 +50,7 @@ async def main(room_url:str, token):
async for response in llm.run_llm_async_sentences(messages):
full_response += 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})