110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
from typing import AsyncGenerator, Callable
|
|
from dailyai.pipeline.frame_processor import FrameProcessor
|
|
from dailyai.pipeline.frames import (
|
|
Frame,
|
|
LLMResponseEndFrame,
|
|
LLMResponseStartFrame,
|
|
OpenAILLMContextFrame,
|
|
TextFrame,
|
|
TranscriptionFrame,
|
|
UserStartedSpeakingFrame,
|
|
UserStoppedSpeakingFrame,
|
|
)
|
|
from dailyai.services.openai_llm_context import OpenAILLMContext
|
|
|
|
from openai.types.chat import ChatCompletionRole
|
|
|
|
|
|
class OpenAIContextAggregator(FrameProcessor):
|
|
|
|
def __init__(
|
|
self,
|
|
context: OpenAILLMContext,
|
|
aggregator: Callable[[Frame, str | None], str | None],
|
|
role: ChatCompletionRole,
|
|
start_frame: type,
|
|
end_frame: type,
|
|
accumulator_frame: type,
|
|
pass_through=True,
|
|
):
|
|
if not (
|
|
issubclass(start_frame, Frame)
|
|
and issubclass(end_frame, Frame)
|
|
and issubclass(accumulator_frame, Frame)
|
|
):
|
|
raise TypeError(
|
|
"start_frame, end_frame and accumulator_frame must be instances of Frame"
|
|
)
|
|
|
|
self._context: OpenAILLMContext = context
|
|
self._aggregator: Callable[[Frame, str | None], None] = aggregator
|
|
self._role: ChatCompletionRole = role
|
|
self._start_frame = start_frame
|
|
self._end_frame = end_frame
|
|
self._accumulator_frame = accumulator_frame
|
|
self._pass_through = pass_through
|
|
|
|
self._aggregating = False
|
|
self._aggregation = None
|
|
|
|
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
|
|
if isinstance(frame, self._start_frame):
|
|
self._aggregating = True
|
|
elif isinstance(frame, self._end_frame):
|
|
self._aggregating = False
|
|
if self._aggregation:
|
|
self._context.add_message(
|
|
{
|
|
"role": self._role,
|
|
"content": self._aggregation,
|
|
"name": self._role,
|
|
} # type: ignore
|
|
)
|
|
self._aggregation = None
|
|
yield OpenAILLMContextFrame(self._context)
|
|
elif isinstance(frame, self._accumulator_frame) and self._aggregating:
|
|
self._aggregation = self._aggregator(frame, self._aggregation)
|
|
if self._pass_through:
|
|
yield frame
|
|
else:
|
|
yield frame
|
|
|
|
def string_aggregator(
|
|
self,
|
|
frame: Frame,
|
|
aggregation: str | None) -> str | None:
|
|
if not isinstance(frame, TextFrame):
|
|
raise TypeError(
|
|
"Frame must be a TextFrame instance to be aggregated by a string aggregator."
|
|
)
|
|
if not aggregation:
|
|
aggregation = ""
|
|
return " ".join([aggregation, frame.text])
|
|
|
|
|
|
class OpenAIUserContextAggregator(OpenAIContextAggregator):
|
|
def __init__(self, context: OpenAILLMContext):
|
|
super().__init__(
|
|
context=context,
|
|
aggregator=self.string_aggregator,
|
|
role="user",
|
|
start_frame=UserStartedSpeakingFrame,
|
|
end_frame=UserStoppedSpeakingFrame,
|
|
accumulator_frame=TranscriptionFrame,
|
|
pass_through=False,
|
|
)
|
|
|
|
|
|
class OpenAIAssistantContextAggregator(OpenAIContextAggregator):
|
|
|
|
def __init__(self, context: OpenAILLMContext):
|
|
super().__init__(
|
|
context,
|
|
aggregator=self.string_aggregator,
|
|
role="assistant",
|
|
start_frame=LLMResponseStartFrame,
|
|
end_frame=LLMResponseEndFrame,
|
|
accumulator_frame=TextFrame,
|
|
pass_through=True,
|
|
)
|