Fix off by one errors, add tests and comment
This commit is contained in:
@@ -65,7 +65,7 @@ async def main(room_url: str):
|
||||
simple_tts_pipeline = Pipeline([azure_tts])
|
||||
await simple_tts_pipeline.queue_frames(
|
||||
[
|
||||
TextFrame("My friend the LLM is going to tell a joke about llamas"),
|
||||
TextFrame("My friend the LLM is going to tell a joke about llamas."),
|
||||
EndPipeFrame(),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -18,7 +18,8 @@ class Pipeline:
|
||||
self,
|
||||
processors: List[FrameProcessor],
|
||||
source: asyncio.Queue | None = None,
|
||||
sink: asyncio.Queue[Frame] | None = None
|
||||
sink: asyncio.Queue[Frame] | None = None,
|
||||
name: str | None = None,
|
||||
):
|
||||
"""Create a new pipeline. By default we create the sink and source queues
|
||||
if they're not provided, but these can be overridden to point to other
|
||||
@@ -31,6 +32,9 @@ class Pipeline:
|
||||
self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
|
||||
|
||||
self._logger = logging.getLogger("dailyai.pipeline")
|
||||
self._last_log_line = ""
|
||||
self._shown_repeated_log = False
|
||||
self._name = name or str(id(self))
|
||||
|
||||
def set_source(self, source: asyncio.Queue[Frame]):
|
||||
"""Set the source queue for this pipeline. Frames from this queue
|
||||
@@ -85,7 +89,6 @@ class Pipeline:
|
||||
try:
|
||||
while True:
|
||||
initial_frame = await self.source.get()
|
||||
self._log_frame(initial_frame, 0)
|
||||
async for frame in self._run_pipeline_recursively(
|
||||
initial_frame, self._processors
|
||||
):
|
||||
@@ -118,15 +121,29 @@ class Pipeline:
|
||||
yield initial_frame
|
||||
|
||||
def _log_frame(self, frame: Frame, depth: int):
|
||||
if not isinstance(frame, AudioFrame):
|
||||
if depth == 0:
|
||||
source = "source"
|
||||
dest = str(self._processors[depth - 1])
|
||||
elif depth == len(self._processors) + 1:
|
||||
dest = "sink"
|
||||
source = str(self._processors[depth - 2])
|
||||
else:
|
||||
dest = str(self._processors[depth - 1])
|
||||
source = str(self._processors[depth - 2])
|
||||
"""Log a frame as it moves through the pipeline. This is useful for debugging.
|
||||
Note that this function inherits the logging level from the "dailyai" logger.
|
||||
If you want debug output from dailyai in general but not this function (it is
|
||||
noisy) you can silence this function by doing something like this:
|
||||
|
||||
self._logger.debug(" " * (depth + 1) + " -> ".join([source, str(frame), dest]))
|
||||
# enable debug logging for the dailyai package.
|
||||
logger = logging.getLogger("dailyai")
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
# silence the pipeline logging
|
||||
logger = logging.getLogger("dailyai.pipeline")
|
||||
logger.setLevel(logging.WARNING)
|
||||
"""
|
||||
source = str(self._processors[depth - 2]) if depth > 1 else "source"
|
||||
dest = str(self._processors[depth - 1]) if depth < (len(self._processors) + 1) else "sink"
|
||||
prefix = self._name + " " * depth
|
||||
logline = prefix + " -> ".join([source, frame.__class__.__name__, dest])
|
||||
if logline == self._last_log_line:
|
||||
if self._shown_repeated_log:
|
||||
return
|
||||
self._shown_repeated_log = True
|
||||
self._logger.debug(prefix + "... repeated")
|
||||
else:
|
||||
self._shown_repeated_log = False
|
||||
self._last_log_line = logline
|
||||
self._logger.debug(logline)
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import Mock
|
||||
|
||||
from dailyai.pipeline.aggregators import SentenceAggregator, StatelessTextTransformer
|
||||
from dailyai.pipeline.frame_processor import FrameProcessor
|
||||
from dailyai.pipeline.frames import EndFrame, TextFrame
|
||||
|
||||
from dailyai.pipeline.pipeline import Pipeline
|
||||
@@ -57,3 +60,52 @@ class TestDailyPipeline(unittest.IsolatedAsyncioTestCase):
|
||||
TextFrame(" "),
|
||||
)
|
||||
self.assertIsInstance(await outgoing_queue.get(), EndFrame)
|
||||
|
||||
|
||||
class TestLogFrame(unittest.TestCase):
|
||||
class MockProcessor(FrameProcessor):
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
def setUp(self):
|
||||
self.processor1 = self.MockProcessor('processor1')
|
||||
self.processor2 = self.MockProcessor('processor2')
|
||||
self.pipeline = Pipeline(
|
||||
processors=[self.processor1, self.processor2])
|
||||
self.pipeline._name = 'MyClass'
|
||||
self.pipeline._logger = Mock()
|
||||
|
||||
def test_log_frame_from_source(self):
|
||||
frame = Mock(__class__=Mock(__name__='MyFrame'))
|
||||
self.pipeline._log_frame(frame, depth=1)
|
||||
self.pipeline._logger.debug.assert_called_once_with(
|
||||
'MyClass source -> MyFrame -> processor1')
|
||||
|
||||
def test_log_frame_to_sink(self):
|
||||
frame = Mock(__class__=Mock(__name__='MyFrame'))
|
||||
self.pipeline._log_frame(frame, depth=3)
|
||||
self.pipeline._logger.debug.assert_called_once_with(
|
||||
'MyClass processor2 -> MyFrame -> sink')
|
||||
|
||||
def test_log_frame_repeated_log(self):
|
||||
frame = Mock(__class__=Mock(__name__='MyFrame'))
|
||||
self.pipeline._log_frame(frame, depth=2)
|
||||
self.pipeline._logger.debug.assert_called_once_with(
|
||||
'MyClass processor1 -> MyFrame -> processor2')
|
||||
self.pipeline._log_frame(frame, depth=2)
|
||||
self.pipeline._logger.debug.assert_called_with('MyClass ... repeated')
|
||||
|
||||
def test_log_frame_reset_repeated_log(self):
|
||||
frame1 = Mock(__class__=Mock(__name__='MyFrame1'))
|
||||
frame2 = Mock(__class__=Mock(__name__='MyFrame2'))
|
||||
self.pipeline._log_frame(frame1, depth=2)
|
||||
self.pipeline._logger.debug.assert_called_once_with(
|
||||
'MyClass processor1 -> MyFrame1 -> processor2')
|
||||
self.pipeline._log_frame(frame1, depth=2)
|
||||
self.pipeline._logger.debug.assert_called_with('MyClass ... repeated')
|
||||
self.pipeline._log_frame(frame2, depth=2)
|
||||
self.pipeline._logger.debug.assert_called_with(
|
||||
'MyClass processor1 -> MyFrame2 -> processor2')
|
||||
|
||||
Reference in New Issue
Block a user