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,
-        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 = ""