use InterimTranscriptionFrame in LLMUserResponseAggregator

This commit is contained in:
Aleix Conchillo Flaqué
2024-04-11 11:21:42 -07:00
parent e288aa047b
commit eb3bf117b1
3 changed files with 73 additions and 22 deletions

View File

@@ -9,6 +9,7 @@ from dailyai.pipeline.frames import (
EndPipeFrame, EndPipeFrame,
Frame, Frame,
ImageFrame, ImageFrame,
InterimTranscriptionFrame,
LLMMessagesFrame, LLMMessagesFrame,
LLMResponseEndFrame, LLMResponseEndFrame,
LLMResponseStartFrame, LLMResponseStartFrame,
@@ -107,8 +108,8 @@ class LLMResponseAggregator(FrameProcessor):
start_frame, start_frame,
end_frame, end_frame,
accumulator_frame, accumulator_frame,
interim_accumulator_frame=None,
pass_through=True, pass_through=True,
end_frame_threshold=0.75,
): ):
self.aggregation = "" self.aggregation = ""
self.aggregating = False self.aggregating = False
@@ -117,42 +118,75 @@ class LLMResponseAggregator(FrameProcessor):
self._start_frame = start_frame self._start_frame = start_frame
self._end_frame = end_frame self._end_frame = end_frame
self._accumulator_frame = accumulator_frame self._accumulator_frame = accumulator_frame
self._interim_accumulator_frame = interim_accumulator_frame
self._pass_through = pass_through self._pass_through = pass_through
self._end_frame_threshold = end_frame_threshold self._seen_start_frame = False
self._last_end_frame_time = 0 self._seen_end_frame = False
self._seen_interim_results = False
# Use cases implemented:
#
# S: Start, E: End, T: Transcription, I: Interim, X: Text
#
# S E -> None
# S T E -> X
# S I T E -> X
# S I E T -> X
# S I E I T -> X
#
# The following case would not be supported:
#
# S I E T1 I T2 -> X
#
# and T2 would be dropped.
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]: async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if not self.messages: if not self.messages:
return return
send_aggregation = False
if isinstance(frame, self._start_frame): if isinstance(frame, self._start_frame):
self._seen_start_frame = True
self.aggregating = True self.aggregating = True
elif isinstance(frame, self._end_frame): elif isinstance(frame, self._end_frame):
self.aggregating = False self._seen_end_frame = True
# Sometimes VAD triggers quickly on and off. If we don't get any transcription,
# it creates empty LLM message queue frames # We might have received the end frame but we might still be
if len(self.aggregation) > 0: # aggregating (i.e. we have seen interim results but not the final
self.messages.append({"role": self._role, "content": self.aggregation}) # text).
self.aggregation = "" self.aggregating = self._seen_interim_results
yield self._end_frame()
yield LLMMessagesFrame(self.messages) # Send the aggregation if we are not aggregating anymore (i.e. no
self._last_end_frame_time = time.time() # more interim results received).
send_aggregation = not self.aggregating
elif isinstance(frame, self._accumulator_frame): elif isinstance(frame, self._accumulator_frame):
# Also accept transcription frames received for a short period after
# the last end frame was received. It might be that transcription
# frames are a bit delayed.
diff_time = time.time() - self._last_end_frame_time
if self.aggregating: if self.aggregating:
self.aggregation += f" {frame.text}" self.aggregation += f" {frame.text}"
elif diff_time <= self._end_frame_threshold: # We have receied a complete sentence, so if we have seen the
self.messages.append({"role": self._role, "content": frame.text}) # end frame and we were still aggregating, it means we should
yield self._end_frame() # send the aggregation.
yield LLMMessagesFrame(self.messages) send_aggregation = self._seen_end_frame
if self._pass_through: if self._pass_through:
yield frame yield frame
# We just got our final result, so let's reset interim results.
self._seen_interim_results = False
elif self._interim_accumulator_frame and isinstance(frame, self._interim_accumulator_frame):
self._seen_interim_results = True
else: else:
yield frame yield frame
if send_aggregation and len(self.aggregation) > 0:
self.messages.append({"role": self._role, "content": self.aggregation})
yield self._end_frame()
yield LLMMessagesFrame(self.messages)
# Reset
self.aggregation = ""
self._seen_start_frame = False
self._seen_end_frame = False
self._seen_interim_results = False
class LLMAssistantResponseAggregator(LLMResponseAggregator): class LLMAssistantResponseAggregator(LLMResponseAggregator):
def __init__(self, messages: list[dict]): def __init__(self, messages: list[dict]):
@@ -173,6 +207,7 @@ class LLMUserResponseAggregator(LLMResponseAggregator):
start_frame=UserStartedSpeakingFrame, start_frame=UserStartedSpeakingFrame,
end_frame=UserStoppedSpeakingFrame, end_frame=UserStoppedSpeakingFrame,
accumulator_frame=TranscriptionFrame, accumulator_frame=TranscriptionFrame,
interim_accumulator_frame=InterimTranscriptionFrame,
pass_through=False, pass_through=False,
) )

View File

@@ -164,6 +164,17 @@ class TranscriptionFrame(TextFrame):
return f"{self.__class__.__name__}, text: '{self.text}' participantId: {self.participantId}, timestamp: {self.timestamp}" return f"{self.__class__.__name__}, text: '{self.text}' participantId: {self.participantId}, timestamp: {self.timestamp}"
@dataclass()
class InterimTranscriptionFrame(TextFrame):
"""A text frame with interim transcription-specific data. Will be placed in
the transport's receive queue when a participant speaks."""
participantId: str
timestamp: str
def __str__(self):
return f"{self.__class__.__name__}, text: '{self.text}' participantId: {self.participantId}, timestamp: {self.timestamp}"
class TTSStartFrame(ControlFrame): class TTSStartFrame(ControlFrame):
"""Used to indicate the beginning of a TTS response. Following AudioFrames """Used to indicate the beginning of a TTS response. Following AudioFrames
are part of the TTS response until an TTEndFrame. These frames can be used are part of the TTS response until an TTEndFrame. These frames can be used

View File

@@ -10,6 +10,7 @@ from functools import partial
from typing import Any from typing import Any
from dailyai.pipeline.frames import ( from dailyai.pipeline.frames import (
InterimTranscriptionFrame,
ReceivedAppMessageFrame, ReceivedAppMessageFrame,
TranscriptionFrame, TranscriptionFrame,
UserImageFrame, UserImageFrame,
@@ -368,8 +369,12 @@ class DailyTransport(ThreadedTransport, EventHandler):
elif "session_id" in message: elif "session_id" in message:
participantId = message["session_id"] participantId = message["session_id"]
if self._my_participant_id and participantId != self._my_participant_id: if self._my_participant_id and participantId != self._my_participant_id:
frame = TranscriptionFrame( is_final = message["rawResponse"]["is_final"]
message["text"], participantId, message["timestamp"]) if is_final:
frame = TranscriptionFrame(message["text"], participantId, message["timestamp"])
else:
frame = InterimTranscriptionFrame(
message["text"], participantId, message["timestamp"])
asyncio.run_coroutine_threadsafe( asyncio.run_coroutine_threadsafe(
self.receive_queue.put(frame), self._loop) self.receive_queue.put(frame), self._loop)