From 5c7e8eba61a7365cb373d87391dca18fcdce3c74 Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Fri, 29 Dec 2023 16:13:01 -0500 Subject: [PATCH] the actual refactor that I missed in the last commit --- .../async_processor/async_processor.py | 52 +++++++++++++------ src/dailyai/orchestrator.py | 12 ++--- src/dailyai/tests/test_asyncprocessor.py | 8 +-- src/samples/simple-sample/simple-sample.py | 2 +- src/samples/static-sprite/sprite-sample.py | 39 ++++++++++---- 5 files changed, 78 insertions(+), 35 deletions(-) diff --git a/src/dailyai/async_processor/async_processor.py b/src/dailyai/async_processor/async_processor.py index 756e2efa4..a4bf9a83a 100644 --- a/src/dailyai/async_processor/async_processor.py +++ b/src/dailyai/async_processor/async_processor.py @@ -181,20 +181,18 @@ class AsyncProcessor: def async_prepare(self) -> None: self.set_state(AsyncProcessorState.PREPARING) - self.preparation_iterator = self.get_preparation_iterator() + self.start_preparation() self.set_state(AsyncProcessorState.READY) - for chunk in self.preparation_iterator: - if self.state not in [ - AsyncProcessorState.READY, - AsyncProcessorState.PLAYING, - ]: - break - - self.process_chunk(chunk) - + self.continue_preparation() self.logger.info(f"Preparation done for {self.__class__.__name__}") self.preparation_done() + def start_preparation(self) -> None: + pass + + def continue_preparation(self) -> None: + pass + def preparation_done(self): pass @@ -213,8 +211,8 @@ class AsyncProcessor: def do_finalization(self) -> None: pass +class OrchestratorResponse(AsyncProcessor): -class Response(AsyncProcessor): def __init__( self, services, @@ -225,6 +223,17 @@ class Response(AsyncProcessor): self.message_handler: MessageHandler = message_handler self.output_queue: Queue = output_queue + + +class LLMResponse(OrchestratorResponse): + def __init__( + self, + services, + message_handler, + output_queue, + ) -> None: + super().__init__(services, message_handler, output_queue) + self.has_sent_first_frame = False self.chunks_in_preparation = Queue() @@ -263,6 +272,19 @@ class Response(AsyncProcessor): for audio_frame in self.services.tts.run_tts(chunk): yield self.get_frames_from_tts_response(audio_frame) + def start_preparation(self) -> None: + self.preparation_iterator = self.get_preparation_iterator() + + def continue_preparation(self) -> None: + for chunk in self.preparation_iterator: + if self.state not in [ + AsyncProcessorState.READY, + AsyncProcessorState.PLAYING, + ]: + break + + self.process_chunk(chunk) + def process_chunk(self, chunk) -> None: self.chunks_in_preparation.put((chunk, self.get_frames_from_chunk(chunk))) @@ -320,7 +342,7 @@ class Response(AsyncProcessor): @dataclass(frozen=True) class ConversationProcessorCollection: - introduction: Optional[Type[Response]] = None - waiting: Optional[Type[Response]] = None - response: Optional[Type[Response]] = None - goodbye: Optional[Type[Response]] = None + introduction: Optional[Type[OrchestratorResponse]] = None + waiting: Optional[Type[OrchestratorResponse]] = None + response: Optional[Type[OrchestratorResponse]] = None + goodbye: Optional[Type[OrchestratorResponse]] = None diff --git a/src/dailyai/orchestrator.py b/src/dailyai/orchestrator.py index 872f9ba1e..920af0422 100644 --- a/src/dailyai/orchestrator.py +++ b/src/dailyai/orchestrator.py @@ -11,7 +11,8 @@ from dailyai.async_processor.async_processor import ( AsyncProcessor, AsyncProcessorState, ConversationProcessorCollection, - Response, + OrchestratorResponse, + LLMResponse, ) from dailyai.services.ai_services import AIServiceConfig from dailyai.message_handler.message_handler import MessageHandler @@ -42,9 +43,9 @@ class OrchestratorConfig: # constructor. The dataclass is defined with Frozen=True, so this should # be safe. default_conversation_collection = ConversationProcessorCollection( - introduction=Response, + introduction=LLMResponse, waiting=None, - response=Response, + response=LLMResponse, goodbye=None, ) @@ -207,7 +208,6 @@ class Orchestrator(EventHandler): def on_response_played(self, response): response.finalize() - self.display_waiting() def on_response_finished(self, response): if not response.was_interrupted: @@ -313,8 +313,8 @@ class Orchestrator(EventHandler): self.message_handler.add_user_message(fragment) - response_type = self.conversation_processors.response or Response - new_response: Response = response_type( + response_type: type[OrchestratorResponse] | type[LLMResponse] = self.conversation_processors.response or LLMResponse + new_response: OrchestratorResponse = response_type( self.services, self.message_handler, self.output_queue ) new_response.set_state_callback( diff --git a/src/dailyai/tests/test_asyncprocessor.py b/src/dailyai/tests/test_asyncprocessor.py index 7c5ed70df..b4a7893b5 100644 --- a/src/dailyai/tests/test_asyncprocessor.py +++ b/src/dailyai/tests/test_asyncprocessor.py @@ -15,7 +15,7 @@ from dailyai.message_handler.message_handler import MessageHandler from dailyai.async_processor.async_processor import ( AsyncProcessor, AsyncProcessorState, - Response, + LLMResponse, ) class MockTTSService(TTSService): @@ -51,7 +51,7 @@ class TestResponse(unittest.TestCase): mock_llm_service = MockLLMService() mock_image_service = MockImageService() message_handler = MessageHandler("Hello World") - processor = Response( + processor = LLMResponse( AIServiceConfig( tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service ), @@ -86,7 +86,7 @@ class TestResponse(unittest.TestCase): mock_llm_service = MockLLMService() mock_image_service = MockImageService() message_handler = MessageHandler("System Message") - processor = Response( + processor = LLMResponse( AIServiceConfig( tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service ), @@ -108,7 +108,7 @@ class TestResponse(unittest.TestCase): mock_llm_service = MockLLMService() mock_image_service = MockImageService() message_handler = MessageHandler("System Message") - processor = Response( + processor = LLMResponse( AIServiceConfig( tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service ), diff --git a/src/samples/simple-sample/simple-sample.py b/src/samples/simple-sample/simple-sample.py index 5bb8a51a6..3a111f067 100644 --- a/src/samples/simple-sample/simple-sample.py +++ b/src/samples/simple-sample/simple-sample.py @@ -6,7 +6,7 @@ import time import urllib.parse from dailyai.async_processor.async_processor import ( - Response, + LLMResponse, ConversationProcessorCollection, ) from dailyai.orchestrator import OrchestratorConfig, Orchestrator diff --git a/src/samples/static-sprite/sprite-sample.py b/src/samples/static-sprite/sprite-sample.py index ba957b8ae..aac0d3d1e 100644 --- a/src/samples/static-sprite/sprite-sample.py +++ b/src/samples/static-sprite/sprite-sample.py @@ -1,26 +1,47 @@ import argparse -from email.mime import image -from re import A +import os import requests import time import urllib.parse +from PIL import Image + from dailyai.async_processor.async_processor import ( - Response, ConversationProcessorCollection, + LLMResponse ) from dailyai.orchestrator import OrchestratorConfig, Orchestrator from dailyai.message_handler.message_handler import MessageHandler from dailyai.services.ai_services import AIServiceConfig from dailyai.services.azure_ai_services import AzureImageGenService, AzureTTSService, AzureLLMService -class SpriteResponse(Response): - def __init__(self, message, image): - super().__init__(message) - self.image = image +class StaticSpriteResponse(LLMResponse): + + def __init__( + self, + services, + message_handler, + output_queue + ) -> None: + super().__init__(services, message_handler, output_queue) + self.image_bytes:bytes | None = None + self.filename = None # override this in subclasses + + def start_preparation(self) -> None: + full_path = os.path.join(os.path.dirname(__file__), "/sprites/") + + with Image.open(full_path) as img: + self.image_bytes = img.tobytes() + + def async_play(self) -> None: + self.output_queue.put({"type": "image_frame", "data": self.image_bytes}) + + +class IntroSpriteResponse(StaticSpriteResponse): + def __init__(self, services, message_handler, output_queue) -> None: + super().__init__(services, message_handler, output_queue) + self.filename = "intro.png" - def get_image(self): - return self.image def add_bot_to_room(room_url, token, expiration) -> None: