diff --git a/src/dailyai/pipeline/aggregators.py b/src/dailyai/pipeline/aggregators.py index f8816da44..7b7056869 100644 --- a/src/dailyai/pipeline/aggregators.py +++ b/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 diff --git a/src/dailyai/pipeline/frame_processor.py b/src/dailyai/pipeline/frame_processor.py index d70294ac0..3d361b4d9 100644 --- a/src/dailyai/pipeline/frame_processor.py +++ b/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.""" diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 0fca44884..6a692fb52 100644 --- a/src/dailyai/services/ai_services.py +++ b/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)])