From 05c77bce25751c469be76ac0acbb397237b647eb Mon Sep 17 00:00:00 2001
From: Moishe Lettvin <moishel@gmail.com>
Date: Wed, 27 Mar 2024 18:48:30 -0400
Subject: [PATCH] Add logging for pipeline

---
 src/dailyai/pipeline/frame_processor.py |  3 +++
 src/dailyai/pipeline/pipeline.py        | 27 +++++++++++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/dailyai/pipeline/frame_processor.py b/src/dailyai/pipeline/frame_processor.py
index 3f42b6987..d8550c237 100644
--- a/src/dailyai/pipeline/frame_processor.py
+++ b/src/dailyai/pipeline/frame_processor.py
@@ -29,3 +29,6 @@ class FrameProcessor:
     async def interrupted(self) -> None:
         """Handle any cleanup if the pipeline was interrupted."""
         pass
+
+    def __str__(self):
+        return self.__class__.__name__
diff --git a/src/dailyai/pipeline/pipeline.py b/src/dailyai/pipeline/pipeline.py
index fbb4488db..0f23b4903 100644
--- a/src/dailyai/pipeline/pipeline.py
+++ b/src/dailyai/pipeline/pipeline.py
@@ -1,8 +1,9 @@
 import asyncio
+import logging
 from typing import AsyncGenerator, AsyncIterable, Iterable, List
 from dailyai.pipeline.frame_processor import FrameProcessor
 
-from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
+from dailyai.pipeline.frames import AudioFrame, EndPipeFrame, EndFrame, Frame
 
 
 class Pipeline:
@@ -29,6 +30,8 @@ class Pipeline:
         self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
         self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
 
+        self._logger = logging.getLogger("dailyai.pipeline")
+
     def set_source(self, source: asyncio.Queue[Frame]):
         """Set the source queue for this pipeline. Frames from this queue
         will be processed by each frame_processor in the pipeline, or order
@@ -82,9 +85,11 @@ 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
                 ):
+                    self._log_frame(frame, len(self._processors) + 1)
                     await self.sink.put(frame)
 
                 if isinstance(initial_frame, EndFrame) or isinstance(
@@ -96,18 +101,32 @@ class Pipeline:
             # here.
             for processor in self._processors:
                 await processor.interrupted()
-            pass
 
     async def _run_pipeline_recursively(
-        self, initial_frame: Frame, processors: List[FrameProcessor]
+        self, initial_frame: Frame, processors: List[FrameProcessor], depth=1
     ) -> AsyncGenerator[Frame, None]:
         """Internal function to add frames to the pipeline as they're yielded
         by each processor."""
         if processors:
+            self._log_frame(initial_frame, depth)
             async for frame in processors[0].process_frame(initial_frame):
                 async for final_frame in self._run_pipeline_recursively(
-                    frame, processors[1:]
+                    frame, processors[1:], depth + 1
                 ):
                     yield final_frame
         else:
             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])
+
+            self._logger.debug("  " * (depth + 1) + " -> ".join([source, str(frame), dest]))