From a3ac0d84e8293a6fae86614b66b6c9372a80fe72 Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Wed, 17 Jan 2024 13:50:55 -0500 Subject: [PATCH] working on making services more consistent/terse/easy --- src/dailyai/services/ai_services.py | 101 ++++++++++++------ src/dailyai/tests/test_ai_services.py | 129 +++++++++++++++++++++++ src/dailyai/tests/test_asyncprocessor.py | 9 +- 3 files changed, 203 insertions(+), 36 deletions(-) create mode 100644 src/dailyai/tests/test_ai_services.py diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 69988a3de..0a82fd821 100644 --- a/src/dailyai/services/ai_services.py +++ b/src/dailyai/services/ai_services.py @@ -5,58 +5,95 @@ import re from dailyai.queue_frame import QueueFrame, FrameType from abc import abstractmethod -from typing import AsyncGenerator +from typing import AsyncGenerator, Iterable from dataclasses import dataclass +from typing import AsyncGenerator + +from collections.abc import Iterable, AsyncIterable class AIService: def __init__( - self, - input_queue: asyncio.Queue[QueueFrame] | None = None, - output_queue: asyncio.Queue[QueueFrame] | None = None, + self ): 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.") + def allowed_input_frame_types(self) -> set[FrameType]: + return set() - while True: - frame = await self.input_queue.get() - self.logger.debug(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 + def possible_output_frame_types(self) -> set[FrameType]: + return set() - output_frame = await self.process_frame(frame) - if output_frame: - await self.output_queue.put(output_frame) - self.input_queue.task_done() + async def run( + self, + requested_frame_types:set[FrameType], + frames:Iterable[QueueFrame] | AsyncIterable[QueueFrame] + ) -> AsyncGenerator[QueueFrame, None]: + if self.possible_output_frame_types().intersection(requested_frame_types) == set(): + raise Exception(f"Requested frame types {requested_frame_types} are not supported by this service.") + + if isinstance(frames, AsyncIterable): + async for frame in frames: + output_frame: QueueFrame | None = await self.process_frame(requested_frame_types, frame) + if output_frame: + yield output_frame + elif isinstance(frames, Iterable): + for frame in frames: + output_frame = await self.process_frame(requested_frame_types, frame) + if output_frame: + yield output_frame + else: + raise Exception("Frames must be an iterable or async iterable") @abstractmethod - async def process_frame(self, frame) -> QueueFrame | None: + async def process_frame(self, requested_frame_types:set[FrameType], frame:QueueFrame) -> QueueFrame | None: pass +class SentenceAggregator(AIService): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.current_sentence = "" + + def allowed_input_frame_types(self) -> set[FrameType]: + return set([FrameType.TEXT_CHUNK, FrameType.SENTENCE]) + + def possible_output_frame_types(self) -> set[FrameType]: + return set([FrameType.SENTENCE]) + + async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> QueueFrame | None: + if not FrameType.SENTENCE in requested_frame_types: + return None + + if frame.frame_type == FrameType.TEXT_CHUNK: + if type(frame.frame_data) != str: + raise Exception("Sentence aggregator requires a string for the data field") + + self.current_sentence += frame.frame_data + if self.current_sentence.endswith((".", "?", "!")): + sentence = self.current_sentence + self.current_sentence = "" + return QueueFrame(FrameType.SENTENCE, sentence) + return None + elif frame.frame_type == FrameType.END_STREAM: + if self.current_sentence: + return QueueFrame(FrameType.SENTENCE, self.current_sentence) + else: + return None + elif frame.frame_type == FrameType.SENTENCE: + return frame + else: + return None + 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]: - # Adding a yield here lets the linter know what this method actually does - yield "" + def allowed_input_frame_types(self) -> set[FrameType]: + return set([FrameType.LLM_MESSAGE, FrameType.SENTENCE, FrameType.TRANSCRIPTION]) - # Generate a responses to a prompt. Returns the response - @abstractmethod - async def run_llm( - self, messages - ) -> str or None: - pass + def allowed_output_frame_types(self) -> set[FrameType]: + return set([FrameType.SENTENCE, FrameType.SENTENCE, FrameType.TEXT_CHUNK]) async def run_llm_async_sentences(self, messages) -> AsyncGenerator[str, None]: current_text = "" diff --git a/src/dailyai/tests/test_ai_services.py b/src/dailyai/tests/test_ai_services.py new file mode 100644 index 000000000..6467442a1 --- /dev/null +++ b/src/dailyai/tests/test_ai_services.py @@ -0,0 +1,129 @@ +from re import A +import unittest + +from typing import AsyncGenerator, Generator + +from dailyai.services.ai_services import AIService, SentenceAggregator +from dailyai.queue_frame import QueueFrame, FrameType + +class SimpleAIService(AIService): + def allowed_input_frame_types(self) -> set[FrameType]: + return set([FrameType.TEXT_CHUNK]) + + def possible_output_frame_types(self) -> set[FrameType]: + return set([FrameType.TEXT_CHUNK]) + + async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> QueueFrame | None: + return frame + +class TestBaseAIService(unittest.IsolatedAsyncioTestCase): + async def test_async_input(self): + service = SimpleAIService() + + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello"), + QueueFrame(FrameType.END_STREAM, None), + ] + async def iterate_frames() -> AsyncGenerator[QueueFrame, None]: + for frame in input_frames: + yield frame + + output_frames = [] + async for frame in service.run(set([FrameType.TEXT_CHUNK]), iterate_frames()): + output_frames.append(frame) + + self.assertEqual(input_frames, output_frames) + + async def test_nonasync_input(self): + service = SimpleAIService() + + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello"), + QueueFrame(FrameType.END_STREAM, None), + ] + + def iterate_frames() -> Generator[QueueFrame, None, None]: + for frame in input_frames: + yield frame + + output_frames = [] + async for frame in service.run(set([FrameType.TEXT_CHUNK]), iterate_frames()): + output_frames.append(frame) + + self.assertEqual(input_frames, output_frames) + + +class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase): + async def test_clause(self) -> None: + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello"), + QueueFrame(FrameType.END_STREAM, None), + ] + + service = SentenceAggregator() + output_frames = [] + async for frame in service.run(set([FrameType.SENTENCE]), input_frames): + output_frames.append(frame) + + self.assertEqual(1, len(output_frames)) + self.assertEqual(QueueFrame(FrameType.SENTENCE, "hello"), output_frames[0]) + + async def test_sentence(self) -> None: + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello, "), + QueueFrame(FrameType.TEXT_CHUNK, "world."), + QueueFrame(FrameType.END_STREAM, None), + ] + + service = SentenceAggregator() + output_frames = [] + async for frame in service.run(set([FrameType.SENTENCE]), input_frames): + output_frames.append(frame) + + self.assertEqual(1, len(output_frames)) + self.assertEqual(QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0]) + + async def test_sentence_and_clause(self) -> None: + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello, "), + QueueFrame(FrameType.TEXT_CHUNK, "world."), + QueueFrame(FrameType.TEXT_CHUNK, " How are"), + QueueFrame(FrameType.END_STREAM, None), + ] + + service = SentenceAggregator() + output_frames = [] + async for frame in service.run(set([FrameType.SENTENCE]), input_frames): + output_frames.append(frame) + + self.assertEqual(2, len(output_frames)) + self.assertEqual( + QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0] + ) + self.assertEqual( + QueueFrame(FrameType.SENTENCE, " How are"), output_frames[1] + ) + + async def test_two_sentences(self) -> None: + input_frames = [ + QueueFrame(FrameType.TEXT_CHUNK, "hello, "), + QueueFrame(FrameType.TEXT_CHUNK, "world."), + QueueFrame(FrameType.TEXT_CHUNK, " How are"), + QueueFrame(FrameType.TEXT_CHUNK, " you doing?"), + QueueFrame(FrameType.END_STREAM, None), + ] + + service = SentenceAggregator() + output_frames = [] + async for frame in service.run(set([FrameType.SENTENCE]), input_frames): + output_frames.append(frame) + + self.assertEqual(2, len(output_frames)) + self.assertEqual( + QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0] + ) + self.assertEqual(QueueFrame(FrameType.SENTENCE, " How are you doing?"), output_frames[1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/dailyai/tests/test_asyncprocessor.py b/src/dailyai/tests/test_asyncprocessor.py index 9c2222791..fcb2781e4 100644 --- a/src/dailyai/tests/test_asyncprocessor.py +++ b/src/dailyai/tests/test_asyncprocessor.py @@ -18,7 +18,7 @@ from dailyai.services.ai_services import ( LLMService, TTSService, ) - +""" class MockTTSService(TTSService): def run_tts(self, sentence): for word in sentence.split(' '): @@ -73,7 +73,7 @@ class TestResponse(unittest.TestCase): while expected_words: actual_word:QueueFrame = output_queue.get() word = expected_words.pop(0) - self.assertEqual(actual_word.frame_type, FrameType.AUDIO) + self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME) self.assertEqual(actual_word.frame_data, bytes(word, "utf-8")) output_queue.task_done() @@ -128,10 +128,10 @@ class TestResponse(unittest.TestCase): while expected_words and not stop_processing_output_queue.is_set(): try: actual_word:QueueFrame = output_queue.get_nowait() - if actual_word.frame_type == FrameType.AUDIO: + if actual_word.frame_type == FrameType.AUDIO_FRAME: time.sleep(0.1) word = expected_words.pop(0) - self.assertEqual(actual_word.frame_type, FrameType.AUDIO) + self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME) self.assertEqual(actual_word.frame_data, bytes(word, "utf-8")) output_queue.task_done() except Empty: @@ -177,3 +177,4 @@ class TestResponse(unittest.TestCase): if __name__ == '__main__': unittest.main() +"""