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

refactor party tonight

Browse Source
This commit is contained in:
Moishe Lettvin
2024-01-17 18:42:08 -05:00
parent a3ac0d84e8
commit 13f2f792af
10 changed files with 187 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
src/dailyai/services/ai_services.py
View File

@@ -2,6 +2,8 @@ import asyncio
import logging
import re
from httpx import request
from dailyai.queue_frame import QueueFrame, FrameType
from abc import abstractmethod
@@ -13,9 +15,7 @@ from collections.abc import Iterable, AsyncIterable
class AIService:
def __init__(
self
):
def __init__(self):
self.logger = logging.getLogger("dailyai")
def stop(self):
@@ -27,30 +27,60 @@ class AIService:
def possible_output_frame_types(self) -> set[FrameType]:
return set()
async def run_to_queue(self, queue: asyncio.Queue, frames, add_end_of_stream=False) -> None:
async for frame in self.run(frames):
print("got frame", frame.frame_type)
await queue.put(frame)
if add_end_of_stream:
await queue.put(QueueFrame(FrameType.END_STREAM, None))
async def run(
self,
requested_frame_types:set[FrameType],
frames:Iterable[QueueFrame] | AsyncIterable[QueueFrame]
) -> AsyncGenerator[QueueFrame, None]:
if self.possible_output_frame_types().intersection(requested_frame_types) == set():
self,
frames: Iterable[QueueFrame]
| AsyncIterable[QueueFrame]
| asyncio.Queue[QueueFrame],
requested_frame_types: set[FrameType] | None=None,
) -> AsyncGenerator[QueueFrame, None]:
if requested_frame_types and self.possible_output_frame_types().intersection(requested_frame_types) == set():
raise Exception(f"Requested frame types {requested_frame_types} are not supported by this service.")
if not requested_frame_types:
requested_frame_types = self.possible_output_frame_types()
print("running", self.__class__.__name__, "with frame types", requested_frame_types)
if isinstance(frames, AsyncIterable):
async for frame in frames:
output_frame: QueueFrame | None = await self.process_frame(requested_frame_types, frame)
if output_frame:
async for output_frame in self.process_frame(requested_frame_types, frame):
print(
"yielding frame", self.__class__.__name__, output_frame.frame_type
)
yield output_frame
elif isinstance(frames, Iterable):
for frame in frames:
output_frame = await self.process_frame(requested_frame_types, frame)
if output_frame:
async for output_frame in self.process_frame(requested_frame_types, frame):
print(
"yielding frame", self.__class__.__name__, output_frame.frame_type
)
yield output_frame
elif isinstance(frames, asyncio.Queue):
while True:
frame = await frames.get()
async for output_frame in self.process_frame(requested_frame_types, frame):
print(
"yielding frame", self.__class__.__name__, output_frame.frame_type
)
yield output_frame
if frame.frame_type == FrameType.END_STREAM:
break
else:
raise Exception("Frames must be an iterable or async iterable")
@abstractmethod
async def process_frame(self, requested_frame_types:set[FrameType], frame:QueueFrame) -> QueueFrame | None:
pass
async def process_frame(self, requested_frame_types:set[FrameType], frame:QueueFrame) -> AsyncGenerator[QueueFrame, None]:
# Yield something so the linter can deduce what should happen here.
yield QueueFrame(FrameType.END_STREAM, None)
class SentenceAggregator(AIService):
def __init__(self, **kwargs):
@@ -63,29 +93,26 @@ class SentenceAggregator(AIService):
def possible_output_frame_types(self) -> set[FrameType]:
return set([FrameType.SENTENCE])
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> QueueFrame | None:
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not FrameType.SENTENCE in requested_frame_types:
return None
return
if frame.frame_type == FrameType.TEXT_CHUNK:
if type(frame.frame_data) != str:
raise Exception("Sentence aggregator requires a string for the data field")
raise Exception(
"Sentence aggregator requires a string for the data field"
)
self.current_sentence += frame.frame_data
if self.current_sentence.endswith((".", "?", "!")):
sentence = self.current_sentence
self.current_sentence = ""
return QueueFrame(FrameType.SENTENCE, sentence)
return None
yield QueueFrame(FrameType.SENTENCE, sentence)
elif frame.frame_type == FrameType.END_STREAM:
if self.current_sentence:
return QueueFrame(FrameType.SENTENCE, self.current_sentence)
else:
return None
yield QueueFrame(FrameType.SENTENCE, self.current_sentence)
elif frame.frame_type == FrameType.SENTENCE:
return frame
else:
return None
yield frame
class LLMService(AIService):
@@ -93,30 +120,29 @@ class LLMService(AIService):
return set([FrameType.LLM_MESSAGE, FrameType.SENTENCE, FrameType.TRANSCRIPTION])
def allowed_output_frame_types(self) -> set[FrameType]:
return set([FrameType.SENTENCE, FrameType.SENTENCE, FrameType.TEXT_CHUNK])
return set([FrameType.SENTENCE, FrameType.TEXT_CHUNK])
async def run_llm_async_sentences(self, messages) -> AsyncGenerator[str, None]:
current_text = ""
async for text in self.run_llm_async(messages):
current_text += text
if re.match(r"^.*[.!?]$", text):
yield current_text
current_text = ""
@abstractmethod
async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
yield ""
if current_text:
yield current_text
async def process_frame(self, frame:QueueFrame) -> QueueFrame | None:
if not self.output_queue:
raise Exception("Output queue must be set before using the run method.")
@abstractmethod
async def run_llm(self, messages) -> str:
pass
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if frame.frame_type == FrameType.LLM_MESSAGE:
if type(frame.frame_data) != list:
raise Exception("LLM service requires a dict for the data field")
messages: list[dict[str, str]] = frame.frame_data
async for message in self.run_llm_async_sentences(messages):
await self.output_queue.put(QueueFrame(FrameType.SENTENCE, message))
if FrameType.SENTENCE in requested_frame_types:
yield QueueFrame(FrameType.SENTENCE, await self.run_llm(messages))
else:
async for text_chunk in self.run_llm_async(messages):
yield QueueFrame(FrameType.TEXT_CHUNK, text_chunk)
# TODO: handle other frame types! Need to aggregate into messages
class TTSService(AIService):
@@ -124,6 +150,12 @@ class TTSService(AIService):
def get_mic_sample_rate(self):
return 16000
def allowed_input_frame_types(self) -> set[FrameType]:
return set([FrameType.SENTENCE, FrameType.TRANSCRIPTION, FrameType.TEXT_CHUNK])
def possible_output_frame_types(self) -> set[FrameType]:
return set([FrameType.AUDIO])
# Converts the sentence to audio. Yields a list of audio frames that can
# be sent to the microphone device
@abstractmethod
@@ -131,25 +163,48 @@ class TTSService(AIService):
# yield empty bytes here, so linting can infer what this method does
yield bytes()
async def process_frame(self, frame:QueueFrame) -> QueueFrame | None:
if not self.output_queue:
raise Exception("Output queue must be set before using the run method.")
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not FrameType.AUDIO in requested_frame_types:
return
if frame.frame_type == FrameType.SENTENCE:
if type(frame.frame_data) != str:
raise Exception("TTS service requires a string for the data field")
if type(frame.frame_data) != str:
raise Exception("TTS service requires a string for the data field")
text = frame.frame_data
async for audio in self.run_tts(text):
await self.output_queue.put(QueueFrame(FrameType.AUDIO, audio))
async for audio_chunk in self.run_tts(frame.frame_data):
yield QueueFrame(FrameType.AUDIO, audio_chunk)
# Convenience function to send the audio for a sentence to the given queue
async def say(self, sentence, queue: asyncio.Queue):
async for audio_chunk in self.run_tts(sentence):
await queue.put(QueueFrame(FrameType.AUDIO, audio_chunk))
class ImageGenService(AIService):
def __init__(self, image_size, **kwargs):
super().__init__(**kwargs)
self.image_size = image_size
def allowed_input_frame_types(self) -> set[FrameType]:
return set([FrameType.SENTENCE, FrameType.TRANSCRIPTION, FrameType.TEXT_CHUNK, FrameType.IMAGE_DESCRIPTION])
def possible_output_frame_types(self) -> set[FrameType]:
return set([FrameType.IMAGE])
# Renders the image. Returns an Image object.
@abstractmethod
async def run_image_gen(self, sentence, size) -> tuple[str, bytes]:
async def run_image_gen(self, sentence) -> tuple[str, bytes]:
pass
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not FrameType.IMAGE in requested_frame_types:
return
if type(frame.frame_data) != str:
raise Exception("Image service requires a string for the data field")
(_, image_data) = await self.run_image_gen(frame.frame_data)
yield QueueFrame(FrameType.IMAGE, image_data)
@dataclass
class AIServiceConfig: