214 lines
8.1 KiB
Python
214 lines
8.1 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
|
|
from httpx import request
|
|
|
|
from dailyai.queue_frame import QueueFrame, FrameType
|
|
|
|
from abc import abstractmethod
|
|
from typing import AsyncGenerator, Iterable
|
|
from dataclasses import dataclass
|
|
from typing import AsyncGenerator
|
|
|
|
from collections.abc import Iterable, AsyncIterable
|
|
|
|
class AIService:
|
|
|
|
def __init__(self):
|
|
self.logger = logging.getLogger("dailyai")
|
|
|
|
def stop(self):
|
|
pass
|
|
|
|
def allowed_input_frame_types(self) -> set[FrameType]:
|
|
return set()
|
|
|
|
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,
|
|
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:
|
|
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:
|
|
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) -> 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):
|
|
super().__init__(**kwargs)
|
|
self.current_sentence = ""
|
|
|
|
def allowed_input_frame_types(self) -> set[FrameType]:
|
|
return set([FrameType.TEXT_CHUNK, FrameType.SENTENCE])
|
|
|
|
def possible_output_frame_types(self) -> set[FrameType]:
|
|
return set([FrameType.SENTENCE])
|
|
|
|
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
|
|
if not FrameType.SENTENCE in requested_frame_types:
|
|
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"
|
|
)
|
|
|
|
self.current_sentence += frame.frame_data
|
|
if self.current_sentence.endswith((".", "?", "!")):
|
|
sentence = self.current_sentence
|
|
self.current_sentence = ""
|
|
yield QueueFrame(FrameType.SENTENCE, sentence)
|
|
elif frame.frame_type == FrameType.END_STREAM:
|
|
if self.current_sentence:
|
|
yield QueueFrame(FrameType.SENTENCE, self.current_sentence)
|
|
elif frame.frame_type == FrameType.SENTENCE:
|
|
yield frame
|
|
|
|
|
|
class LLMService(AIService):
|
|
def allowed_input_frame_types(self) -> set[FrameType]:
|
|
return set([FrameType.LLM_MESSAGE, FrameType.SENTENCE, FrameType.TRANSCRIPTION])
|
|
|
|
def allowed_output_frame_types(self) -> set[FrameType]:
|
|
return set([FrameType.SENTENCE, FrameType.TEXT_CHUNK])
|
|
|
|
@abstractmethod
|
|
async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
|
|
yield ""
|
|
|
|
@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
|
|
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):
|
|
# Some TTS services require a specific sample rate. We default to 16k
|
|
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
|
|
async def run_tts(self, sentence) -> AsyncGenerator[bytes, None]:
|
|
# yield empty bytes here, so linting can infer what this method does
|
|
yield bytes()
|
|
|
|
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 type(frame.frame_data) != str:
|
|
raise Exception("TTS service requires a string for the data field")
|
|
|
|
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) -> 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:
|
|
tts: TTSService
|
|
image: ImageGenService
|
|
llm: LLMService
|