Adding queue transportation to services

This commit is contained in:
Moishe Lettvin
2024-01-11 19:14:19 -05:00
parent 7ca7764be3
commit b9b82695c6
18 changed files with 194 additions and 87 deletions

View File

@@ -9,7 +9,7 @@ from queue import Queue, PriorityQueue, Empty
from threading import Event, Semaphore, Thread
from typing import Any, Generator, Iterator, Optional, Type
from dailyai.output_queue import OutputQueueFrame, FrameType
from dailyai.queue_frame import QueueFrame, FrameType
from dailyai.message_handler.message_handler import MessageHandler
from dailyai.services.ai_services import AIServiceConfig
@@ -268,10 +268,10 @@ class LLMResponse(OrchestratorResponse):
if out.strip():
yield out.strip()
def get_frames_from_tts_response(self, audio_frame) -> list[OutputQueueFrame]:
return [OutputQueueFrame(FrameType.AUDIO_FRAME, audio_frame)]
def get_frames_from_tts_response(self, audio_frame) -> list[QueueFrame]:
return [QueueFrame(FrameType.AUDIO_FRAME, audio_frame)]
def get_frames_from_chunk(self, chunk) -> Generator[list[OutputQueueFrame], Any, None]:
def get_frames_from_chunk(self, chunk) -> Generator[list[QueueFrame], Any, None]:
for audio_frame in self.services.tts.run_tts(chunk):
yield self.get_frames_from_tts_response(audio_frame)
@@ -317,7 +317,7 @@ class LLMResponse(OrchestratorResponse):
break
if not self.has_sent_first_frame:
self.output_queue.put(OutputQueueFrame(FrameType.START_STREAM, None))
self.output_queue.put(QueueFrame(FrameType.START_STREAM, None))
self.has_sent_first_frame = True
for frame in frames:

View File

@@ -15,7 +15,7 @@ from dailyai.async_processor.async_processor import (
OrchestratorResponse,
LLMResponse,
)
from dailyai.output_queue import OutputQueueFrame, FrameType
from dailyai.queue_frame import QueueFrame, FrameType
from dailyai.services.ai_services import AIServiceConfig
from dailyai.message_handler.message_handler import MessageHandler
@@ -197,7 +197,7 @@ class Orchestrator(EventHandler):
self.logger.info("Camera thread stopped")
self.logger.info("Put stop in output queue")
self.output_queue.put(OutputQueueFrame(FrameType.END_STREAM, None))
self.output_queue.put(QueueFrame(FrameType.END_STREAM, None))
self.frame_consumer_thread.join()
self.logger.info("Orchestrator stopped.")
@@ -367,7 +367,7 @@ class Orchestrator(EventHandler):
all_audio_frames = bytearray()
while True:
try:
frame:OutputQueueFrame = self.output_queue.get()
frame:QueueFrame = self.output_queue.get()
if frame.frame_type == FrameType.END_STREAM:
self.logger.info("Stopping frame consumer thread")
return

View File

@@ -1,14 +0,0 @@
from enum import Enum
from dataclasses import dataclass
class FrameType(Enum):
AUDIO_FRAME = 1
IMAGE_FRAME = 2
START_STREAM = 3
END_STREAM = 4
@dataclass(frozen=True)
class OutputQueueFrame:
frame_type: FrameType
frame_data: bytes | None

View File

@@ -0,0 +1,18 @@
from enum import Enum
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
IMAGE_DESCRIPTION = 8
@dataclass(frozen=True)
class QueueFrame:
frame_type: FrameType
frame_data: str | dict | bytes | list | None

View File

@@ -1,23 +1,56 @@
import asyncio
import logging
import re
from tkinter import END
from dailyai.queue_frame import QueueFrame, FrameType
from asyncio import Queue
from abc import abstractmethod
from typing import AsyncGenerator
from dataclasses import dataclass
class AIService:
def __init__(self):
self.logger = logging.getLogger("dailyai")
def close(self):
def __init__(
self,
input_queue: asyncio.Queue[QueueFrame] | None = None,
output_queue: asyncio.Queue[QueueFrame] | None = None,
):
self.logger = logging.getLogger("dailyai")
self.input_queue: asyncio.Queue[QueueFrame] | None = input_queue
self.output_queue: asyncio.Queue[QueueFrame] | None = output_queue
def stop(self):
pass
async def run(self) -> None:
if self.input_queue is None or self.output_queue is None:
raise Exception("Input and output queues must be set before using the run method.")
while True:
frame = await self.input_queue.get()
print(f"{self.__class__.__name__} got frame:", frame.frame_type)
if frame.frame_type == FrameType.END_STREAM:
self.input_queue.task_done()
await self.output_queue.put(QueueFrame(FrameType.END_STREAM, None))
break
output_frame = await self.process_frame(frame)
if output_frame:
await self.output_queue.put(output_frame)
self.input_queue.task_done()
@abstractmethod
async def process_frame(self, frame) -> QueueFrame | None:
pass
class LLMService(AIService):
# Generate a set of responses to a prompt. Yields a list of responses.
@abstractmethod
async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
pass
# Adding a yield here lets the linter know what this method actually does
yield ""
# Generate a responses to a prompt. Returns the response
@abstractmethod
@@ -26,6 +59,30 @@ class LLMService(AIService):
) -> str or None:
pass
async def run_llm_async_sentences(self, messages) -> AsyncGenerator[str, None]:
current_text = ""
async for text in self.run_llm_async(messages):
current_text += text
if re.match(r"^.*[.!?]$", text):
yield current_text
current_text = ""
if current_text:
yield current_text
async def process_frame(self, frame:QueueFrame) -> QueueFrame | None:
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 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):
print("got message", message)
await self.output_queue.put(QueueFrame(FrameType.SENTENCE_FRAME, message))
class TTSService(AIService):
# Some TTS services require a specific sample rate. We default to 16k
@@ -36,7 +93,21 @@ class TTSService(AIService):
# be sent to the microphone device
@abstractmethod
async def run_tts(self, sentence) -> AsyncGenerator[bytes, None]:
pass
# yield empty bytes here, so linting can infer what this method does
yield bytes()
async def process_frame(self, frame:QueueFrame) -> QueueFrame | None:
if not self.output_queue:
raise Exception("Output queue must be set before using the run method.")
print(frame.frame_type)
if frame.frame_type == FrameType.SENTENCE_FRAME:
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))
class ImageGenService(AIService):

View File

@@ -16,8 +16,8 @@ from PIL import Image
from azure.cognitiveservices.speech import SpeechSynthesizer, SpeechConfig, ResultReason, CancellationReason
class AzureTTSService(TTSService):
def __init__(self, speech_key=None, speech_region=None):
super().__init__()
def __init__(self, input_queue=None, output_queue=None, speech_key=None, speech_region=None):
super().__init__(input_queue, output_queue)
speech_key = speech_key or os.getenv("AZURE_SPEECH_SERVICE_KEY")
speech_region = speech_region or os.getenv("AZURE_SPEECH_SERVICE_REGION")
@@ -48,8 +48,8 @@ class AzureTTSService(TTSService):
self.logger.info("Error details: {}".format(cancellation_details.error_details))
class AzureLLMService(LLMService):
def __init__(self, api_key=None, azure_endpoint=None, api_version=None, model=None):
super().__init__()
def __init__(self, input_queue=None, output_queue=None, api_key=None, azure_endpoint=None, api_version=None, model=None):
super().__init__(input_queue, output_queue)
api_key = api_key or os.getenv("AZURE_CHATGPT_KEY")
azure_endpoint = azure_endpoint or os.getenv("AZURE_CHATGPT_ENDPOINT")

View File

@@ -7,7 +7,7 @@ import types
from functools import partial
from queue import Queue, Empty
from dailyai.output_queue import OutputQueueFrame, FrameType
from dailyai.queue_frame import QueueFrame, FrameType
from threading import Thread, Event, Timer
@@ -48,6 +48,12 @@ class DailyTransportService(EventHandler):
self.camera_thread = None
self.frame_consumer_thread = None
# 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 = {}
@@ -162,6 +168,7 @@ class DailyTransportService(EventHandler):
)
if self.token:
self.transcription_queue = Queue()
self.client.start_transcription(
{
"language": "en",
@@ -178,11 +185,29 @@ class DailyTransportService(EventHandler):
self.my_participant_id = self.client.participants()["local"]["id"]
def get_transcriptions(self):
while True:
transcript = self.transcription_queue.get()
yield transcript
def get_async_output_queue(self):
return self.async_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()
if frame.frame_type == FrameType.END_STREAM:
break
async def run(self) -> None:
self.configure_daily()
self.participant_left = False
async_output_queue_marshal_task = asyncio.create_task(self.marshal_frames())
try:
participant_count: int = len(self.client.participants())
self.logger.info(f"{participant_count} participants in room")
@@ -194,10 +219,13 @@ class DailyTransportService(EventHandler):
self.client.leave()
self.stop_threads.set()
await self.async_output_queue.put(QueueFrame(FrameType.END_STREAM, None))
await async_output_queue_marshal_task
if self.camera_thread and self.camera_thread.is_alive():
self.camera_thread.join()
if self.frame_consumer_thread and self.frame_consumer_thread.is_alive():
self.output_queue.put(OutputQueueFrame(FrameType.END_STREAM, None))
self.frame_consumer_thread.join()
def stop(self):
@@ -224,6 +252,7 @@ class DailyTransportService(EventHandler):
pass
def on_transcription_message(self, message):
self.transcription_queue.put(message["text"])
pass
def on_transcription_stopped(self, stopped_by, stopped_by_error):
@@ -255,11 +284,11 @@ class DailyTransportService(EventHandler):
all_audio_frames = bytearray()
while True:
try:
frames_or_frame: OutputQueueFrame | list[OutputQueueFrame] = self.output_queue.get()
if type(frames_or_frame) == OutputQueueFrame:
frames: list[OutputQueueFrame] = [frames_or_frame]
frames_or_frame: QueueFrame | list[QueueFrame] = self.output_queue.get()
if type(frames_or_frame) == QueueFrame:
frames: list[QueueFrame] = [frames_or_frame]
elif type(frames_or_frame) == list:
frames: list[OutputQueueFrame] = frames_or_frame
frames: list[QueueFrame] = frames_or_frame
else:
raise Exception("Unknown type in output queue")

View File

@@ -9,11 +9,11 @@ from dailyai.services.ai_services import TTSService
class ElevenLabsTTSService(TTSService):
def __init__(self):
super().__init__()
def __init__(self, input_queue, output_queue, api_key=None, voice_id=None):
super().__init__(input_queue, output_queue)
self.api_key = os.getenv("ELEVENLABS_API_KEY")
self.voice_id = os.getenv("ELEVENLABS_VOICE_ID")
self.api_key = api_key or os.getenv("ELEVENLABS_API_KEY")
self.voice_id = voice_id or os.getenv("ELEVENLABS_VOICE_ID")
async def run_tts(self, sentence) -> AsyncGenerator[bytes, None]:
async with aiohttp.ClientSession() as session:

View File

@@ -11,7 +11,7 @@ from dailyai.async_processor.async_processor import (
LLMResponse,
)
from dailyai.message_handler.message_handler import MessageHandler
from dailyai.output_queue import OutputQueueFrame, FrameType
from dailyai.queue_frame import QueueFrame, FrameType
from dailyai.services.ai_services import (
AIServiceConfig,
ImageGenService,
@@ -71,7 +71,7 @@ class TestResponse(unittest.TestCase):
output_queue.task_done()
while expected_words:
actual_word:OutputQueueFrame = output_queue.get()
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_data, bytes(word, "utf-8"))
@@ -127,7 +127,7 @@ class TestResponse(unittest.TestCase):
expected_words = ["Hello", "there.", "How", "are", "you?", "I", "hope", "you", "are", "well."]
while expected_words and not stop_processing_output_queue.is_set():
try:
actual_word:OutputQueueFrame = output_queue.get_nowait()
actual_word:QueueFrame = output_queue.get_nowait()
if actual_word.frame_type == FrameType.AUDIO_FRAME:
time.sleep(0.1)
word = expected_words.pop(0)