Remove Queue in frame names

This commit is contained in:
Moishe Lettvin
2024-03-06 14:09:06 -05:00
parent b9556716dd
commit 62fd371b97
25 changed files with 241 additions and 240 deletions

View File

@@ -6,14 +6,14 @@ import wave
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
AudioFrame,
EndFrame,
ImageFrame,
LLMMessagesQueueFrame,
LLMResponseEndQueueFrame,
LLMResponseStartQueueFrame,
QueueFrame,
TextQueueFrame,
LLMResponseEndFrame,
LLMResponseStartFrame,
Frame,
TextFrame,
TranscriptionQueueFrame,
)
@@ -33,14 +33,14 @@ class AIService(FrameProcessor):
await queue.put(frame)
if add_end_of_stream:
await queue.put(EndStreamQueueFrame())
await queue.put(EndFrame())
async def run(
self,
frames: Iterable[QueueFrame]
| AsyncIterable[QueueFrame]
| asyncio.Queue[QueueFrame],
) -> AsyncGenerator[QueueFrame, None]:
frames: Iterable[Frame]
| AsyncIterable[Frame]
| asyncio.Queue[Frame],
) -> AsyncGenerator[Frame, None]:
try:
if isinstance(frames, AsyncIterable):
async for frame in frames:
@@ -55,7 +55,7 @@ class AIService(FrameProcessor):
frame = await frames.get()
async for output_frame in self.process_frame(frame):
yield output_frame
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
else:
raise Exception("Frames must be an iterable or async iterable")
@@ -76,12 +76,12 @@ class LLMService(AIService):
async def run_llm(self, messages) -> str:
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, LLMMessagesQueueFrame):
yield LLMResponseStartQueueFrame()
yield LLMResponseStartFrame()
async for text_chunk in self.run_llm_async(frame.messages):
yield TextQueueFrame(text_chunk)
yield LLMResponseEndQueueFrame()
yield TextFrame(text_chunk)
yield LLMResponseEndFrame()
else:
yield frame
@@ -103,8 +103,8 @@ class TTSService(AIService):
# yield empty bytes here, so linting can infer what this method does
yield bytes()
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not isinstance(frame, TextQueueFrame):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if not isinstance(frame, TextFrame):
yield frame
return
@@ -119,16 +119,16 @@ class TTSService(AIService):
if text:
async for audio_chunk in self.run_tts(text):
yield AudioQueueFrame(audio_chunk)
yield AudioFrame(audio_chunk)
async def finalize(self):
if self.current_sentence:
async for audio_chunk in self.run_tts(self.current_sentence):
yield AudioQueueFrame(audio_chunk)
yield AudioFrame(audio_chunk)
# Convenience function to send the audio for a sentence to the given queue
async def say(self, sentence, queue: asyncio.Queue):
await self.run_to_queue(queue, [TextQueueFrame(sentence)])
await self.run_to_queue(queue, [TextFrame(sentence)])
class ImageGenService(AIService):
@@ -141,13 +141,13 @@ class ImageGenService(AIService):
async def run_image_gen(self, sentence: str) -> tuple[str, bytes]:
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not isinstance(frame, TextQueueFrame):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if not isinstance(frame, TextFrame):
yield frame
return
(url, image_data) = await self.run_image_gen(frame.text)
yield ImageQueueFrame(url, image_data)
yield ImageFrame(url, image_data)
class STTService(AIService):
@@ -164,9 +164,9 @@ class STTService(AIService):
"""Returns transcript as a string"""
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
"""Processes a frame of audio data, either buffering or transcribing it."""
if not isinstance(frame, AudioQueueFrame):
if not isinstance(frame, AudioFrame):
return
data = frame.data
@@ -187,8 +187,8 @@ class FrameLogger(AIService):
super().__init__(**kwargs)
self.prefix = prefix
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, (AudioQueueFrame, ImageQueueFrame)):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, (AudioFrame, ImageFrame)):
self.logger.info(f"{self.prefix}: {type(frame)}")
else:
print(f"{self.prefix}: {frame}")

View File

@@ -12,12 +12,12 @@ from typing import AsyncGenerator
from enum import Enum
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
QueueFrame,
SpriteQueueFrame,
StartStreamQueueFrame,
AudioFrame,
EndFrame,
ImageFrame,
Frame,
SpriteFrame,
StartFrame,
TranscriptionQueueFrame,
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame
@@ -159,7 +159,7 @@ class BaseTransportService():
self._stop_threads.set()
await self.send_queue.put(EndStreamQueueFrame())
await self.send_queue.put(EndFrame())
await async_output_queue_marshal_task
await self.send_queue.join()
self._frame_consumer_thread.join()
@@ -182,7 +182,7 @@ class BaseTransportService():
pipeline.set_sink(self.send_queue)
pipeline_task = asyncio.create_task(pipeline.run_pipeline())
async def yield_frame(frame:QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def yield_frame(frame:Frame) -> AsyncGenerator[Frame, None]:
yield frame
async def post_process(post_processor):
@@ -194,7 +194,7 @@ class BaseTransportService():
print("post-processing frame: ", frame.__class__.__name__)
await post_processor.process_frame(frame)
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
post_process_task = asyncio.create_task(post_process(post_processor))
@@ -214,7 +214,7 @@ class BaseTransportService():
async for frame in frame_generator:
await source_queue.put(frame)
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
await asyncio.gather(pipeline_task, post_process_task)
@@ -303,20 +303,20 @@ class BaseTransportService():
async def _marshal_frames(self):
while True:
frame: QueueFrame | list = await self.send_queue.get()
frame: Frame | list = await self.send_queue.get()
self._threadsafe_send_queue.put(frame)
self.send_queue.task_done()
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
def interrupt(self):
self._is_interrupted.set()
async def get_receive_frames(self) -> AsyncGenerator[QueueFrame, None]:
async def get_receive_frames(self) -> AsyncGenerator[Frame, None]:
while True:
frame = await self.receive_queue.get()
yield frame
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
def _receive_audio(self):
@@ -329,13 +329,13 @@ class BaseTransportService():
while not self._stop_threads.is_set():
buffer = self.read_audio_frames(desired_frame_count)
if len(buffer) > 0:
frame = AudioQueueFrame(buffer)
frame = AudioFrame(buffer)
asyncio.run_coroutine_threadsafe(
self.receive_queue.put(frame), self._loop
)
asyncio.run_coroutine_threadsafe(
self.receive_queue.put(EndStreamQueueFrame()), self._loop
self.receive_queue.put(EndFrame()), self._loop
)
def _set_image(self, image: bytes):
@@ -363,18 +363,18 @@ class BaseTransportService():
all_audio_frames = bytearray()
while True:
try:
frames_or_frame: QueueFrame | list[QueueFrame] = (
frames_or_frame: Frame | list[Frame] = (
self._threadsafe_send_queue.get()
)
if isinstance(frames_or_frame, QueueFrame):
frames: list[QueueFrame] = [frames_or_frame]
if isinstance(frames_or_frame, Frame):
frames: list[Frame] = [frames_or_frame]
elif isinstance(frames_or_frame, list):
frames: list[QueueFrame] = frames_or_frame
frames: list[Frame] = frames_or_frame
else:
raise Exception("Unknown type in output queue")
for frame in frames:
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
self._logger.info("Stopping frame consumer thread")
self._threadsafe_send_queue.task_done()
if self._loop:
@@ -386,7 +386,7 @@ class BaseTransportService():
# if interrupted, we just pull frames off the queue and discard them
if not self._is_interrupted.is_set():
if frame:
if isinstance(frame, AudioQueueFrame):
if isinstance(frame, AudioFrame):
chunk = frame.data
all_audio_frames.extend(chunk)
@@ -398,9 +398,9 @@ class BaseTransportService():
if truncated_length:
self.write_frame_to_mic(bytes(b[:truncated_length]))
b = b[truncated_length:]
elif isinstance(frame, ImageQueueFrame):
elif isinstance(frame, ImageFrame):
self._set_image(frame.image)
elif isinstance(frame, SpriteQueueFrame):
elif isinstance(frame, SpriteFrame):
self._set_images(frame.images)
elif len(b):
self.write_frame_to_mic(bytes(b))
@@ -418,7 +418,7 @@ class BaseTransportService():
self.write_frame_to_mic(bytes(b[:truncated_length]))
b = bytearray()
if isinstance(frame, StartStreamQueueFrame):
if isinstance(frame, StartFrame):
self._is_interrupted.clear()
self._threadsafe_send_queue.task_done()

View File

@@ -4,7 +4,7 @@ import math
import time
from typing import AsyncGenerator
import wave
from dailyai.pipeline.frames import AudioQueueFrame, QueueFrame, TranscriptionQueueFrame
from dailyai.pipeline.frames import AudioFrame, Frame, TranscriptionQueueFrame
from dailyai.services.ai_services import STTService
@@ -39,9 +39,9 @@ class LocalSTTService(STTService):
ww.setframerate(self._frame_rate)
self._wave = ww
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
"""Processes a frame of audio data, either buffering or transcribing it."""
if not isinstance(frame, AudioQueueFrame):
if not isinstance(frame, AudioFrame):
return
data = frame.data