examples: fix storytelling example

examples/storytelling-chatbot/src/processors.py

@@ -1,19 +1,13 @@
-from typing import AsyncGenerator
 import re
 
-from dailyai.pipeline.frames import TextFrame, Frame, AudioFrame
-from dailyai.pipeline.frame_processor import FrameProcessor
-from dailyai.pipeline.frames import (
-    Frame,
-    TextFrame,
-    SendAppMessageFrame,
-    LLMResponseEndFrame,
-    UserStoppedSpeakingFrame,
-)
+from async_timeout import timeout
+
+from pipecat.frames.frames import Frame, LLMResponseEndFrame, TextFrame, UserStoppedSpeakingFrame
+from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.transports.services.daily import DailyTransportMessageFrame
 
 from utils.helpers import load_sounds
 from prompts import IMAGE_GEN_PROMPT, CUE_USER_TURN, CUE_ASSISTANT_TURN
-import asyncio
 
 sounds = load_sounds(["talking.wav", "listening.wav", "ding.wav"])
 
@@ -42,7 +36,7 @@ class StoryImageProcessor(FrameProcessor):
     Processor for image prompt frames that will be sent to the FAL service.
 
     This processor is responsible for consuming frames of type `StoryImageFrame`.
-    It processes the by passing it to the FAL service
+    It processes them by passing it to the FAL service.
     The processed frames are then yielded back.
 
     Attributes:
@@ -50,25 +44,26 @@ class StoryImageProcessor(FrameProcessor):
     """
 
     def __init__(self, fal_service):
+        super().__init__()
         self._fal_service = fal_service
 
-    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
         if isinstance(frame, StoryImageFrame):
             try:
-                async with asyncio.timeout(7):
-                    async for i in self._fal_service.process_frame(TextFrame(IMAGE_GEN_PROMPT % frame.text)):
-                        yield i
+                async with timeout(7):
+                    async for i in self._fal_service.run_image_gen(IMAGE_GEN_PROMPT % frame.text):
+                        await self.push_frame(i)
             except TimeoutError:
                 pass
             pass
         else:
-            yield frame
+            await self.push_frame(frame)
 
 
 class StoryProcessor(FrameProcessor):
     """
     Primary frame processor. It takes the frames generated by the LLM
-    and processes them into image prompts and story pages (sentences.)
+    and processes them into image prompts and story pages (sentences).
     For a clearer picture of how this works, reference prompts.py
 
     Attributes:
@@ -81,15 +76,16 @@ class StoryProcessor(FrameProcessor):
     """
 
     def __init__(self, messages, story):
+        super().__init__()
         self._messages = messages
         self._text = ""
         self._story = story
 
-    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
         if isinstance(frame, UserStoppedSpeakingFrame):
             # Send an app message to the UI
-            yield SendAppMessageFrame(CUE_ASSISTANT_TURN, None)
-            yield AudioFrame(sounds["talking"])
+            await self.push_frame(DailyTransportMessageFrame(CUE_ASSISTANT_TURN))
+            await self.push_frame(sounds["talking"])
 
         elif isinstance(frame, TextFrame):
             # We want to look for sentence breaks in the text
@@ -111,7 +107,7 @@ class StoryProcessor(FrameProcessor):
                 # Remove the image prompt from the text
                 self._text = re.sub(r"<.*?>", '', self._text, count=1)
                 # Process the image prompt frame
-                yield StoryImageFrame(image_prompt)
+                await self.push_frame(StoryImageFrame(image_prompt))
 
             # STORY PAGE
             # Looking for: [break] in the LLM response
@@ -126,9 +122,9 @@ class StoryProcessor(FrameProcessor):
                 if len(self._text) > 2:
                     # Append the sentence to the story
                     self._story.append(self._text)
-                    yield StoryPageFrame(self._text)
+                    await self.push_frame(StoryPageFrame(self._text))
                     # Assert that it's the LLMs turn, until we're finished
-                    yield SendAppMessageFrame(CUE_ASSISTANT_TURN, None)
+                    await self.push_frame(DailyTransportMessageFrame(CUE_ASSISTANT_TURN))
                 # Clear the buffer
                 self._text = ""
 
@@ -136,13 +132,13 @@ class StoryProcessor(FrameProcessor):
         # Driven by the prompt, the LLM should have asked the user for input
         elif isinstance(frame, LLMResponseEndFrame):
             # We use a different frame type, as to avoid image generation ingest
-            yield StoryPromptFrame(self._text)
+            await self.push_frame(StoryPromptFrame(self._text))
             self._text = ""
-            yield frame
+            await self.push_frame(frame)
             # Send an app message to the UI
-            yield SendAppMessageFrame(CUE_USER_TURN, None)
-            yield AudioFrame(sounds["listening"])
+            await self.push_frame(DailyTransportMessageFrame(CUE_USER_TURN))
+            await self.push_frame(sounds["listening"])
 
         # Anything that is not a TextFrame pass through
         else:
-            yield frame
+            await self.push_frame(frame)