working on making services more consistent/terse/easy
This commit is contained in:
Moishe Lettvin
@@ -5,58 +5,95 @@ import re
|
||||
from dailyai.queue_frame import QueueFrame, FrameType
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import AsyncGenerator
|
||||
from typing import AsyncGenerator, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from collections.abc import Iterable, AsyncIterable
|
||||
|
||||
class AIService:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_queue: asyncio.Queue[QueueFrame] | None = None,
|
||||
output_queue: asyncio.Queue[QueueFrame] | None = None,
|
||||
self
|
||||
):
|
||||
self.logger = logging.getLogger("dailyai")
|
||||
self.input_queue: asyncio.Queue[QueueFrame] | None = input_queue
|
||||
self.output_queue: asyncio.Queue[QueueFrame] | None = output_queue
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
async def run(self) -> None:
|
||||
if self.input_queue is None or self.output_queue is None:
|
||||
raise Exception("Input and output queues must be set before using the run method.")
|
||||
def allowed_input_frame_types(self) -> set[FrameType]:
|
||||
return set()
|
||||
|
||||
while True:
|
||||
frame = await self.input_queue.get()
|
||||
self.logger.debug(f"{self.__class__.__name__} got frame:", frame.frame_type)
|
||||
if frame.frame_type == FrameType.END_STREAM:
|
||||
self.input_queue.task_done()
|
||||
await self.output_queue.put(QueueFrame(FrameType.END_STREAM, None))
|
||||
break
|
||||
def possible_output_frame_types(self) -> set[FrameType]:
|
||||
return set()
|
||||
|
||||
output_frame = await self.process_frame(frame)
|
||||
if output_frame:
|
||||
await self.output_queue.put(output_frame)
|
||||
self.input_queue.task_done()
|
||||
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():
|
||||
raise Exception(f"Requested frame types {requested_frame_types} are not supported by this service.")
|
||||
|
||||
if isinstance(frames, AsyncIterable):
|
||||
async for frame in frames:
|
||||
output_frame: QueueFrame | None = await self.process_frame(requested_frame_types, frame)
|
||||
if output_frame:
|
||||
yield output_frame
|
||||
elif isinstance(frames, Iterable):
|
||||
for frame in frames:
|
||||
output_frame = await self.process_frame(requested_frame_types, frame)
|
||||
if output_frame:
|
||||
yield output_frame
|
||||
else:
|
||||
raise Exception("Frames must be an iterable or async iterable")
|
||||
|
||||
@abstractmethod
|
||||
async def process_frame(self, frame) -> QueueFrame | None:
|
||||
async def process_frame(self, requested_frame_types:set[FrameType], frame:QueueFrame) -> QueueFrame | None:
|
||||
pass
|
||||
|
||||
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) -> QueueFrame | None:
|
||||
if not FrameType.SENTENCE in requested_frame_types:
|
||||
return None
|
||||
|
||||
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 = ""
|
||||
return QueueFrame(FrameType.SENTENCE, sentence)
|
||||
return None
|
||||
elif frame.frame_type == FrameType.END_STREAM:
|
||||
if self.current_sentence:
|
||||
return QueueFrame(FrameType.SENTENCE, self.current_sentence)
|
||||
else:
|
||||
return None
|
||||
elif frame.frame_type == FrameType.SENTENCE:
|
||||
return frame
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
class LLMService(AIService):
|
||||
# Generate a set of responses to a prompt. Yields a list of responses.
|
||||
@abstractmethod
|
||||
async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
|
||||
# Adding a yield here lets the linter know what this method actually does
|
||||
yield ""
|
||||
def allowed_input_frame_types(self) -> set[FrameType]:
|
||||
return set([FrameType.LLM_MESSAGE, FrameType.SENTENCE, FrameType.TRANSCRIPTION])
|
||||
|
||||
# Generate a responses to a prompt. Returns the response
|
||||
@abstractmethod
|
||||
async def run_llm(
|
||||
self, messages
|
||||
) -> str or None:
|
||||
pass
|
||||
def allowed_output_frame_types(self) -> set[FrameType]:
|
||||
return set([FrameType.SENTENCE, FrameType.SENTENCE, FrameType.TEXT_CHUNK])
|
||||
|
||||
async def run_llm_async_sentences(self, messages) -> AsyncGenerator[str, None]:
|
||||
current_text = ""
|
||||
|
||||
129
src/dailyai/tests/test_ai_services.py
Normal file
129
src/dailyai/tests/test_ai_services.py
Normal file
@@ -0,0 +1,129 @@
|
||||
from re import A
|
||||
import unittest
|
||||
|
||||
from typing import AsyncGenerator, Generator
|
||||
|
||||
from dailyai.services.ai_services import AIService, SentenceAggregator
|
||||
from dailyai.queue_frame import QueueFrame, FrameType
|
||||
|
||||
class SimpleAIService(AIService):
|
||||
def allowed_input_frame_types(self) -> set[FrameType]:
|
||||
return set([FrameType.TEXT_CHUNK])
|
||||
|
||||
def possible_output_frame_types(self) -> set[FrameType]:
|
||||
return set([FrameType.TEXT_CHUNK])
|
||||
|
||||
async def process_frame(self, requested_frame_types: set[FrameType], frame: QueueFrame) -> QueueFrame | None:
|
||||
return frame
|
||||
|
||||
class TestBaseAIService(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_async_input(self):
|
||||
service = SimpleAIService()
|
||||
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello"),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
async def iterate_frames() -> AsyncGenerator[QueueFrame, None]:
|
||||
for frame in input_frames:
|
||||
yield frame
|
||||
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.TEXT_CHUNK]), iterate_frames()):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(input_frames, output_frames)
|
||||
|
||||
async def test_nonasync_input(self):
|
||||
service = SimpleAIService()
|
||||
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello"),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
|
||||
def iterate_frames() -> Generator[QueueFrame, None, None]:
|
||||
for frame in input_frames:
|
||||
yield frame
|
||||
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.TEXT_CHUNK]), iterate_frames()):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(input_frames, output_frames)
|
||||
|
||||
|
||||
class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_clause(self) -> None:
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello"),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
|
||||
service = SentenceAggregator()
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.SENTENCE]), input_frames):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(1, len(output_frames))
|
||||
self.assertEqual(QueueFrame(FrameType.SENTENCE, "hello"), output_frames[0])
|
||||
|
||||
async def test_sentence(self) -> None:
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello, "),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "world."),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
|
||||
service = SentenceAggregator()
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.SENTENCE]), input_frames):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(1, len(output_frames))
|
||||
self.assertEqual(QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0])
|
||||
|
||||
async def test_sentence_and_clause(self) -> None:
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello, "),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "world."),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, " How are"),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
|
||||
service = SentenceAggregator()
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.SENTENCE]), input_frames):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(2, len(output_frames))
|
||||
self.assertEqual(
|
||||
QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0]
|
||||
)
|
||||
self.assertEqual(
|
||||
QueueFrame(FrameType.SENTENCE, " How are"), output_frames[1]
|
||||
)
|
||||
|
||||
async def test_two_sentences(self) -> None:
|
||||
input_frames = [
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "hello, "),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, "world."),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, " How are"),
|
||||
QueueFrame(FrameType.TEXT_CHUNK, " you doing?"),
|
||||
QueueFrame(FrameType.END_STREAM, None),
|
||||
]
|
||||
|
||||
service = SentenceAggregator()
|
||||
output_frames = []
|
||||
async for frame in service.run(set([FrameType.SENTENCE]), input_frames):
|
||||
output_frames.append(frame)
|
||||
|
||||
self.assertEqual(2, len(output_frames))
|
||||
self.assertEqual(
|
||||
QueueFrame(FrameType.SENTENCE, "hello, world."), output_frames[0]
|
||||
)
|
||||
self.assertEqual(QueueFrame(FrameType.SENTENCE, " How are you doing?"), output_frames[1])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -18,7 +18,7 @@ from dailyai.services.ai_services import (
|
||||
LLMService,
|
||||
TTSService,
|
||||
)
|
||||
|
||||
"""
|
||||
class MockTTSService(TTSService):
|
||||
def run_tts(self, sentence):
|
||||
for word in sentence.split(' '):
|
||||
@@ -73,7 +73,7 @@ class TestResponse(unittest.TestCase):
|
||||
while expected_words:
|
||||
actual_word:QueueFrame = output_queue.get()
|
||||
word = expected_words.pop(0)
|
||||
self.assertEqual(actual_word.frame_type, FrameType.AUDIO)
|
||||
self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME)
|
||||
self.assertEqual(actual_word.frame_data, bytes(word, "utf-8"))
|
||||
output_queue.task_done()
|
||||
|
||||
@@ -128,10 +128,10 @@ class TestResponse(unittest.TestCase):
|
||||
while expected_words and not stop_processing_output_queue.is_set():
|
||||
try:
|
||||
actual_word:QueueFrame = output_queue.get_nowait()
|
||||
if actual_word.frame_type == FrameType.AUDIO:
|
||||
if actual_word.frame_type == FrameType.AUDIO_FRAME:
|
||||
time.sleep(0.1)
|
||||
word = expected_words.pop(0)
|
||||
self.assertEqual(actual_word.frame_type, FrameType.AUDIO)
|
||||
self.assertEqual(actual_word.frame_type, FrameType.AUDIO_FRAME)
|
||||
self.assertEqual(actual_word.frame_data, bytes(word, "utf-8"))
|
||||
output_queue.task_done()
|
||||
except Empty:
|
||||
@@ -177,3 +177,4 @@ class TestResponse(unittest.TestCase):
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user