Add support for audio frame processor

livekit-agents/livekit/agents/voice/room_io/_input.py

@@ -29,6 +29,7 @@ def __init__(
         room: rtc.Room,
         *,
         track_source: rtc.TrackSource.ValueType | list[rtc.TrackSource.ValueType],
+        processor: rtc.FrameProcessor[T] | None = None,
     ) -> None:
         self._room = room
         self._accepted_sources = (
@@ -48,6 +49,9 @@ def __init__(
 
         self._room.on("track_subscribed", self._on_track_available)
         self._room.on("track_unpublished", self._on_track_unavailable)
+        self._room.on("token_refreshed", self._on_token_refreshed)
+
+        self._processor = processor
 
     async def __anext__(self) -> T:
         return await self._data_ch.__anext__()
@@ -122,6 +126,8 @@ async def aclose(self) -> None:
 
         self._room.off("track_subscribed", self._on_track_available)
         self._data_ch.close()
+        if self._processor:
+            self._processor._close()
 
     @log_exceptions(logger=logger)
     async def _forward_task(
@@ -174,6 +180,16 @@ def _on_track_available(
         self._close_stream()
         self._stream = self._create_stream(track)
         self._publication = publication
+        if self._processor:
+            self._processor._update_stream_info(
+                room_name=self._room.name,
+                participant_identity=participant.identity,
+                publication_sid=publication.sid,
+            )
+            if self._room._token is not None and self._room._server_url is not None:
+                self._processor._update_credentials(
+                    token=self._room._token, url=self._room._server_url
+                )
         self._forward_atask = asyncio.create_task(
             self._forward_task(self._forward_atask, self._stream, publication, participant)
         )
@@ -198,6 +214,14 @@ def _on_track_unavailable(
             if self._on_track_available(publication.track, publication, participant):
                 return
 
+    def _on_token_refreshed(self) -> None:
+        if (
+            self._processor is not None
+            and self._room._token is not None
+            and self._room._server_url is not None
+        ):
+            self._processor._update_credentials(token=self._room._token, url=self._room._server_url)
+
 
 class _ParticipantAudioInputStream(_ParticipantInputStream[rtc.AudioFrame], AudioInput):
     def __init__(
@@ -206,11 +230,20 @@ def __init__(
         *,
         sample_rate: int,
         num_channels: int,
-        noise_cancellation: rtc.NoiseCancellationOptions | None,
+        noise_cancellation: rtc.NoiseCancellationOptions
+        | rtc.FrameProcessor[rtc.AudioFrame]
+        | None,
         pre_connect_audio_handler: PreConnectAudioHandler | None,
     ) -> None:
+        audio_processor: rtc.FrameProcessor[rtc.AudioFrame] | None = None
+        if isinstance(noise_cancellation, rtc.FrameProcessor):
+            audio_processor = noise_cancellation
+
         _ParticipantInputStream.__init__(
-            self, room=room, track_source=rtc.TrackSource.SOURCE_MICROPHONE
+            self,
+            room=room,
+            track_source=rtc.TrackSource.SOURCE_MICROPHONE,
+            processor=audio_processor,
         )
         AudioInput.__init__(self, label="RoomIO")
         self._sample_rate = sample_rate
@@ -251,7 +284,7 @@ async def _forward_task(
             try:
                 duration: float = 0
                 frames = await self._pre_connect_audio_handler.wait_for_data(publication.track.sid)
-                for frame in self._resample_frames(frames):
+                for frame in self._resample_frames(self._apply_audio_processor(frames)):
                     if self._attached:
                         await self._data_ch.send(frame)
                         duration += frame.duration
@@ -305,6 +338,13 @@ def _resample_frames(self, frames: Iterable[rtc.AudioFrame]) -> Iterable[rtc.Aud
         if resampler:
             yield from resampler.flush()
 
+    def _apply_audio_processor(self, frames: Iterable[rtc.AudioFrame]) -> Iterable[rtc.AudioFrame]:
+        for frame in frames:
+            if self._processor is not None:
+                yield self._processor._process(frame)
+            else:
+                yield frame
+
 
 class _ParticipantVideoInputStream(_ParticipantInputStream[rtc.VideoFrame], VideoInput):
     def __init__(self, room: rtc.Room) -> None: