diff --git a/src/dailyai/pipeline/aggregators.py b/src/dailyai/pipeline/aggregators.py index fd9f5a849..f8816da44 100644 --- a/src/dailyai/pipeline/aggregators.py +++ b/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__( diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 129eed5cf..61f043bee 100644 --- a/src/dailyai/services/ai_services.py +++ b/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): diff --git a/src/dailyai/services/azure_ai_services.py b/src/dailyai/services/azure_ai_services.py index 1f3e67cd2..a69095d02 100644 --- a/src/dailyai/services/azure_ai_services.py +++ b/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()) diff --git a/src/dailyai/services/base_transport_service.py b/src/dailyai/services/base_transport_service.py index 76b6d8646..3cec85d8e 100644 --- a/src/dailyai/services/base_transport_service.py +++ b/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): diff --git a/src/dailyai/services/open_ai_services.py b/src/dailyai/services/open_ai_services.py index 2aea60fa9..05e1e9860 100644 --- a/src/dailyai/services/open_ai_services.py +++ 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 diff --git a/src/examples/foundational/07-interruptible.py b/src/examples/foundational/07-interruptible.py index cb722db93..9c22a1f17 100644 --- a/src/examples/foundational/07-interruptible.py +++ b/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