wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

services: run_* now return async generators

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2024-05-14 00:30:07 -07:00
parent e018d5b47a
commit 922cdefee5
9 changed files with 68 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

110
src/pipecat/processors/aggregators/openai_llm_context.py
View File

@@ -6,23 +6,13 @@
from dataclasses import dataclass
from typing import AsyncGenerator, Callable, List
from typing import List
from pipecat.frames.frames import (
Frame,
LLMResponseEndFrame,
LLMResponseStartFrame,
TextFrame,
TranscriptionFrame,
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame,
)
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.frames.frames import Frame
from openai._types import NOT_GIVEN, NotGiven
from openai.types.chat import (
ChatCompletionRole,
ChatCompletionToolParam,
ChatCompletionToolChoiceOptionParam,
ChatCompletionMessageParam
@@ -42,7 +32,7 @@ class OpenAILLMContext:
self.tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = tool_choice
self.tools: List[ChatCompletionToolParam] | NotGiven = tools
@ staticmethod
@staticmethod
def from_messages(messages: List[dict]) -> "OpenAILLMContext":
context = OpenAILLMContext()
for message in messages:
@@ -71,100 +61,6 @@ class OpenAILLMContext:
self.tools = tools
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,
)
@dataclass
class OpenAILLMContextFrame(Frame):
"""Like an LLMMessagesFrame, but with extra context specific to the OpenAI

28
src/pipecat/services/ai_services.py
View File

@@ -10,11 +10,12 @@ import math
import wave
from abc import abstractmethod
from typing import BinaryIO
from typing import AsyncGenerator, BinaryIO
from pipecat.frames.frames import (
AudioRawFrame,
EndFrame,
ErrorFrame,
Frame,
TextFrame,
VisionImageRawFrame,
@@ -26,6 +27,13 @@ class AIService(FrameProcessor):
def __init__(self):
super().__init__()
async def process_generator(self, generator: AsyncGenerator[Frame, None]):
async for f in generator:
if isinstance(f, ErrorFrame):
await self.push_error(f)
else:
await self.push_frame(f)
class LLMService(AIService):
"""This class is a no-op but serves as a base class for LLM services."""
@@ -42,7 +50,7 @@ class TTSService(AIService):
# Converts the text to audio.
@abstractmethod
async def run_tts(self, text: str):
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
pass
async def say(self, text: str):
@@ -59,14 +67,14 @@ class TTSService(AIService):
self._current_sentence = ""
if text:
await self.run_tts(text)
await self.process_generator(self.run_tts(text))
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, TextFrame):
await self._process_text_frame(frame)
elif isinstance(frame, EndFrame):
if self._current_sentence:
await self.run_tts(self._current_sentence)
await self.process_generator(self.run_tts(self._current_sentence))
await self.push_frame(frame)
else:
await self.push_frame(frame, direction)
@@ -89,7 +97,7 @@ class STTService(AIService):
(self._content, self._wave) = self._new_wave()
@abstractmethod
async def run_stt(self, audio: BinaryIO):
async def run_stt(self, audio: BinaryIO) -> AsyncGenerator[Frame, None]:
"""Returns transcript as a string"""
pass
@@ -130,7 +138,7 @@ class STTService(AIService):
self._current_silence_frames = 0
self._wave.close()
self._content.seek(0)
await self.run_stt(self._content)
await self.process_generator(self.run_stt(self._content))
(self._content, self._wave) = self._new_wave()
# If we get this far, this is a frame of silence
self._current_silence_frames += 1
@@ -143,12 +151,12 @@ class ImageGenService(AIService):
# Renders the image. Returns an Image object.
@abstractmethod
async def run_image_gen(self, prompt: str):
async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
pass
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, TextFrame):
await self.run_image_gen(frame.text)
await self.process_generator(self.run_image_gen(frame.text))
else:
await self.push_frame(frame, direction)
@@ -161,11 +169,11 @@ class VisionService(AIService):
self._describe_text = None
@abstractmethod
async def run_vision(self, frame: VisionImageRawFrame):
async def run_vision(self, frame: VisionImageRawFrame) -> AsyncGenerator[Frame, None]:
pass
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, VisionImageRawFrame):
await self.run_vision(frame)
await self.process_generator(self.run_vision(frame))
else:
await self.push_frame(frame, direction)

15
src/pipecat/services/azure.py
View File

@@ -9,10 +9,11 @@ import asyncio
import io
from PIL import Image
from typing import AsyncGenerator
from openai import AsyncAzureOpenAI
from pipecat.frames.frames import AudioRawFrame, ErrorFrame, URLImageRawFrame
from pipecat.frames.frames import AudioRawFrame, ErrorFrame, Frame, URLImageRawFrame
from pipecat.services.ai_services import TTSService, ImageGenService
from pipecat.services.openai import BaseOpenAILLMService
@@ -43,7 +44,7 @@ class AzureTTSService(TTSService):
)
self._voice = voice
async def run_tts(self, text: str):
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Transcribing text: {text}")
ssml = (
@@ -60,7 +61,7 @@ class AzureTTSService(TTSService):
if result.reason == ResultReason.SynthesizingAudioCompleted:
# Azure always sends a 44-byte header. Strip it off.
await self.push_frame(AudioRawFrame(audio=result.audio_data[44:], sample_rate=16000, num_channels=1))
yield AudioRawFrame(audio=result.audio_data[44:], sample_rate=16000, num_channels=1)
elif result.reason == ResultReason.Canceled:
cancellation_details = result.cancellation_details
logger.warning(f"Speech synthesis canceled: {cancellation_details.reason}")
@@ -110,7 +111,7 @@ class AzureImageGenServiceREST(ImageGenService):
self._aiohttp_session = aiohttp_session
self._image_size = image_size
async def run_image_gen(self, prompt: str):
async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
url = f"{self._azure_endpoint}openai/images/generations:submit?api-version={self._api_version}"
headers = {
@@ -136,7 +137,7 @@ class AzureImageGenServiceREST(ImageGenService):
attempts_left -= 1
if attempts_left == 0:
logger.error("Image generation timed out")
await self.push_error(ErrorFrame("Image generation timed out"))
yield ErrorFrame("Image generation timed out")
return
await asyncio.sleep(1)
@@ -149,7 +150,7 @@ class AzureImageGenServiceREST(ImageGenService):
image_url = json_response["result"]["data"][0]["url"] if json_response else None
if not image_url:
logger.error("Image generation failed")
await self.push_error(ErrorFrame("Image generation failed"))
yield ErrorFrame("Image generation failed")
return
# Load the image from the url
@@ -161,4 +162,4 @@ class AzureImageGenServiceREST(ImageGenService):
image=image.tobytes(),
size=image.size,
format=image.format)
await self.push_frame(frame)
yield frame

8
src/pipecat/services/deepgram.py
View File

@@ -4,7 +4,9 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
from pipecat.frames.frames import AudioRawFrame
from typing import AsyncGenerator
from pipecat.frames.frames import AudioRawFrame, Frame
from pipecat.services.ai_services import TTSService
from loguru import logger
@@ -24,7 +26,7 @@ class DeepgramTTSService(TTSService):
self._api_key = api_key
self._aiohttp_session = aiohttp_session
async def run_tts(self, text: str):
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
logger.info(f"Running Deepgram TTS for {text}")
base_url = "https://api.beta.deepgram.com/v1/speak"
request_url = f"{base_url}?model={self._voice}&encoding=linear16&container=none&sample_rate=16000"
@@ -33,4 +35,4 @@ class DeepgramTTSService(TTSService):
async with self._aiohttp_session.post(request_url, headers=headers, json=body) as r:
async for data in r.content:
frame = AudioRawFrame(audio=data, sample_rate=16000, num_channels=1)
await self.push_frame(frame)
yield frame

13
src/pipecat/services/elevenlabs.py
View File

@@ -6,7 +6,9 @@
import aiohttp
from pipecat.frames.frames import AudioRawFrame, TTSStartedFrame, TTSStoppedFrame
from typing import AsyncGenerator
from pipecat.frames.frames import AudioRawFrame, ErrorFrame, Frame, TTSStartedFrame, TTSStoppedFrame
from pipecat.services.ai_services import TTSService
from loguru import logger
@@ -29,7 +31,7 @@ class ElevenLabsTTSService(TTSService):
self._aiohttp_session = aiohttp_session
self._model = model
async def run_tts(self, text: str):
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Transcribing text: {text}")
url = f"https://api.elevenlabs.io/v1/text-to-speech/{self._voice_id}/stream"
@@ -48,11 +50,12 @@ class ElevenLabsTTSService(TTSService):
async with self._aiohttp_session.post(url, json=payload, headers=headers, params=querystring) as r:
if r.status != 200:
logger.error(f"Audio fetch status code: {r.status}, error: {r.text}")
yield ErrorFrame(f"Audio fetch status code: {r.status}, error: {r.text}")
return
await self.push_frame(TTSStartedFrame())
yield TTSStartedFrame()
async for chunk in r.content:
if len(chunk) > 0:
frame = AudioRawFrame(chunk, 16000, 1)
await self.push_frame(frame)
await self.push_frame(TTSStoppedFrame())
yield frame
yield TTSStoppedFrame()

10
src/pipecat/services/fal.py
View File

@@ -9,11 +9,10 @@ import io
import os
from PIL import Image
from numpy import result_type
from pydantic import BaseModel
from typing import Optional, Union, Dict
from typing import AsyncGenerator, Optional, Union, Dict
from pipecat.frames.frames import URLImageRawFrame
from pipecat.frames.frames import ErrorFrame, Frame, URLImageRawFrame
from pipecat.services.ai_services import ImageGenService
from loguru import logger
@@ -52,7 +51,7 @@ class FalImageGenService(ImageGenService):
if key:
os.environ["FAL_KEY"] = key
async def run_image_gen(self, prompt: str):
async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Generating image from prompt: {prompt}")
response = await fal_client.run_async(
@@ -64,6 +63,7 @@ class FalImageGenService(ImageGenService):
if not image_url:
logger.error("Image generation failed")
yield ErrorFrame("Image generation failed")
return
logger.debug(f"Image generated at: {image_url}")
@@ -80,4 +80,4 @@ class FalImageGenService(ImageGenService):
image=image.tobytes(),
size=image.size,
format=image.format)
await self.push_frame(frame)
yield frame

13
src/pipecat/services/moondream.py
View File

@@ -6,11 +6,13 @@
import asyncio
from pipecat.frames.frames import TextFrame, VisionImageRawFrame
from pipecat.services.ai_services import VisionService
from PIL import Image
from typing import AsyncGenerator
from pipecat.frames.frames import ErrorFrame, Frame, TextFrame, VisionImageRawFrame
from pipecat.services.ai_services import VisionService
from loguru import logger
try:
@@ -67,9 +69,10 @@ class MoondreamService(VisionService):
logger.debug("Loaded Moondream model")
async def run_vision(self, frame: VisionImageRawFrame):
async def run_vision(self, frame: VisionImageRawFrame) -> AsyncGenerator[Frame, None]:
if not self._model:
logger.error("Moondream model not available")
yield ErrorFrame("Moondream model not available")
return
logger.debug(f"Analyzing image: {frame}")
@@ -85,4 +88,4 @@ class MoondreamService(VisionService):
description = await asyncio.to_thread(get_image_description, frame)
await self.push_frame(TextFrame(text=description))
yield TextFrame(text=description)

12
src/pipecat/services/openai.py
View File

@@ -8,11 +8,13 @@ import io
import json
import time
import aiohttp
from PIL import Image
from typing import List, Literal
from typing import AsyncGenerator, List, Literal
from pipecat.frames.frames import (
ErrorFrame,
Frame,
LLMMessagesFrame,
LLMResponseEndFrame,
@@ -174,7 +176,7 @@ class OpenAIImageGenService(ImageGenService):
self._client = AsyncOpenAI(api_key=api_key)
self._aiohttp_session = aiohttp_session
async def run_image_gen(self, prompt: str):
async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Generating image from prompt: {prompt}")
image = await self._client.images.generate(
@@ -187,11 +189,13 @@ class OpenAIImageGenService(ImageGenService):
image_url = image.data[0].url
if not image_url:
logger.error(f"no image provided in response: {image}")
logger.error(f"No image provided in response: {image}")
yield ErrorFrame("Image generation failed")
return
# Load the image from the url
async with self._aiohttp_session.get(image_url) as response:
image_stream = io.BytesIO(await response.content.read())
image = Image.open(image_stream)
frame = URLImageRawFrame(image_url, image.tobytes(), image.size, image.format)
await self.push_frame(frame)
yield frame

8
src/pipecat/services/playht.py
View File

@@ -7,7 +7,9 @@
import io
import struct
from pipecat.frames.frames import AudioRawFrame
from typing import AsyncGenerator
from pipecat.frames.frames import AudioRawFrame, Frame
from pipecat.services.ai_services import TTSService
from loguru import logger
@@ -44,7 +46,7 @@ class PlayHTAIService(TTSService):
def __del__(self):
self._client.close()
async def run_tts(self, text: str):
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
b = bytearray()
in_header = True
for chunk in self._client.tts(text, self._options):
@@ -69,4 +71,4 @@ class PlayHTAIService(TTSService):
else:
if len(chunk):
frame = AudioRawFrame(chunk, 16000, 1)
await self.push_frame(frame)
yield frame