Merge pull request #500 from pipecat-ai/aleix/rtvi-action-frames-task

RTVI action frames task
This commit is contained in:
Aleix Conchillo Flaqué
2024-09-24 10:13:43 -07:00
committed by GitHub

View File

@@ -242,6 +242,22 @@ class RTVILLMFunctionCallResultData(BaseModel):
result: dict | str
class RTVITextMessageData(BaseModel):
text: str
class RTVILLMTextMessage(BaseModel):
label: Literal["rtvi-ai"] = "rtvi-ai"
type: Literal["llm-text"] = "llm-text"
data: RTVITextMessageData
class RTVITTSTextMessage(BaseModel):
label: Literal["rtvi-ai"] = "rtvi-ai"
type: Literal["tts-text"] = "tts-text"
data: RTVITextMessageData
class RTVITranscriptionMessageData(BaseModel):
text: str
user_id: str
@@ -300,6 +316,11 @@ class RTVIProcessor(FrameProcessor):
self._registered_actions: Dict[str, RTVIAction] = {}
self._registered_services: Dict[str, RTVIService] = {}
# A task to process incoming action frames.
self._action_task = self.get_event_loop().create_task(self._action_task_handler())
self._action_queue = asyncio.Queue()
# A task to process incoming transport messages.
self._message_task = self.get_event_loop().create_task(self._message_task_handler())
self._message_queue = asyncio.Queue()
@@ -385,7 +406,7 @@ class RTVIProcessor(FrameProcessor):
elif isinstance(frame, TransportMessageFrame):
await self._message_queue.put(frame)
elif isinstance(frame, RTVIActionFrame):
await self._handle_action(frame.message_id, frame.rtvi_action_run)
await self._action_queue.put(frame)
# Other frames
else:
await self.push_frame(frame, direction)
@@ -399,12 +420,16 @@ class RTVIProcessor(FrameProcessor):
await self._maybe_send_bot_ready()
async def _stop(self, frame: EndFrame):
# We need to cancel the message task handler because that one is not
# processing EndFrames.
self._action_task.cancel()
await self._action_task
self._message_task.cancel()
await self._message_task
async def _cancel(self, frame: CancelFrame):
self._action_task.cancel()
await self._action_task
self._message_task.cancel()
await self._message_task
@@ -455,6 +480,15 @@ class RTVIProcessor(FrameProcessor):
if message:
await self._push_transport_message(message)
async def _action_task_handler(self):
while True:
try:
frame = await self._action_queue.get()
await self._handle_action(frame.message_id, frame.rtvi_action_run)
self._action_queue.task_done()
except asyncio.CancelledError:
break
async def _message_task_handler(self):
while True:
try: