Merge pull request #39 from daily-co/docstrings

Docstrings

src/dailyai/pipeline/aggregators.py

@@ -1,16 +1,15 @@
 import asyncio
 import re
 
-from tblib import Frame
 from dailyai.pipeline.frame_processor import FrameProcessor
 
 from dailyai.pipeline.frames import (
-    ControlFrame,
-    EndPipeFrame,
     EndFrame,
+    EndPipeFrame,
+    Frame,
+    ImageFrame,
     LLMMessagesQueueFrame,
     LLMResponseEndFrame,
-    Frame,
     LLMResponseStartFrame,
     TextFrame,
     TranscriptionQueueFrame,
@@ -18,7 +17,7 @@ from dailyai.pipeline.frames import (
 from dailyai.pipeline.pipeline import Pipeline
 from dailyai.services.ai_services import AIService
 
-from typing import AsyncGenerator, Coroutine, List, Text
+from typing import AsyncGenerator, Coroutine, List
 
 class LLMResponseAggregator(FrameProcessor):
     def __init__(self, messages: list[dict]):
@@ -94,13 +93,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 +116,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 +154,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 +199,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
 
@@ -177,6 +231,23 @@ class StatelessTextTransformer(FrameProcessor):
             yield frame
 
 class ParallelPipeline(FrameProcessor):
+    """ Run multiple pipelines in parallel.
+
+    This class takes frames from its source queue and sends them to each
+    sub-pipeline. Each sub-pipeline emits its frames into this class's
+    sink queue. No guarantees are made about the ordering of frames in
+    the sink queue (that is, no sub-pipeline has higher priority than
+    any other, frames are put on the sink in the order they're emitted
+    by the sub-pipelines).
+
+    After each frame is taken from this class's source queue and placed
+    in each sub-pipeline's source queue, an EndPipeFrame is put on each
+    sub-pipeline's source queue. This indicates to the sub-pipe runner
+    that it should exit.
+
+    Since frame handlers pass through unhandled frames by convention, this
+    class de-dupes frames in its sink before yielding them.
+    """
     def __init__(self, pipeline_definitions: List[List[FrameProcessor]]):
         self.sources = [asyncio.Queue() for _ in pipeline_definitions]
         self.sink: asyncio.Queue[Frame] = asyncio.Queue()
@@ -213,6 +284,30 @@ class ParallelPipeline(FrameProcessor):
                 yield frame
 
 class GatedAggregator(FrameProcessor):
+    """Accumulate frames, with custom functions to start and stop accumulation.
+    Yields gate-opening frame before any accumulated frames, then ensuing frames
+    until and not including the gate-closed 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 = GatedAggregator(
+    ...     gate_close_fn=lambda x: isinstance(x, LLMResponseStartFrame),
+    ...     gate_open_fn=lambda x: isinstance(x, ImageFrame),
+    ...     start_open=False)
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Hello")))
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Hello again.")))
+    >>> asyncio.run(print_frames(aggregator, ImageFrame(url='', image=bytes([]))))
+    ImageFrame
+    Hello
+    Hello again.
+    >>> asyncio.run(print_frames(aggregator, TextFrame("Goodbye.")))
+    Goodbye.
+    """
     def __init__(self, gate_open_fn, gate_close_fn, start_open):
         self.gate_open_fn = gate_open_fn
         self.gate_close_fn = gate_close_fn

src/dailyai/pipeline/frame_processor.py

@@ -3,35 +3,31 @@ from typing import AsyncGenerator
 
 from dailyai.pipeline.frames import ControlFrame, Frame
 
-"""
-This is the base class for all frame processors. Frame processors consume a frame
-and yield 0 or more frames. Generally frame processors are used as part of a pipeline,
-where frames come from a source queue, are processed by a series of frame processors,
-then placed on a sink queue.
-
-By convention, FrameProcessors should immediately yield any frames they don't process.
-
-Stateful FrameProcessors should watch for the EndStreamQueueFrame and finalize their
-output, eg. yielding an unfinished sentence if they're aggregating LLM output to full
-sentences. EndStreamQueueFrame is also a chance to clean up any services that need to
-be closed, del'd, etc.
-"""
 
 class FrameProcessor:
+    """This is the base class for all frame processors. Frame processors consume a frame
+    and yield 0 or more frames. Generally frame processors are used as part of a pipeline
+    where frames come from a source queue, are processed by a series of frame processors,
+    then placed on a sink queue.
+
+    By convention, FrameProcessors should immediately yield any frames they don't process.
+
+    Stateful FrameProcessors should watch for the EndStreamQueueFrame and finalize their
+    output, eg. yielding an unfinished sentence if they're aggregating LLM output to full
+    sentences. EndStreamQueueFrame is also a chance to clean up any services that need to
+    be closed, del'd, etc.
+    """
+
     @abstractmethod
     async def process_frame(
         self, frame: Frame
     ) -> AsyncGenerator[Frame, None]:
+        """Process a single frame and yield 0 or more frames."""
         if isinstance(frame, ControlFrame):
             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."""
         pass
-

src/dailyai/pipeline/pipeline.py

@@ -4,14 +4,14 @@ from dailyai.pipeline.frame_processor import FrameProcessor
 
 from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
 
-"""
-This class manages a pipe of FrameProcessors, and runs them in sequence. The "source"
-and "sink" queues are managed by the caller. You can use this class stand-alone to
-perform specialized processing, or you can use the Transport's run_pipeline method to
-instantiate and run a pipeline with the Transport's sink and source queues.
-"""
 
 class Pipeline:
+    """
+    This class manages a pipe of FrameProcessors, and runs them in sequence. The "source"
+    and "sink" queues are managed by the caller. You can use this class stand-alone to
+    perform specialized processing, or you can use the Transport's run_pipeline method to
+    instantiate and run a pipeline with the Transport's sink and source queues.
+    """
 
     def __init__(
         self,
@@ -19,22 +19,47 @@ class Pipeline:
         source: asyncio.Queue | None = None,
         sink: asyncio.Queue[Frame] | None = None,
     ):
+        """ Create a new pipeline. By default neither the source nor sink
+        queues are set, so you'll need to pass them to this constructor or
+        call set_source and set_sink before using the pipeline. Note that
+        the transport's run_*_pipeline methods will set the source and sink
+        queues on the pipeline for you.
+        """
         self.processors = processors
         self.source: asyncio.Queue[Frame] | None = source
         self.sink: asyncio.Queue[Frame] | None = sink
 
     def set_source(self, source: asyncio.Queue[Frame]):
+        """ Set the source queue for this pipeline. Frames from this queue
+        will be processed by each frame_processor in the pipeline, or order
+        from first to last. """
         self.source = source
 
     def set_sink(self, sink: asyncio.Queue[Frame]):
+        """ Set the sink queue for this pipeline. After the last frame_processor
+        has processed a frame, its output will be placed on this queue."""
         self.sink = sink
 
     async def get_next_source_frame(self) -> AsyncGenerator[Frame, None]:
+        """ Convenience function to get the next frame from the source queue. This
+        lets us consistently have an AsyncGenerator yield frames, from either the
+        source queue or a frame_processor."""
         if self.source is None:
             raise ValueError("Source queue not set")
         yield await self.source.get()
 
     async def run_pipeline(self):
+        """ Run the pipeline. Take each frame from the source queue, pass it to
+        the first frame_processor, pass the output of that frame_processor to the
+        next in the list, etc. until the last frame_processor has processed the
+        resulting frames, then place those frames in the sink queue.
+
+        The source and sink queues must be set before calling this method.
+
+        This method will exit when an EndStreamQueueFrame is placed on the sink queue.
+        No more frames will be placed on the sink queue after an EndStreamQueueFrame, even
+        if it's not the last frame yielded by the last frame_processor in the pipeline.."""
+
         if self.source is None or self.sink is None:
             raise ValueError("Source or sink queue not set")
 

src/dailyai/services/ai_services.py

@@ -59,9 +59,6 @@ class AIService(FrameProcessor):
                         break
             else:
                 raise Exception("Frames must be an iterable or async iterable")
-
-            async for output_frame in self.finalize():
-                yield output_frame
         except Exception as e:
             self.logger.error("Exception occurred while running AI service", e)
             raise e
@@ -108,7 +105,7 @@ class TTSService(AIService):
             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):
             yield frame

src/dailyai/tests/test_aggregators.py

@@ -1,4 +1,5 @@
 import asyncio
+import doctest
 import functools
 import unittest
 
@@ -118,3 +119,10 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase):
         while not sink.empty():
             frame = await sink.get()
             self.assertEqual(frame, expected_output_frames.pop(0))
+
+
+def load_tests(loader, tests, ignore):
+    """ Run doctests on the aggregators module. """
+    from dailyai.pipeline import aggregators
+    tests.addTests(doctest.DocTestSuite(aggregators))
+    return tests