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

Remove Queue in frame names

Browse Source
This commit is contained in:
Moishe Lettvin
2024-03-06 14:09:06 -05:00
parent b9556716dd
commit 62fd371b97
25 changed files with 241 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
src/dailyai/pipeline/aggregators.py
View File

@@ -5,13 +5,13 @@ from tblib import Frame
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import (
ControlQueueFrame,
EndParallelPipeQueueFrame,
EndStreamQueueFrame,
ControlFrame,
EndPipeFrame,
EndFrame,
LLMMessagesQueueFrame,
LLMResponseEndQueueFrame,
QueueFrame,
TextQueueFrame,
LLMResponseEndFrame,
Frame,
TextFrame,
TranscriptionQueueFrame,
)
from dailyai.pipeline.pipeline import Pipeline
@@ -38,10 +38,10 @@ class LLMContextAggregator(AIService):
self.pass_through = pass_through
async def process_frame(
self, frame: QueueFrame
) -> AsyncGenerator[QueueFrame, None]:
self, frame: Frame
) -> AsyncGenerator[Frame, None]:
# We don't do anything with non-text frames, pass it along to next in the pipeline.
if not isinstance(frame, TextQueueFrame):
if not isinstance(frame, TextFrame):
yield frame
return
@@ -71,7 +71,7 @@ class LLMContextAggregator(AIService):
self.messages.append({"role": self.role, "content": frame.text})
yield LLMMessagesQueueFrame(self.messages)
async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
async def finalize(self) -> AsyncGenerator[Frame, None]:
# Send any dangling words that weren't finished with punctuation.
if self.complete_sentences and self.sentence:
self.messages.append({"role": self.role, "content": self.sentence})
@@ -106,18 +106,18 @@ class SentenceAggregator(FrameProcessor):
self.aggregation = ""
async def process_frame(
self, frame: QueueFrame
) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, TextQueueFrame):
self, frame: Frame
) -> AsyncGenerator[Frame, None]:
if isinstance(frame, TextFrame):
m = re.search("(.*[?.!])(.*)", frame.text)
if m:
yield TextQueueFrame(self.aggregation + m.group(1))
yield TextFrame(self.aggregation + m.group(1))
self.aggregation = m.group(2)
else:
self.aggregation += frame.text
elif isinstance(frame, EndStreamQueueFrame):
elif isinstance(frame, EndFrame):
if self.aggregation:
yield TextQueueFrame(self.aggregation)
yield TextFrame(self.aggregation)
yield frame
else:
yield frame
@@ -128,12 +128,12 @@ class LLMFullResponseAggregator(FrameProcessor):
self.aggregation = ""
async def process_frame(
self, frame: QueueFrame
) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, TextQueueFrame):
self, frame: Frame
) -> AsyncGenerator[Frame, None]:
if isinstance(frame, TextFrame):
self.aggregation += frame.text
elif isinstance(frame, LLMResponseEndQueueFrame):
yield TextQueueFrame(self.aggregation)
elif isinstance(frame, LLMResponseEndFrame):
yield TextFrame(self.aggregation)
self.aggregation = ""
else:
yield frame
@@ -143,20 +143,20 @@ class StatelessTextTransformer(FrameProcessor):
def __init__(self, transform_fn):
self.transform_fn = transform_fn
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, TextQueueFrame):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, TextFrame):
result = self.transform_fn(frame.text)
if isinstance(result, Coroutine):
result = await result
yield TextQueueFrame(result)
yield TextFrame(result)
else:
yield frame
class ParallelPipeline(FrameProcessor):
def __init__(self, pipeline_definitions: List[List[FrameProcessor]]):
self.sources = [asyncio.Queue() for _ in pipeline_definitions]
self.sink: asyncio.Queue[QueueFrame] = asyncio.Queue()
self.sink: asyncio.Queue[Frame] = asyncio.Queue()
self.pipelines: list[Pipeline] = [
Pipeline(
pipeline_definition,
@@ -166,10 +166,10 @@ class ParallelPipeline(FrameProcessor):
for source, pipeline_definition in zip(self.sources, pipeline_definitions)
]
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
for source in self.sources:
await source.put(frame)
await source.put(EndParallelPipeQueueFrame())
await source.put(EndPipeFrame())
await asyncio.gather(*[pipeline.run_pipeline() for pipeline in self.pipelines])
@@ -186,7 +186,7 @@ class ParallelPipeline(FrameProcessor):
seen_ids.add(id(frame))
# Skip passing along EndParallelPipeQueueFrame, because we use them for our own flow control.
if not isinstance(frame, EndParallelPipeQueueFrame):
if not isinstance(frame, EndPipeFrame):
yield frame
class GatedAggregator(FrameProcessor):
@@ -194,9 +194,9 @@ class GatedAggregator(FrameProcessor):
self.gate_open_fn = gate_open_fn
self.gate_close_fn = gate_close_fn
self.gate_open = start_open
self.accumulator: List[QueueFrame] = []
self.accumulator: List[Frame] = []
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if self.gate_open:
if self.gate_close_fn(frame):
self.gate_open = False

12
src/dailyai/pipeline/frame_processor.py
View File

@@ -1,7 +1,7 @@
from abc import abstractmethod
from typing import AsyncGenerator
from dailyai.pipeline.frames import ControlQueueFrame, QueueFrame
from dailyai.pipeline.frames import ControlFrame, Frame
"""
This is the base class for all frame processors. Frame processors consume a frame
@@ -20,16 +20,16 @@ be closed, del'd, etc.
class FrameProcessor:
@abstractmethod
async def process_frame(
self, frame: QueueFrame
) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, ControlQueueFrame):
self, frame: Frame
) -> AsyncGenerator[Frame, None]:
if isinstance(frame, ControlFrame):
yield frame
@abstractmethod
async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
async def finalize(self) -> AsyncGenerator[Frame, None]:
# This is a trick for the interpreter (and linter) to know that this is a generator.
if False:
yield QueueFrame()
yield Frame()
@abstractmethod
async def interrupted(self) -> None:

45
src/dailyai/pipeline/frames.py
View File

@@ -2,72 +2,73 @@ from dataclasses import dataclass
from typing import Any
class QueueFrame:
class Frame:
pass
class ControlFrame(Frame):
# Control frames should contain no instance data, so
# equality is based solely on the class.
def __eq__(self, other):
return isinstance(other, self.__class__)
return type(other) == self.__class__
class ControlQueueFrame(QueueFrame):
class StartFrame(ControlFrame):
pass
class StartStreamQueueFrame(ControlQueueFrame):
class EndFrame(ControlFrame):
pass
class EndPipeFrame(ControlFrame):
pass
class EndStreamQueueFrame(ControlQueueFrame):
pass
class EndParallelPipeQueueFrame(ControlQueueFrame):
class LLMResponseStartFrame(ControlFrame):
pass
class LLMResponseStartQueueFrame(QueueFrame):
pass
class LLMResponseEndQueueFrame(QueueFrame):
class LLMResponseEndFrame(ControlFrame):
pass
@dataclass()
class AudioQueueFrame(QueueFrame):
class AudioFrame(Frame):
data: bytes
@dataclass()
class ImageQueueFrame(QueueFrame):
class ImageFrame(Frame):
url: str | None
image: bytes
@dataclass()
class SpriteQueueFrame(QueueFrame):
class SpriteFrame(Frame):
images: list[bytes]
@dataclass()
class TextQueueFrame(QueueFrame):
class TextFrame(Frame):
text: str
@dataclass()
class TranscriptionQueueFrame(TextQueueFrame):
class TranscriptionQueueFrame(TextFrame):
participantId: str
timestamp: str
@dataclass()
class LLMMessagesQueueFrame(QueueFrame):
class LLMMessagesQueueFrame(Frame):
messages: list[dict[str, str]] # TODO: define this more concretely!
class AppMessageQueueFrame(QueueFrame):
class AppMessageQueueFrame(Frame):
message: Any
participantId: str
class UserStartedSpeakingFrame(QueueFrame):
class UserStartedSpeakingFrame(Frame):
pass
class UserStoppedSpeakingFrame(QueueFrame):
class UserStoppedSpeakingFrame(Frame):
pass

18
src/dailyai/pipeline/pipeline.py
View File

@@ -2,7 +2,7 @@ import asyncio
from typing import AsyncGenerator, List
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import EndParallelPipeQueueFrame, EndStreamQueueFrame, QueueFrame
from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
"""
This class manages a pipe of FrameProcessors, and runs them in sequence. The "source"
@@ -17,19 +17,19 @@ class Pipeline:
self,
processors: List[FrameProcessor],
source: asyncio.Queue | None = None,
sink: asyncio.Queue[QueueFrame] | None = None,
sink: asyncio.Queue[Frame] | None = None,
):
self.processors = processors
self.source: asyncio.Queue[QueueFrame] | None = source
self.sink: asyncio.Queue[QueueFrame] | None = sink
self.source: asyncio.Queue[Frame] | None = source
self.sink: asyncio.Queue[Frame] | None = sink
def set_source(self, source: asyncio.Queue[QueueFrame]):
def set_source(self, source: asyncio.Queue[Frame]):
self.source = source
def set_sink(self, sink: asyncio.Queue[QueueFrame]):
def set_sink(self, sink: asyncio.Queue[Frame]):
self.sink = sink
async def get_next_source_frame(self) -> AsyncGenerator[QueueFrame, None]:
async def get_next_source_frame(self) -> AsyncGenerator[Frame, None]:
if self.source is None:
raise ValueError("Source queue not set")
yield await self.source.get()
@@ -52,9 +52,9 @@ class Pipeline:
async for frame in frame_generator:
await self.sink.put(frame)
if isinstance(
frame, EndStreamQueueFrame
frame, EndFrame
) or isinstance(
frame, EndParallelPipeQueueFrame
frame, EndPipeFrame
):
return
except asyncio.CancelledError:

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

@@ -6,14 +6,14 @@ import wave
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
AudioFrame,
EndFrame,
ImageFrame,
LLMMessagesQueueFrame,
LLMResponseEndQueueFrame,
LLMResponseStartQueueFrame,
QueueFrame,
TextQueueFrame,
LLMResponseEndFrame,
LLMResponseStartFrame,
Frame,
TextFrame,
TranscriptionQueueFrame,
)
@@ -33,14 +33,14 @@ class AIService(FrameProcessor):
await queue.put(frame)
if add_end_of_stream:
await queue.put(EndStreamQueueFrame())
await queue.put(EndFrame())
async def run(
self,
frames: Iterable[QueueFrame]
| AsyncIterable[QueueFrame]
| asyncio.Queue[QueueFrame],
) -> AsyncGenerator[QueueFrame, None]:
frames: Iterable[Frame]
| AsyncIterable[Frame]
| asyncio.Queue[Frame],
) -> AsyncGenerator[Frame, None]:
try:
if isinstance(frames, AsyncIterable):
async for frame in frames:
@@ -55,7 +55,7 @@ class AIService(FrameProcessor):
frame = await frames.get()
async for output_frame in self.process_frame(frame):
yield output_frame
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
else:
raise Exception("Frames must be an iterable or async iterable")
@@ -76,12 +76,12 @@ class LLMService(AIService):
async def run_llm(self, messages) -> str:
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, LLMMessagesQueueFrame):
yield LLMResponseStartQueueFrame()
yield LLMResponseStartFrame()
async for text_chunk in self.run_llm_async(frame.messages):
yield TextQueueFrame(text_chunk)
yield LLMResponseEndQueueFrame()
yield TextFrame(text_chunk)
yield LLMResponseEndFrame()
else:
yield frame
@@ -103,8 +103,8 @@ class TTSService(AIService):
# yield empty bytes here, so linting can infer what this method does
yield bytes()
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not isinstance(frame, TextQueueFrame):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if not isinstance(frame, TextFrame):
yield frame
return
@@ -119,16 +119,16 @@ class TTSService(AIService):
if text:
async for audio_chunk in self.run_tts(text):
yield AudioQueueFrame(audio_chunk)
yield AudioFrame(audio_chunk)
async def finalize(self):
if self.current_sentence:
async for audio_chunk in self.run_tts(self.current_sentence):
yield AudioQueueFrame(audio_chunk)
yield AudioFrame(audio_chunk)
# Convenience function to send the audio for a sentence to the given queue
async def say(self, sentence, queue: asyncio.Queue):
await self.run_to_queue(queue, [TextQueueFrame(sentence)])
await self.run_to_queue(queue, [TextFrame(sentence)])
class ImageGenService(AIService):
@@ -141,13 +141,13 @@ class ImageGenService(AIService):
async def run_image_gen(self, sentence: str) -> tuple[str, bytes]:
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if not isinstance(frame, TextQueueFrame):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if not isinstance(frame, TextFrame):
yield frame
return
(url, image_data) = await self.run_image_gen(frame.text)
yield ImageQueueFrame(url, image_data)
yield ImageFrame(url, image_data)
class STTService(AIService):
@@ -164,9 +164,9 @@ class STTService(AIService):
"""Returns transcript as a string"""
pass
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
"""Processes a frame of audio data, either buffering or transcribing it."""
if not isinstance(frame, AudioQueueFrame):
if not isinstance(frame, AudioFrame):
return
data = frame.data
@@ -187,8 +187,8 @@ class FrameLogger(AIService):
super().__init__(**kwargs)
self.prefix = prefix
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, (AudioQueueFrame, ImageQueueFrame)):
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, (AudioFrame, ImageFrame)):
self.logger.info(f"{self.prefix}: {type(frame)}")
else:
print(f"{self.prefix}: {frame}")

50
src/dailyai/services/base_transport_service.py
View File

@@ -12,12 +12,12 @@ from typing import AsyncGenerator
from enum import Enum
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
QueueFrame,
SpriteQueueFrame,
StartStreamQueueFrame,
AudioFrame,
EndFrame,
ImageFrame,
Frame,
SpriteFrame,
StartFrame,
TranscriptionQueueFrame,
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame
@@ -159,7 +159,7 @@ class BaseTransportService():
self._stop_threads.set()
await self.send_queue.put(EndStreamQueueFrame())
await self.send_queue.put(EndFrame())
await async_output_queue_marshal_task
await self.send_queue.join()
self._frame_consumer_thread.join()
@@ -182,7 +182,7 @@ class BaseTransportService():
pipeline.set_sink(self.send_queue)
pipeline_task = asyncio.create_task(pipeline.run_pipeline())
async def yield_frame(frame:QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def yield_frame(frame:Frame) -> AsyncGenerator[Frame, None]:
yield frame
async def post_process(post_processor):
@@ -194,7 +194,7 @@ class BaseTransportService():
print("post-processing frame: ", frame.__class__.__name__)
await post_processor.process_frame(frame)
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
post_process_task = asyncio.create_task(post_process(post_processor))
@@ -214,7 +214,7 @@ class BaseTransportService():
async for frame in frame_generator:
await source_queue.put(frame)
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
await asyncio.gather(pipeline_task, post_process_task)
@@ -303,20 +303,20 @@ class BaseTransportService():
async def _marshal_frames(self):
while True:
frame: QueueFrame | list = await self.send_queue.get()
frame: Frame | list = await self.send_queue.get()
self._threadsafe_send_queue.put(frame)
self.send_queue.task_done()
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
def interrupt(self):
self._is_interrupted.set()
async def get_receive_frames(self) -> AsyncGenerator[QueueFrame, None]:
async def get_receive_frames(self) -> AsyncGenerator[Frame, None]:
while True:
frame = await self.receive_queue.get()
yield frame
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
break
def _receive_audio(self):
@@ -329,13 +329,13 @@ class BaseTransportService():
while not self._stop_threads.is_set():
buffer = self.read_audio_frames(desired_frame_count)
if len(buffer) > 0:
frame = AudioQueueFrame(buffer)
frame = AudioFrame(buffer)
asyncio.run_coroutine_threadsafe(
self.receive_queue.put(frame), self._loop
)
asyncio.run_coroutine_threadsafe(
self.receive_queue.put(EndStreamQueueFrame()), self._loop
self.receive_queue.put(EndFrame()), self._loop
)
def _set_image(self, image: bytes):
@@ -363,18 +363,18 @@ class BaseTransportService():
all_audio_frames = bytearray()
while True:
try:
frames_or_frame: QueueFrame | list[QueueFrame] = (
frames_or_frame: Frame | list[Frame] = (
self._threadsafe_send_queue.get()
)
if isinstance(frames_or_frame, QueueFrame):
frames: list[QueueFrame] = [frames_or_frame]
if isinstance(frames_or_frame, Frame):
frames: list[Frame] = [frames_or_frame]
elif isinstance(frames_or_frame, list):
frames: list[QueueFrame] = frames_or_frame
frames: list[Frame] = frames_or_frame
else:
raise Exception("Unknown type in output queue")
for frame in frames:
if isinstance(frame, EndStreamQueueFrame):
if isinstance(frame, EndFrame):
self._logger.info("Stopping frame consumer thread")
self._threadsafe_send_queue.task_done()
if self._loop:
@@ -386,7 +386,7 @@ class BaseTransportService():
# if interrupted, we just pull frames off the queue and discard them
if not self._is_interrupted.is_set():
if frame:
if isinstance(frame, AudioQueueFrame):
if isinstance(frame, AudioFrame):
chunk = frame.data
all_audio_frames.extend(chunk)
@@ -398,9 +398,9 @@ class BaseTransportService():
if truncated_length:
self.write_frame_to_mic(bytes(b[:truncated_length]))
b = b[truncated_length:]
elif isinstance(frame, ImageQueueFrame):
elif isinstance(frame, ImageFrame):
self._set_image(frame.image)
elif isinstance(frame, SpriteQueueFrame):
elif isinstance(frame, SpriteFrame):
self._set_images(frame.images)
elif len(b):
self.write_frame_to_mic(bytes(b))
@@ -418,7 +418,7 @@ class BaseTransportService():
self.write_frame_to_mic(bytes(b[:truncated_length]))
b = bytearray()
if isinstance(frame, StartStreamQueueFrame):
if isinstance(frame, StartFrame):
self._is_interrupted.clear()
self._threadsafe_send_queue.task_done()

6
src/dailyai/services/local_stt_service.py
View File

@@ -4,7 +4,7 @@ import math
import time
from typing import AsyncGenerator
import wave
from dailyai.pipeline.frames import AudioQueueFrame, QueueFrame, TranscriptionQueueFrame
from dailyai.pipeline.frames import AudioFrame, Frame, TranscriptionQueueFrame
from dailyai.services.ai_services import STTService
@@ -39,9 +39,9 @@ class LocalSTTService(STTService):
ww.setframerate(self._frame_rate)
self._wave = ww
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
"""Processes a frame of audio data, either buffering or transcribing it."""
if not isinstance(frame, AudioQueueFrame):
if not isinstance(frame, AudioFrame):
return
data = frame.data

76
src/dailyai/tests/test_aggregators.py
View File

@@ -9,13 +9,13 @@ from dailyai.pipeline.aggregators import (
StatelessTextTransformer,
)
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
LLMResponseEndQueueFrame,
LLMResponseStartQueueFrame,
QueueFrame,
TextQueueFrame,
AudioFrame,
EndFrame,
ImageFrame,
LLMResponseEndFrame,
LLMResponseStartFrame,
Frame,
TextFrame,
)
from dailyai.pipeline.pipeline import Pipeline
@@ -27,46 +27,46 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase):
expected_sentences = ["Hello, world.", " How are you?", " I am fine "]
aggregator = SentenceAggregator()
for word in sentence.split(" "):
async for sentence in aggregator.process_frame(TextQueueFrame(word + " ")):
self.assertIsInstance(sentence, TextQueueFrame)
if isinstance(sentence, TextQueueFrame):
async for sentence in aggregator.process_frame(TextFrame(word + " ")):
self.assertIsInstance(sentence, TextFrame)
if isinstance(sentence, TextFrame):
self.assertEqual(sentence.text, expected_sentences.pop(0))
async for sentence in aggregator.process_frame(EndStreamQueueFrame()):
async for sentence in aggregator.process_frame(EndFrame()):
if len(expected_sentences):
self.assertIsInstance(sentence, TextQueueFrame)
if isinstance(sentence, TextQueueFrame):
self.assertIsInstance(sentence, TextFrame)
if isinstance(sentence, TextFrame):
self.assertEqual(sentence.text, expected_sentences.pop(0))
else:
self.assertIsInstance(sentence, EndStreamQueueFrame)
self.assertIsInstance(sentence, EndFrame)
self.assertEqual(expected_sentences, [])
async def test_gated_accumulator(self):
gated_aggregator = GatedAggregator(
gate_open_fn=lambda frame: isinstance(frame, ImageQueueFrame),
gate_close_fn=lambda frame: isinstance(frame, LLMResponseStartQueueFrame),
gate_open_fn=lambda frame: isinstance(frame, ImageFrame),
gate_close_fn=lambda frame: isinstance(frame, LLMResponseStartFrame),
start_open=False,
)
frames = [
LLMResponseStartQueueFrame(),
TextQueueFrame("Hello, "),
TextQueueFrame("world."),
AudioQueueFrame(b"hello"),
ImageQueueFrame("image", b"image"),
AudioQueueFrame(b"world"),
LLMResponseEndQueueFrame(),
LLMResponseStartFrame(),
TextFrame("Hello, "),
TextFrame("world."),
AudioFrame(b"hello"),
ImageFrame("image", b"image"),
AudioFrame(b"world"),
LLMResponseEndFrame(),
]
expected_output_frames = [
ImageQueueFrame("image", b"image"),
LLMResponseStartQueueFrame(),
TextQueueFrame("Hello, "),
TextQueueFrame("world."),
AudioQueueFrame(b"hello"),
AudioQueueFrame(b"world"),
LLMResponseEndQueueFrame(),
ImageFrame("image", b"image"),
LLMResponseStartFrame(),
TextFrame("Hello, "),
TextFrame("world."),
AudioFrame(b"hello"),
AudioFrame(b"world"),
LLMResponseEndFrame(),
]
for frame in frames:
async for out_frame in gated_aggregator.process_frame(frame):
@@ -98,16 +98,16 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase):
)
frames = [
TextQueueFrame("Hello, "),
TextQueueFrame("world."),
EndStreamQueueFrame()
TextFrame("Hello, "),
TextFrame("world."),
EndFrame()
]
expected_output_frames: list[QueueFrame] = [
TextQueueFrame(text='Hello, :pipe1.'),
TextQueueFrame(text='world.:pipe1.'),
TextQueueFrame(text='Hello, world.:pipe2.'),
EndStreamQueueFrame()
expected_output_frames: list[Frame] = [
TextFrame(text='Hello, :pipe1.'),
TextFrame(text='world.:pipe1.'),
TextFrame(text='Hello, world.:pipe2.'),
EndFrame()
]
for frame in frames:

14
src/dailyai/tests/test_ai_services.py
View File

@@ -3,11 +3,11 @@ import unittest
from typing import AsyncGenerator, Generator
from dailyai.services.ai_services import AIService
from dailyai.pipeline.frames import EndStreamQueueFrame, QueueFrame, TextQueueFrame
from dailyai.pipeline.frames import EndFrame, Frame, TextFrame
class SimpleAIService(AIService):
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
yield frame
@@ -16,11 +16,11 @@ class TestBaseAIService(unittest.IsolatedAsyncioTestCase):
service = SimpleAIService()
input_frames = [
TextQueueFrame("hello"),
EndStreamQueueFrame()
TextFrame("hello"),
EndFrame()
]
async def iterate_frames() -> AsyncGenerator[QueueFrame, None]:
async def iterate_frames() -> AsyncGenerator[Frame, None]:
for frame in input_frames:
yield frame
@@ -33,9 +33,9 @@ class TestBaseAIService(unittest.IsolatedAsyncioTestCase):
async def test_nonasync_input(self):
service = SimpleAIService()
input_frames = [TextQueueFrame("hello"), EndStreamQueueFrame()]
input_frames = [TextFrame("hello"), EndFrame()]
def iterate_frames() -> Generator[QueueFrame, None, None]:
def iterate_frames() -> Generator[Frame, None, None]:
for frame in input_frames:
yield frame

2
src/dailyai/tests/test_daily_transport_service.py
View File

@@ -3,7 +3,7 @@ import unittest
from unittest.mock import MagicMock, patch
from dailyai.pipeline.frames import AudioQueueFrame, ImageQueueFrame
from dailyai.pipeline.frames import AudioFrame, ImageFrame
class TestDailyTransport(unittest.IsolatedAsyncioTestCase):

24
src/dailyai/tests/test_pipeline.py
View File

@@ -2,7 +2,7 @@ import asyncio
from doctest import OutputChecker
import unittest
from dailyai.pipeline.aggregators import SentenceAggregator, StatelessTextTransformer
from dailyai.pipeline.frames import EndStreamQueueFrame, TextQueueFrame
from dailyai.pipeline.frames import EndFrame, TextFrame
from dailyai.pipeline.pipeline import Pipeline
@@ -16,14 +16,14 @@ class TestDailyPipeline(unittest.IsolatedAsyncioTestCase):
incoming_queue = asyncio.Queue()
pipeline = Pipeline([aggregator], incoming_queue, outgoing_queue)
await incoming_queue.put(TextQueueFrame("Hello, "))
await incoming_queue.put(TextQueueFrame("world."))
await incoming_queue.put(EndStreamQueueFrame())
await incoming_queue.put(TextFrame("Hello, "))
await incoming_queue.put(TextFrame("world."))
await incoming_queue.put(EndFrame())
await pipeline.run_pipeline()
self.assertEqual(await outgoing_queue.get(), TextQueueFrame("Hello, world."))
self.assertIsInstance(await outgoing_queue.get(), EndStreamQueueFrame)
self.assertEqual(await outgoing_queue.get(), TextFrame("Hello, world."))
self.assertIsInstance(await outgoing_queue.get(), EndFrame)
async def test_pipeline_multiple_stages(self):
sentence_aggregator = SentenceAggregator()
@@ -40,21 +40,21 @@ class TestDailyPipeline(unittest.IsolatedAsyncioTestCase):
sentence = "Hello, world. It's me, a pipeline."
for c in sentence:
await incoming_queue.put(TextQueueFrame(c))
await incoming_queue.put(EndStreamQueueFrame())
await incoming_queue.put(TextFrame(c))
await incoming_queue.put(EndFrame())
await pipeline.run_pipeline()
self.assertEqual(
await outgoing_queue.get(), TextQueueFrame("H E L L O , W O R L D .")
await outgoing_queue.get(), TextFrame("H E L L O , W O R L D .")
)
self.assertEqual(
await outgoing_queue.get(),
TextQueueFrame(" I T ' S M E , A P I P E L I N E ."),
TextFrame(" I T ' S M E , A P I P E L I N E ."),
)
# leftover little bit because of the spacing
self.assertEqual(
await outgoing_queue.get(),
TextQueueFrame(" "),
TextFrame(" "),
)
self.assertIsInstance(await outgoing_queue.get(), EndStreamQueueFrame)
self.assertIsInstance(await outgoing_queue.get(), EndFrame)