130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
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()
|