working on making services more consistent/terse/easy

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,
+        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:
+    def allowed_input_frame_types(self) -> set[FrameType]:
-        if self.input_queue is None or self.output_queue is None:
+        return set()
-            raise Exception("Input and output queues must be set before using the run method.")
 
-        while True:
+    def possible_output_frame_types(self) -> set[FrameType]:
-            frame = await self.input_queue.get()
+        return set()
-            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
 
-            output_frame = await self.process_frame(frame)
+    async def run(
-            if output_frame:
+            self,
-                await self.output_queue.put(output_frame)
+            requested_frame_types:set[FrameType],
-            self.input_queue.task_done()
+            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.
+    def allowed_input_frame_types(self) -> set[FrameType]:
-    @abstractmethod
+        return set([FrameType.LLM_MESSAGE, FrameType.SENTENCE, FrameType.TRANSCRIPTION])
-    async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
-        # Adding a yield here lets the linter know what this method actually does
-        yield ""
 
-    # Generate a responses to a prompt. Returns the response
+    def allowed_output_frame_types(self) -> set[FrameType]:
-    @abstractmethod
+        return set([FrameType.SENTENCE, FrameType.SENTENCE, FrameType.TEXT_CHUNK])
-    async def run_llm(
-        self, messages
-    ) -> str or None:
-        pass
 
     async def run_llm_async_sentences(self, messages) -> AsyncGenerator[str, None]:
         current_text = ""

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()

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()
+"""