wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #8 from daily-co/queueframe-refactor

Browse Source 
Refactor QueueFrame
This commit is contained in:
Moishe Lettvin
2024-01-23 13:15:11 -05:00
committed by GitHub
parent b51abd2969 ba89e41c5b
commit 690cf2e47d
10 changed files with 113 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
src/dailyai/services/ai_services.py
View File

@@ -1,10 +1,14 @@
import asyncio
import logging
import re
from httpx import request
from dailyai.queue_frame import QueueFrame, FrameType
from dailyai.queue_frame import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
LLMMessagesQueueFrame,
QueueFrame,
TextQueueFrame,
)
from abc import abstractmethod
from typing import AsyncGenerator, AsyncIterable, Iterable
@@ -24,7 +28,7 @@ class AIService:
await queue.put(frame)
if add_end_of_stream:
await queue.put(QueueFrame(FrameType.END_STREAM, None))
await queue.put(EndStreamQueueFrame())
async def run(
self,
@@ -46,7 +50,7 @@ class AIService:
frame = await frames.get()
async for output_frame in self.process_frame(frame):
yield output_frame
if frame.frame_type == FrameType.END_STREAM:
if isinstance(frame, EndStreamQueueFrame):
break
else:
raise Exception("Frames must be an iterable or async iterable")
@@ -61,21 +65,15 @@ class AIService:
async def process_frame(self, frame:QueueFrame) -> AsyncGenerator[QueueFrame, None]:
# This is a trick for the interpreter (and linter) to know that this is a generator.
if False:
yield QueueFrame(FrameType.NOOP, None)
yield QueueFrame()
@abstractmethod
async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
# This is a trick for the interpreter (and linter) to know that this is a generator.
if False:
yield QueueFrame(FrameType.NOOP, None)
yield QueueFrame()
class LLMService(AIService):
def allowed_input_frame_types(self) -> set[FrameType]:
return set([FrameType.LLM_MESSAGE])
def allowed_output_frame_types(self) -> set[FrameType]:
return set([FrameType.TEXT])
@abstractmethod
async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
yield ""
@@ -85,13 +83,9 @@ class LLMService(AIService):
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
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 text_chunk in self.run_llm_async(messages):
yield QueueFrame(FrameType.TEXT, text_chunk)
if isinstance(frame, LLMMessagesQueueFrame):
async for text_chunk in self.run_llm_async(frame.messages):
yield TextQueueFrame(text_chunk)
class TTSService(AIService):
@@ -112,34 +106,31 @@ class TTSService(AIService):
yield bytes()
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if frame.frame_type != FrameType.TEXT:
if not isinstance(frame, TextQueueFrame):
yield frame
return
if not isinstance(frame.frame_data, str):
raise(Exception(f"Invalid data type in frame type: {frame.frame_type}, type: {type(frame.frame_data)}"))
text: str | None = None
if not self.aggregate_sentences:
text = frame.frame_data
text = frame.text
else:
self.current_sentence += frame.frame_data
self.current_sentence += frame.text
if self.current_sentence.endswith((".", "?", "!")):
text = self.current_sentence
self.current_sentence = ""
if text:
async for audio_chunk in self.run_tts(text):
yield QueueFrame(FrameType.AUDIO, audio_chunk)
yield AudioQueueFrame(audio_chunk)
async def finalize(self):
if self.current_sentence:
async for audio_chunk in self.run_tts(self.current_sentence):
yield QueueFrame(FrameType.AUDIO, audio_chunk)
yield AudioQueueFrame(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, [QueueFrame(FrameType.TEXT, sentence)])
await self.run_to_queue(queue, [TextQueueFrame(sentence)])
class ImageGenService(AIService):
@@ -149,15 +140,15 @@ class ImageGenService(AIService):
# Renders the image. Returns an Image object.
@abstractmethod
async def run_image_gen(self, sentence) -> tuple[str, bytes]:
async def run_image_gen(self, sentence:str) -> tuple[str, bytes]:
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if type(frame.frame_data) != str:
raise Exception("Image service requires a string for the data field")
if not isinstance(frame, TextQueueFrame):
return
(_, image_data) = await self.run_image_gen(frame.frame_data)
yield QueueFrame(FrameType.IMAGE, image_data)
(url, image_data) = await self.run_image_gen(frame.text)
yield ImageQueueFrame(url, image_data)
@dataclass

42
src/dailyai/services/daily_transport_service.py
View File

@@ -8,9 +8,16 @@ import types
from functools import partial
from queue import Queue, Empty
from dailyai.queue_frame import QueueFrame, FrameType
from dailyai.queue_frame import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
QueueFrame,
StartStreamQueueFrame,
TranscriptionQueueFrame,
)
from threading import Thread, Event, Timer
from threading import Thread, Event
from daily import (
EventHandler,
@@ -201,7 +208,7 @@ class DailyTransportService(EventHandler):
while True:
frame = await self.receive_queue.get()
yield frame
if frame.frame_type == FrameType.END_STREAM:
if isinstance(frame, EndStreamQueueFrame):
break
def get_async_send_queue(self):
@@ -212,7 +219,7 @@ class DailyTransportService(EventHandler):
frame: QueueFrame | list = await self.send_queue.get()
self.threadsafe_send_queue.put(frame)
self.send_queue.task_done()
if type(frame) == QueueFrame and frame.frame_type == FrameType.END_STREAM:
if isinstance(frame, EndStreamQueueFrame):
break
async def wait_for_send_queue_to_empty(self):
@@ -242,8 +249,8 @@ class DailyTransportService(EventHandler):
self.stop_threads.set()
await self.receive_queue.put(QueueFrame(FrameType.END_STREAM, None))
await self.send_queue.put(QueueFrame(FrameType.END_STREAM, None))
await self.receive_queue.put(EndStreamQueueFrame())
await self.send_queue.put(EndStreamQueueFrame())
await async_output_queue_marshal_task
if self.camera_thread and self.camera_thread.is_alive():
@@ -281,7 +288,12 @@ class DailyTransportService(EventHandler):
def on_transcription_message(self, message:dict):
if self.loop:
frame = QueueFrame(FrameType.TRANSCRIPTION, message)
participantId = ""
if "participantId" in message:
participantId = message["participantId"]
elif "session_id" in message:
participantId = message["session_id"]
frame = TranscriptionQueueFrame(message["text"], participantId, message["timestamp"])
asyncio.run_coroutine_threadsafe(self.receive_queue.put(frame), self.loop)
def on_transcription_stopped(self, stopped_by, stopped_by_error):
@@ -314,15 +326,15 @@ class DailyTransportService(EventHandler):
while True:
try:
frames_or_frame: QueueFrame | list[QueueFrame] = self.threadsafe_send_queue.get()
if type(frames_or_frame) == QueueFrame:
if isinstance(frames_or_frame, QueueFrame):
frames: list[QueueFrame] = [frames_or_frame]
elif type(frames_or_frame) == list:
elif isinstance(frames_or_frame, list):
frames: list[QueueFrame] = frames_or_frame
else:
raise Exception("Unknown type in output queue")
for frame in frames:
if frame.frame_type == FrameType.END_STREAM:
if isinstance(frame, EndStreamQueueFrame):
self.logger.info("Stopping frame consumer thread")
self.threadsafe_send_queue.task_done()
return
@@ -330,8 +342,8 @@ class DailyTransportService(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:
chunk = frame.frame_data
if isinstance(frame, AudioQueueFrame):
chunk = frame.data
all_audio_frames.extend(chunk)
@@ -340,8 +352,8 @@ class DailyTransportService(EventHandler):
if l:
self.mic.write_frames(bytes(b[:l]))
b = b[l:]
elif frame.frame_type == FrameType.IMAGE:
self.set_image(frame.frame_data)
elif isinstance(frame, ImageQueueFrame):
self.set_image(frame.image)
elif len(b):
self.mic.write_frames(bytes(b))
b = bytearray()
@@ -352,7 +364,7 @@ class DailyTransportService(EventHandler):
)
self.interrupt_time = None
if frame.frame_type == FrameType.START_STREAM:
if isinstance(frame, StartStreamQueueFrame):
self.is_interrupted.clear()
self.threadsafe_send_queue.task_done()