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,
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__(