Fix interruptible pipeline runner and aggregator.

This commit is contained in:
Moishe Lettvin
2024-03-07 08:31:55 -05:00
parent b4505b7eff
commit d5683c4f24
6 changed files with 55 additions and 25 deletions

View File

@@ -11,6 +11,7 @@ from dailyai.pipeline.frames import (
LLMMessagesQueueFrame, LLMMessagesQueueFrame,
LLMResponseEndFrame, LLMResponseEndFrame,
Frame, Frame,
LLMResponseStartFrame,
TextFrame, TextFrame,
TranscriptionQueueFrame, TranscriptionQueueFrame,
) )
@@ -19,6 +20,28 @@ from dailyai.services.ai_services import AIService
from typing import AsyncGenerator, Coroutine, List, Text 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): class LLMContextAggregator(AIService):
def __init__( def __init__(

View File

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

View File

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

View File

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

View File

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

View File

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