livekit 1.2.1__py3-none-win_arm64.whl

livekit/rtc/audio_stream.py

@@ -0,0 +1,354 @@
+# Copyright 2023 LiveKit, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import asyncio
+import json
+from dataclasses import dataclass
+from typing import Any, AsyncIterator, Optional
+
+from ._ffi_client import FfiClient, FfiHandle
+from ._proto import audio_frame_pb2 as proto_audio_frame
+from ._proto import ffi_pb2 as proto_ffi
+from ._proto.track_pb2 import TrackSource
+from ._utils import RingQueue, task_done_logger
+from .audio_frame import AudioFrame
+from .log import logger
+from .participant import Participant
+from .track import Track
+from .frame_processor import FrameProcessor
+
+
+@dataclass
+class AudioFrameEvent:
+    """An event representing a received audio frame.
+
+    Attributes:
+        frame (AudioFrame): The received audio frame.
+    """
+
+    frame: AudioFrame
+
+
+@dataclass
+class NoiseCancellationOptions:
+    module_id: str
+    options: dict[str, Any]
+
+
+class AudioStream:
+    """An asynchronous audio stream for receiving audio frames from a participant or track.
+
+    The `AudioStream` class provides an asynchronous iterator over audio frames received from
+    a specific track or participant. It allows you to receive audio frames in real-time with
+    customizable sample rates and channel configurations.
+    """
+
+    def __init__(
+        self,
+        track: Track,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+        capacity: int = 0,
+        sample_rate: int = 48000,
+        num_channels: int = 1,
+        frame_size_ms: int | None = None,
+        noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize an `AudioStream` instance.
+
+        Args:
+            track (Optional[Track]): The audio track from which to receive audio. If not provided,
+                you must specify `participant` and `track_source` in `kwargs`.
+            loop (Optional[asyncio.AbstractEventLoop], optional): The event loop to use.
+                Defaults to the current event loop.
+            capacity (int, optional): The capacity of the internal frame queue. Defaults to 0 (unbounded).
+            sample_rate (int, optional): The sample rate for the audio stream in Hz.
+                Defaults to 48000.
+            num_channels (int, optional): The number of audio channels. Defaults to 1.
+            noise_cancellation (Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]], optional):
+                If noise cancellation is used, pass a `NoiseCancellationOptions` or `FrameProcessor[AudioFrame]` instance
+                created by the noise cancellation module.
+            auto_close_noise_cancellation (bool):
+                When the audio stream closes, should the FrameProcessor's close method be run? If
+                False, then the frame processor can be reused with another AudioStream.
+
+        Example:
+            ```python
+            audio_stream = AudioStream(
+                track=audio_track,
+                sample_rate=44100,
+                num_channels=2,
+            )
+
+            audio_stream = AudioStream.from_track(
+                track=audio_track,
+                sample_rate=44100,
+                num_channels=2,
+            )
+            ```
+        """
+        self._track: Track | None = track
+        self._sample_rate = sample_rate
+        self._num_channels = num_channels
+        self._frame_size_ms = frame_size_ms
+        self._loop = loop or asyncio.get_event_loop()
+        # Subscribe before creating the stream so none of its events are missed. Its
+        # handle is only known once the request returns; until then every audio stream
+        # event is let through, and _run drops the ones of other streams.
+        handle_id: int | None = None
+        self._ffi_queue = FfiClient.instance.queue.subscribe(
+            self._loop,
+            filter_fn=lambda e: (
+                e.WhichOneof("message") == "audio_stream_event"
+                and (handle_id is None or e.audio_stream_event.stream_handle == handle_id)
+            ),
+        )
+        self._queue: RingQueue[AudioFrameEvent | None] = RingQueue(capacity)
+
+        self._audio_filter_module: str | None = None
+        self._audio_filter_options: dict[str, Any] | None = None
+        self._processor: FrameProcessor[AudioFrame] | None = None
+        self._processor_auto_close = True
+        if isinstance(noise_cancellation, NoiseCancellationOptions):
+            self._audio_filter_module = noise_cancellation.module_id
+            self._audio_filter_options = noise_cancellation.options
+        elif isinstance(noise_cancellation, FrameProcessor):
+            self._processor = noise_cancellation
+            self._processor_auto_close = auto_close_noise_cancellation
+
+        stream: Any = None
+        if "participant" in kwargs:
+            stream = self._create_owned_stream_from_participant(
+                participant=kwargs["participant"], track_source=kwargs["track_source"]
+            )
+        else:
+            stream = self._create_owned_stream()
+        self._ffi_handle = FfiHandle(stream.handle.id)
+        handle_id = stream.handle.id
+        self._info = stream.info
+
+        # Start _run only after _ffi_handle is set, so a failure above doesn't orphan
+        # a task that dereferences an unassigned _ffi_handle.
+        self._task = self._loop.create_task(self._run())
+        self._task.add_done_callback(task_done_logger)
+
+        if self._track is not None:
+            self._track._register_audio_stream(self)
+
+    @classmethod
+    def from_participant(
+        cls,
+        *,
+        participant: Participant,
+        track_source: TrackSource.ValueType,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+        capacity: int = 0,
+        sample_rate: int = 48000,
+        num_channels: int = 1,
+        frame_size_ms: int | None = None,
+        noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = True,
+    ) -> AudioStream:
+        """Create an `AudioStream` from a participant's audio track.
+
+        Args:
+            participant (Participant): The participant from whom to receive audio.
+            track_source (TrackSource.ValueType): The source of the audio track (e.g., microphone, screen share).
+            loop (Optional[asyncio.AbstractEventLoop], optional): The event loop to use. Defaults to the current event loop.
+            capacity (int, optional): The capacity of the internal frame queue. Defaults to 0 (unbounded).
+            sample_rate (int, optional): The sample rate for the audio stream in Hz. Defaults to 48000.
+            num_channels (int, optional): The number of audio channels. Defaults to 1.
+            noise_cancellation (Optional[NoiseCancellationOptions], optional):
+                If noise cancellation is used, pass a `NoiseCancellationOptions` instance
+                created by the noise cancellation module.
+
+        Returns:
+            AudioStream: An instance of `AudioStream` that can be used to receive audio frames.
+
+        Example:
+            ```python
+            audio_stream = AudioStream.from_participant(
+                participant=participant,
+                track_source=TrackSource.MICROPHONE,
+                sample_rate=24000,
+                num_channels=1,
+            )
+            ```
+        """
+        return AudioStream(
+            participant=participant,
+            track_source=track_source,
+            loop=loop,
+            capacity=capacity,
+            track=None,  # type: ignore
+            sample_rate=sample_rate,
+            num_channels=num_channels,
+            frame_size_ms=frame_size_ms,
+            noise_cancellation=noise_cancellation,
+            auto_close_noise_cancellation=auto_close_noise_cancellation,
+        )
+
+    @classmethod
+    def from_track(
+        cls,
+        *,
+        track: Track,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+        capacity: int = 0,
+        sample_rate: int = 48000,
+        num_channels: int = 1,
+        frame_size_ms: int | None = None,
+        noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = False,
+    ) -> AudioStream:
+        """Create an `AudioStream` from an existing audio track.
+
+        Args:
+            track (Track): The audio track from which to receive audio.
+            loop (Optional[asyncio.AbstractEventLoop], optional): The event loop to use. Defaults to the current event loop.
+            capacity (int, optional): The capacity of the internal frame queue. Defaults to 0 (unbounded).
+            sample_rate (int, optional): The sample rate for the audio stream in Hz. Defaults to 48000.
+            num_channels (int, optional): The number of audio channels. Defaults to 1.
+            noise_cancellation (Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]], optional):
+                If noise cancellation is used, pass a `NoiseCancellationOptions` or `FrameProcessor[AudioFrame]` instance
+                created by the noise cancellation module.
+            auto_close_noise_cancellation (bool):
+                When the audio stream closes, leaves the FrameProcessor in an unclosed state so it
+                can be used with another AudioStream.
+
+        Returns:
+            AudioStream: An instance of `AudioStream` that can be used to receive audio frames.
+
+        Example:
+            ```python
+            audio_stream = AudioStream.from_track(
+                track=audio_track,
+                sample_rate=44100,
+                num_channels=2,
+            )
+            ```
+        """
+        return AudioStream(
+            track=track,
+            loop=loop,
+            capacity=capacity,
+            sample_rate=sample_rate,
+            num_channels=num_channels,
+            frame_size_ms=frame_size_ms,
+            noise_cancellation=noise_cancellation,
+            auto_close_noise_cancellation=auto_close_noise_cancellation,
+        )
+
+    def __del__(self) -> None:
+        # _ffi_queue is unset if __init__ raised early or the instance was built via __new__
+        ffi_queue = getattr(self, "_ffi_queue", None)
+        if ffi_queue is not None:
+            FfiClient.instance.queue.unsubscribe(ffi_queue)
+
+    def _create_owned_stream(self) -> Any:
+        assert self._track is not None
+        req = proto_ffi.FfiRequest()
+        new_audio_stream = req.new_audio_stream
+        new_audio_stream.track_handle = self._track._ffi_handle.handle
+        new_audio_stream.sample_rate = self._sample_rate
+        new_audio_stream.num_channels = self._num_channels
+        new_audio_stream.queue_size_frames = 0
+        if self._frame_size_ms:
+            new_audio_stream.frame_size_ms = self._frame_size_ms
+        new_audio_stream.type = proto_audio_frame.AudioStreamType.AUDIO_STREAM_NATIVE
+        if self._audio_filter_module is not None:
+            new_audio_stream.audio_filter_module_id = self._audio_filter_module
+        if self._audio_filter_options is not None:
+            new_audio_stream.audio_filter_options = json.dumps(self._audio_filter_options)
+        resp = FfiClient.instance.request(req)
+        return resp.new_audio_stream.stream
+
+    def _create_owned_stream_from_participant(
+        self, participant: Participant, track_source: TrackSource.ValueType
+    ) -> Any:
+        req = proto_ffi.FfiRequest()
+        audio_stream_from_participant = req.audio_stream_from_participant
+        audio_stream_from_participant.participant_handle = participant._ffi_handle.handle
+        audio_stream_from_participant.sample_rate = self._sample_rate
+        audio_stream_from_participant.num_channels = self._num_channels
+        audio_stream_from_participant.queue_size_frames = 0
+        audio_stream_from_participant.type = proto_audio_frame.AudioStreamType.AUDIO_STREAM_NATIVE
+        audio_stream_from_participant.track_source = track_source
+        if self._frame_size_ms:
+            audio_stream_from_participant.frame_size_ms = self._frame_size_ms
+
+        if self._audio_filter_module is not None:
+            audio_stream_from_participant.audio_filter_module_id = self._audio_filter_module
+        if self._audio_filter_options is not None:
+            audio_stream_from_participant.audio_filter_options = json.dumps(
+                self._audio_filter_options
+            )
+        resp = FfiClient.instance.request(req)
+        return resp.audio_stream_from_participant.stream
+
+    async def _run(self) -> None:
+        while True:
+            event = await self._ffi_queue.wait_for(self._is_event)
+            audio_event: proto_audio_frame.AudioStreamEvent = event.audio_stream_event
+
+            if audio_event.HasField("frame_received"):
+                owned_buffer_info = audio_event.frame_received.frame
+                frame = AudioFrame._from_owned_info(owned_buffer_info)
+                if self._processor is not None and self._processor.enabled:
+                    try:
+                        frame = self._processor._process(frame)
+                    except Exception:
+                        logger.warning(
+                            "Frame processing failed, passing through original frame",
+                            exc_info=True,
+                        )
+                event = AudioFrameEvent(frame)
+                self._queue.put(event)
+            elif audio_event.HasField("eos"):
+                self._queue.put(None)
+                break
+
+        FfiClient.instance.queue.unsubscribe(self._ffi_queue)
+
+    async def aclose(self) -> None:
+        """Asynchronously close the audio stream.
+
+        This method cleans up resources associated with the audio stream and waits for
+        any pending operations to complete.
+        """
+        if self._track is not None:
+            self._track._unregister_audio_stream(self)
+        self._ffi_handle.dispose()
+        await self._task
+        if self._processor is not None and self._processor_auto_close:
+            self._processor._close()
+
+    def _is_event(self, e: proto_ffi.FfiEvent) -> bool:
+        return e.audio_stream_event.stream_handle == self._ffi_handle.handle
+
+    def __aiter__(self) -> AsyncIterator[AudioFrameEvent]:
+        return self
+
+    async def __anext__(self) -> AudioFrameEvent:
+        if self._task.done():
+            raise StopAsyncIteration
+
+        item = await self._queue.get()
+        if item is None:
+            raise StopAsyncIteration
+
+        return item