Remove Queue in frame names
This commit is contained in:
@@ -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}")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user