wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #39 from daily-co/docstrings

Browse Source 
Docstrings
This commit is contained in:
Moishe Lettvin
2024-03-07 15:39:50 -05:00
committed by GitHub
parent eb430621f1 b14f08a7d5
commit 2bcb4966d3
5 changed files with 163 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
src/dailyai/pipeline/aggregators.py
View File

@@ -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

36
src/dailyai/pipeline/frame_processor.py
View File

@@ -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
@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()
yield frame
@abstractmethod
async def interrupted(self) -> None:
"""Handle any cleanup if the pipeline was interrupted."""
pass

37
src/dailyai/pipeline/pipeline.py
View File

@@ -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")

5
src/dailyai/services/ai_services.py
View File

@@ -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

8
src/dailyai/tests/test_aggregators.py
View File

@@ -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