Fix interruptible pipeline runner and aggregator.

src/dailyai/pipeline/aggregators.py

@@ -11,6 +11,7 @@ from dailyai.pipeline.frames import (
     LLMMessagesQueueFrame,
     LLMResponseEndFrame,
     Frame,
+    LLMResponseStartFrame,
     TextFrame,
     TranscriptionQueueFrame,
 )
@@ -19,6 +20,28 @@ from dailyai.services.ai_services import AIService
 
 from typing import AsyncGenerator, Coroutine, List, Text
 
+class LLMResponseAggregator(FrameProcessor):
+    def __init__(self, messages: list[dict]):
+        self.aggregation = ""
+        self.aggregating = False
+        self.messages = messages
+
+    async def process_frame(
+        self, frame: Frame
+    ) -> AsyncGenerator[Frame, None]:
+        if isinstance(frame, LLMResponseStartFrame):
+            self.aggregating = True
+        elif isinstance(frame, LLMResponseEndFrame):
+            self.aggregating = False
+            self.messages.append({"role": "assistant", "content": self.aggregation})
+            self.aggregation = ""
+            yield LLMMessagesQueueFrame(self.messages)
+        elif isinstance(frame, TextFrame) and self.aggregating:
+            self.aggregation += frame.text
+            yield frame
+        else:
+            yield frame
+
 
 class LLMContextAggregator(AIService):
     def __init__(

src/dailyai/services/ai_services.py

@@ -104,6 +104,12 @@ class TTSService(AIService):
         yield bytes()
 
     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
+
         if not isinstance(frame, TextFrame):
             yield frame
             return
@@ -121,10 +127,8 @@ class TTSService(AIService):
             async for audio_chunk in self.run_tts(text):
                 yield AudioFrame(audio_chunk)
 
-    async def finalize(self):
-        if self.current_sentence:
-            async for audio_chunk in self.run_tts(self.current_sentence):
-                yield AudioFrame(audio_chunk)
+            # note we pass along the text frame *after* the audio, so the text frame is completed after the audio is processed.
+            yield TextFrame(text)
 
     # Convenience function to send the audio for a sentence to the given queue
     async def say(self, sentence, queue: asyncio.Queue):

src/dailyai/services/azure_ai_services.py

@@ -112,11 +112,9 @@ class AzureImageGenServiceREST(ImageGenService):
         async with self._aiohttp_session.post(
             url, headers=headers, json=body
         ) as submission:
-            print(f"submission: {submission}")
             # We never get past this line, because this header isn't
             # defined on a 429 response, but something is eating our exceptions!
             operation_location = submission.headers['operation-location']
-            print(f"submission status: {submission.status}")
             status = ""
             attempts_left = 120
             json_response = None
@@ -139,5 +137,4 @@ class AzureImageGenServiceREST(ImageGenService):
             async with self._aiohttp_session.get(image_url) as response:
                 image_stream = io.BytesIO(await response.content.read())
                 image = Image.open(image_stream)
-                print("i got an image file!")
                 return (image_url, image.tobytes())

src/dailyai/services/base_transport_service.py

@@ -10,6 +10,7 @@ import threading
 import time
 from typing import AsyncGenerator
 from enum import Enum
+from dailyai.pipeline.frame_processor import FrameProcessor
 
 from dailyai.pipeline.frames import (
     AudioFrame,
@@ -180,7 +181,13 @@ class BaseTransportService():
         pipeline.set_source(self.receive_queue)
         await pipeline.run_pipeline()
 
-    async def run_interruptible_pipeline(self, pipeline: Pipeline, allow_interruptions=True, pre_processor=None, post_processor=None):
+    async def run_interruptible_pipeline(
+        self,
+        pipeline: Pipeline,
+        allow_interruptions=True,
+        pre_processor=None,
+        post_processor: FrameProcessor | None = None,
+    ):
         pipeline.set_sink(self.send_queue)
         source_queue = asyncio.Queue()
         pipeline.set_source(source_queue)
@@ -190,24 +197,24 @@ class BaseTransportService():
         async def yield_frame(frame: Frame) -> AsyncGenerator[Frame, None]:
             yield frame
 
-        async def post_process(post_processor):
-            if not post_processor:
-                return
-
+        async def post_process(post_processor: FrameProcessor):
             while True:
                 frame = await self.completed_queue.get()
-                print("post-processing frame: ", frame.__class__.__name__)
-                await post_processor.process_frame(frame)
+
+                # We ignore the output of the post_processor's process frame;
+                # this is called to update the post-processor's state.
+                async for frame in post_processor.process_frame(frame):
+                    pass
 
                 if isinstance(frame, EndFrame):
                     break
 
-        post_process_task = asyncio.create_task(post_process(post_processor))
+        if post_processor:
+            post_process_task = asyncio.create_task(post_process(post_processor))
 
         started = False
 
         async for frame in self.get_receive_frames():
-            print("Got frame: ", frame.__class__.__name__)
             if isinstance(frame, UserStartedSpeakingFrame):
                 pipeline_task.cancel()
                 self.interrupt()
@@ -425,11 +432,6 @@ class BaseTransportService():
                         elif len(b):
                             self.write_frame_to_mic(bytes(b))
                             b = bytearray()
-
-                        if self._loop:
-                            asyncio.run_coroutine_threadsafe(
-                                self.completed_queue.put(frame), self._loop
-                            )
                     else:
                         # if there are leftover audio bytes, write them now; failing to do so
                         # can cause static in the audio stream.
@@ -442,6 +444,11 @@ class BaseTransportService():
                         if isinstance(frame, StartFrame):
                             self._is_interrupted.clear()
 
+                    if self._loop:
+                        asyncio.run_coroutine_threadsafe(
+                            self.completed_queue.put(frame), self._loop
+                        )
+
                 self._threadsafe_send_queue.task_done()
             except queue.Empty:
                 if len(b):

src/dailyai/services/open_ai_services.py

@@ -57,7 +57,6 @@ class OpenAIImageGenService(ImageGenService):
     ):
         super().__init__(image_size=image_size)
         self._model = model
-        print(f"api key: {api_key}")
         self._client = AsyncOpenAI(api_key=api_key)
         self._aiohttp_session = aiohttp_session
 

src/examples/foundational/07-interruptible.py

@@ -1,7 +1,7 @@
 import asyncio
 import aiohttp
 import os
-from dailyai.pipeline.aggregators import LLMAssistantContextAggregator, LLMUserContextAggregator
+from dailyai.pipeline.aggregators import LLMAssistantContextAggregator, LLMResponseAggregator, LLMUserContextAggregator
 
 from dailyai.pipeline.pipeline import Pipeline
 from dailyai.services.ai_services import FrameLogger
@@ -46,8 +46,8 @@ async def main(room_url: str, token):
 
             await transport.run_interruptible_pipeline(
                 pipeline,
-                post_processor=LLMAssistantContextAggregator(
-                    messages, transport._my_participant_id
+                post_processor=LLMResponseAggregator(
+                    messages
                 ),
                 pre_processor=LLMUserContextAggregator(
                     messages, transport._my_participant_id, complete_sentences=False