Getting started on docstrings

src/dailyai/pipeline/aggregators.py

@@ -1,11 +1,9 @@
 import asyncio
 import re
 
-from tblib import Frame
 from dailyai.pipeline.frame_processor import FrameProcessor
 
 from dailyai.pipeline.frames import (
-    ControlFrame,
     EndPipeFrame,
     EndFrame,
     LLMMessagesQueueFrame,
@@ -94,13 +92,6 @@ class LLMContextAggregator(AIService):
             self.messages.append({"role": self.role, "content": frame.text})
             yield LLMMessagesQueueFrame(self.messages)
 
-    async def finalize(self) -> AsyncGenerator[Frame, None]:
-        # Send any dangling words that weren't finished with punctuation.
-        if self.complete_sentences and self.sentence:
-            self.messages.append({"role": self.role, "content": self.sentence})
-            yield LLMMessagesQueueFrame(self.messages)
-
-
 class LLMUserContextAggregator(LLMContextAggregator):
     def __init__(
         self, messages: list[dict], bot_participant_id=None, complete_sentences=True
@@ -124,7 +115,22 @@ class LLMAssistantContextAggregator(LLMContextAggregator):
 
 
 class SentenceAggregator(FrameProcessor):
+    """This frame processor aggregates text frames into complete sentences.
 
+    Frame input/output:
+        TextFrame("Hello,") -> None
+        TextFrame(" world.") -> TextFrame("Hello world.")
+
+    Doctest:
+    >>> async def print_frames(aggregator, frame):
+    ...     async for frame in aggregator.process_frame(frame):
+    ...         print(frame.text)
+
+    >>> aggregator = SentenceAggregator()
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Hello,")))
+    >>> asyncio.run(print_frames(aggregator, TextFrame(" world.")))
+    Hello, world.
+    """
     def __init__(self):
         self.aggregation = ""
 
@@ -147,6 +153,41 @@ class SentenceAggregator(FrameProcessor):
 
 
 class LLMFullResponseAggregator(FrameProcessor):
+    """This class aggregates Text frames until it receives a
+    LLMResponseEndFrame, then emits the concatenated text as
+    a single text frame.
+
+    given the following frames:
+
+        TextFrame("Hello,")
+        TextFrame(" world.")
+        TextFrame(" I am")
+        TextFrame(" an LLM.")
+        LLMResponseEndFrame()]
+
+    this processor will yield nothing for the first 4 frames, then
+
+        TextFrame("Hello, world. I am an LLM.")
+        LLMResponseEndFrame()
+
+    when passed the last frame.
+
+    >>> async def print_frames(aggregator, frame):
+    ...     async for frame in aggregator.process_frame(frame):
+    ...         if isinstance(frame, TextFrame):
+    ...             print(frame.text)
+    ...         else:
+    ...             print(frame.__class__.__name__)
+
+    >>> aggregator = LLMFullResponseAggregator()
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Hello,")))
+    >>> asyncio.run(print_frames(aggregator, TextFrame(" world.")))
+    >>> asyncio.run(print_frames(aggregator, TextFrame(" I am")))
+    >>> asyncio.run(print_frames(aggregator, TextFrame(" an LLM.")))
+    >>> asyncio.run(print_frames(aggregator, LLMResponseEndFrame()))
+    Hello, world. I am an LLM.
+    LLMResponseEndFrame
+    """
     def __init__(self):
         self.aggregation = ""
 
@@ -157,12 +198,24 @@ class LLMFullResponseAggregator(FrameProcessor):
             self.aggregation += frame.text
         elif isinstance(frame, LLMResponseEndFrame):
             yield TextFrame(self.aggregation)
+            yield frame
             self.aggregation = ""
         else:
             yield frame
 
 
 class StatelessTextTransformer(FrameProcessor):
+    """This processor calls the given function on any text in a text frame.
+
+    >>> async def print_frames(aggregator, frame):
+    ...     async for frame in aggregator.process_frame(frame):
+    ...         print(frame.text)
+
+    >>> aggregator = StatelessTextTransformer(lambda x: x.upper())
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Hello")))
+    HELLO
+    """
+
     def __init__(self, transform_fn):
         self.transform_fn = transform_fn
 

src/dailyai/pipeline/frame_processor.py

@@ -27,12 +27,6 @@ class FrameProcessor:
             yield frame
         yield frame
 
-    @abstractmethod
-    async def finalize(self) -> AsyncGenerator[Frame, None]:
-        # This is a trick for the interpreter (and linter) to know that this is a generator.
-        if False:
-            yield Frame()
-
     @abstractmethod
     async def interrupted(self) -> None:
         """Handle any cleanup if the pipeline was interrupted."""

src/dailyai/services/ai_services.py

@@ -100,23 +100,14 @@ class TTSService(AIService):
         # yield empty bytes here, so linting can infer what this method does
         yield bytes()
 
-<<<<<<< HEAD
     async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
         if isinstance(frame, EndFrame):
             if self.current_sentence:
                 async for audio_chunk in self.run_tts(self.current_sentence):
                     yield AudioFrame(audio_chunk)
-            yield frame
+                yield TextFrame(self.current_sentence)
 
         if not isinstance(frame, TextFrame):
-=======
-    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
-        if not isinstance(frame, TextQueueFrame):
-            if self.current_sentence:
-                async for audio_chunk in self.run_tts(self.current_sentence):
-                    yield AudioQueueFrame(audio_chunk)
-
->>>>>>> fa5f38c (frame and pipeline docstrings)
             yield frame
             return
 
@@ -133,12 +124,9 @@ class TTSService(AIService):
             async for audio_chunk in self.run_tts(text):
                 yield AudioFrame(audio_chunk)
 
-<<<<<<< HEAD
             # note we pass along the text frame *after* the audio, so the text frame is completed after the audio is processed.
             yield TextFrame(text)
 
-=======
->>>>>>> fa5f38c (frame and pipeline docstrings)
     # Convenience function to send the audio for a sentence to the given queue
     async def say(self, sentence, queue: asyncio.Queue):
         await self.run_to_queue(queue, [TextFrame(sentence)])