From 3e738642a7ccd655ac5c48c7c90b5474230ec79b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Thu, 18 Jul 2024 14:52:48 -0700 Subject: [PATCH] processors(realtime-ai): add support for getting/updating LLM context --- src/pipecat/frames/frames.py | 10 ++ .../processors/aggregators/llm_response.py | 10 ++ .../processors/frameworks/realtimeai.py | 109 +++++++++++++----- 3 files changed, 101 insertions(+), 28 deletions(-) diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py index b8a2a6d06..515098a6e 100644 --- a/src/pipecat/frames/frames.py +++ b/src/pipecat/frames/frames.py @@ -158,6 +158,16 @@ class LLMMessagesFrame(DataFrame): messages: List[dict] +@dataclass +class LLMMessagesUpdateFrame(DataFrame): + """A frame containing a list of new LLM messages. These messages will + replace the current context LLM messages and should generate a new + LLMMessagesFrame. + + """ + messages: List[dict] + + @dataclass class TransportMessageFrame(DataFrame): message: Any diff --git a/src/pipecat/processors/aggregators/llm_response.py b/src/pipecat/processors/aggregators/llm_response.py index aa85c5ade..efc8af179 100644 --- a/src/pipecat/processors/aggregators/llm_response.py +++ b/src/pipecat/processors/aggregators/llm_response.py @@ -15,6 +15,7 @@ from pipecat.frames.frames import ( LLMFullResponseEndFrame, LLMFullResponseStartFrame, LLMMessagesFrame, + LLMMessagesUpdateFrame, StartInterruptionFrame, TranscriptionFrame, TextFrame, @@ -120,6 +121,15 @@ class LLMResponseAggregator(FrameProcessor): # Reset anyways self._reset() await self.push_frame(frame, direction) + elif isinstance(frame, LLMMessagesUpdateFrame): + # We push the frame downstream so the assistant aggregator gets + # updated as well. + await self.push_frame(frame) + # We can now reset this one. + self._reset() + self._messages = frame.messages + messages_frame = LLMMessagesFrame(self._messages) + await self.push_frame(messages_frame) else: await self.push_frame(frame, direction) diff --git a/src/pipecat/processors/frameworks/realtimeai.py b/src/pipecat/processors/frameworks/realtimeai.py index b68ce2ee5..688538924 100644 --- a/src/pipecat/processors/frameworks/realtimeai.py +++ b/src/pipecat/processors/frameworks/realtimeai.py @@ -9,7 +9,7 @@ import dataclasses from typing import List, Literal, Optional, Type from pydantic import BaseModel, ValidationError -from pipecat.frames.frames import Frame, StartFrame, TransportMessageFrame +from pipecat.frames.frames import Frame, LLMMessagesFrame, LLMMessagesUpdateFrame, StartFrame, TransportMessageFrame from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.aggregators.llm_response import LLMAssistantResponseAggregator, LLMUserResponseAggregator from pipecat.processors.frame_processor import FrameDirection, FrameProcessor @@ -21,50 +21,61 @@ from pipecat.vad.silero import SileroVAD class RealtimeAILLMConfig(BaseModel): - model: str - messages: List[dict] + model: Optional[str] = None + messages: Optional[List[dict]] = None class RealtimeAITTSConfig(BaseModel): - voice: str + voice: Optional[str] = None class RealtimeAIConfig(BaseModel): - llm: RealtimeAILLMConfig - tts: RealtimeAITTSConfig + llm: Optional[RealtimeAILLMConfig] = None + tts: Optional[RealtimeAITTSConfig] = None + + +class RealtimeAISetup(BaseModel): + config: RealtimeAIConfig class RealtimeAIMessageData(BaseModel): + setup: Optional[RealtimeAISetup] = None config: Optional[RealtimeAIConfig] = None class RealtimeAIMessage(BaseModel): tag: Literal["realtime-ai"] = "realtime-ai" type: str - data: RealtimeAIMessageData + data: Optional[RealtimeAIMessageData] = None -class RealtimeAIResponseMessage(BaseModel): +class RealtimeAIBasicResponse(BaseModel): tag: Literal["realtime-ai"] = "realtime-ai" type: str success: bool error: Optional[str] = None +class RealtimeAILLMContextResponse(BaseModel): + tag: Literal["realtime-ai"] = "realtime-ai" + type: Literal["llm-context"] = "llm-context" + messages: List[dict] + + class RealtimeAIProcessor(FrameProcessor): def __init__( self, *, transport: BaseTransport, - config: RealtimeAIConfig | None = None, + setup: RealtimeAISetup | None = None, llm_api_key: str = "", tts_api_key: str = "", llm_cls: Type[AIService] = OLLamaLLMService, tts_cls: Type[AIService] = CartesiaTTSService): super().__init__() self._transport = transport - self._config = config + self._setup = setup self._llm_api_key = llm_api_key self._tts_api_key = tts_api_key self._llm_cls = llm_cls @@ -82,32 +93,48 @@ class RealtimeAIProcessor(FrameProcessor): if isinstance(frame, StartFrame): self._start_frame = frame - if self._config: - await self._handle_config(self._config) + if self._setup and self._setup.config: + await self._handle_setup(self._setup) async def _handle_message(self, frame: TransportMessageFrame): try: message = RealtimeAIMessage.model_validate(frame.message) - - match message.type: - case "config": - await self._handle_config(RealtimeAIConfig.model_validate(message.data.config)) except ValidationError as e: - await self._send_response("config", False, f"invalid configuration: {e}") + await self._send_response("setup", False, f"invalid message: {e}") + return + + print(message) - async def _handle_config(self, config: RealtimeAIConfig): try: - tma_in = LLMUserResponseAggregator(config.llm.messages) - tma_out = LLMAssistantResponseAggregator(config.llm.messages) + match message.type: + case "setup": + await self._handle_setup(RealtimeAISetup.model_validate(message.data.setup)) + case "llm-get-context": + await self._handle_llm_get_context() + case "llm-update-context": + await self._handle_llm_update_context(RealtimeAIConfig.model_validate(message.data.config)) + except ValidationError as e: + await self._send_response(message.type, False, f"invalid message: {e}") + async def _handle_setup(self, setup: RealtimeAISetup): + try: vad = SileroVAD() - self._llm = self._llm_cls(model=config.llm.model) + self._tma_in = LLMUserResponseAggregator(setup.config.llm.messages) + self._tma_out = LLMAssistantResponseAggregator(setup.config.llm.messages) - self._tts = self._tts_cls(api_key=self._tts_api_key, voice_id=config.tts.voice) + self._llm = self._llm_cls(model=setup.config.llm.model) - pipeline = Pipeline([vad, tma_in, self._llm, self._tts, - self._transport.output(), tma_out]) + self._tts = self._tts_cls(api_key=self._tts_api_key, voice_id=setup.config.tts.voice) + + pipeline = Pipeline([ + vad, + self._tma_in, + self._llm, + self._tts, + self._transport.output(), + self._tma_out + ]) self._pipeline = pipeline parent = self.get_parent() @@ -121,11 +148,37 @@ class RealtimeAIProcessor(FrameProcessor): # We now send a message to indicate we successfully initialized # the pipelines. - await self._send_response("config", True) + await self._send_response("setup", True) except Exception as e: - await self._send_response("config", False, f"unable to create pipeline: {e}") + await self._send_response("setup", False, f"unable to create pipeline: {e}") - async def _send_response(self, type: str, success: bool, error: str | None = None): - response = RealtimeAIResponseMessage(type=type, success=success) + async def _handle_llm_get_context(self): + messages = self._tma_in.messages + response = RealtimeAILLMContextResponse(messages=messages) + message = TransportMessageFrame(message=response.model_dump(exclude_none=True)) + await self.push_frame(message) + + async def _handle_llm_update_context(self, config: RealtimeAIConfig): + if config.llm and config.llm.messages: + frame = LLMMessagesUpdateFrame(config.llm.messages) + await self.push_frame(frame) + + async def _send_response(self, type: str, success: bool, error: str | None = None): + # TODO(aleix): This is a bit hacky, but we might get invalid + # configuration or something might going wrong during setup and we would + # like to send the error to the client. However, if the pipeline is not + # setup yet we don't have an output transport and therefore we can't + # send any messages. So, we setup a super basic pipeline with just the + # output transport so we can send messages. + if not self._pipeline: + # We add the SilerVAD() so the audio doesn't go through. + pipeline = Pipeline([SileroVAD(), self._transport.output()]) + self._pipeline = pipeline + + parent = self.get_parent() + if parent and self._start_frame: + parent.link(pipeline) + + response = RealtimeAIBasicResponse(type=type, success=success, error=error) message = TransportMessageFrame(message=response.model_dump(exclude_none=True)) await self.push_frame(message)