move handler registration back to room

Author: lukasIO

livekit-rtc/livekit/rtc/participant.py

@@ -26,7 +26,6 @@
 
 from ._ffi_client import FfiClient, FfiHandle
 from ._proto import ffi_pb2 as proto_ffi
-from ._proto import room_pb2 as proto_room
 from ._proto import participant_pb2 as proto_participant
 from ._proto.room_pb2 import (
     TrackPublishOptions,
@@ -52,10 +51,6 @@
     ByteStreamWriter,
     ByteStreamInfo,
     STREAM_CHUNK_SIZE,
-    TextStreamReader,
-    ByteStreamReader,
-    TextStreamHandler,
-    ByteStreamHandler,
 )
 
 
@@ -165,10 +160,6 @@ def __init__(
         self._rpc_handlers: Dict[
             str, Callable[[RpcInvocationData], Union[Awaitable[str], str]]
         ] = {}
-        self._text_stream_readers: Dict[str, TextStreamReader] = {}
-        self._byte_stream_readers: Dict[str, ByteStreamReader] = {}
-        self._text_stream_handlers: Dict[str, TextStreamHandler] = {}
-        self._byte_stream_handlers: Dict[str, ByteStreamHandler] = {}
 
     @property
     def track_publications(self) -> Mapping[str, LocalTrackPublication]:
@@ -559,59 +550,6 @@ async def set_attributes(self, attributes: dict[str, str]) -> None:
         finally:
             FfiClient.instance.queue.unsubscribe(queue)
 
-    def _handle_stream_header(
-        self, header: proto_room.DataStream.Header, participant_identity: str
-    ):
-        stream_type = header.WhichOneof("content_header")
-        if stream_type == "text_header":
-            text_stream_handler = self._text_stream_handlers.get(header.topic)
-            if text_stream_handler is None:
-                logging.info(
-                    "ignoring text stream with topic '%s', no callback attached",
-                    header.topic,
-                )
-                return
-
-            text_reader = TextStreamReader(header)
-            self._text_stream_readers[header.stream_id] = text_reader
-            text_stream_handler(text_reader, participant_identity)
-        elif stream_type == "byte_header":
-            logging.warning("received byte header, %s", header.stream_id)
-            byte_stream_handler = self._byte_stream_handlers.get(header.topic)
-            if byte_stream_handler is None:
-                logging.info(
-                    "ignoring byte stream with topic '%s', no callback attached",
-                    header.topic,
-                )
-                return
-
-            byte_reader = ByteStreamReader(header)
-            self._byte_stream_readers[header.stream_id] = byte_reader
-            byte_stream_handler(byte_reader, participant_identity)
-        else:
-            logging.warning("received unknown header type, %s", stream_type)
-        pass
-
-    async def _handle_stream_chunk(self, chunk: proto_room.DataStream.Chunk):
-        text_reader = self._text_stream_readers.get(chunk.stream_id)
-        file_reader = self._byte_stream_readers.get(chunk.stream_id)
-
-        if text_reader:
-            await text_reader._on_chunk_update(chunk)
-        elif file_reader:
-            await file_reader._on_chunk_update(chunk)
-
-    async def _handle_stream_trailer(self, trailer: proto_room.DataStream.Trailer):
-        text_reader = self._text_stream_readers.get(trailer.stream_id)
-        file_reader = self._byte_stream_readers.get(trailer.stream_id)
-
-        if text_reader:
-            await text_reader._on_stream_close(trailer)
-            self._text_stream_readers.pop(trailer.stream_id)
-        elif file_reader:
-            await file_reader._on_stream_close(trailer)
-            self._byte_stream_readers.pop(trailer.stream_id)
-
     async def stream_text(
         self,
         *,
@@ -726,28 +664,6 @@ async def send_file(
 
         return writer.info
 
-    def register_byte_stream_handler(self, topic: str, handler: ByteStreamHandler):
-        existing_handler = self._byte_stream_handlers.get(topic)
-        if existing_handler is None:
-            self._byte_stream_handlers[topic] = handler
-        else:
-            raise ValueError("byte stream handler for topic '%s' already set" % topic)
-
-    def unregister_byte_stream_handler(self, topic: str):
-        if self._byte_stream_handlers.get(topic):
-            self._byte_stream_handlers.pop(topic)
-
-    def register_text_stream_handler(self, topic: str, handler: TextStreamHandler):
-        existing_handler = self._text_stream_handlers.get(topic)
-        if existing_handler is None:
-            self._text_stream_handlers[topic] = handler
-        else:
-            raise ValueError("text stream handler for topic '%s' already set" % topic)
-
-    def unregister_text_stream_handler(self, topic: str):
-        if self._text_stream_handlers.get(topic):
-            self._text_stream_handlers.pop(topic)
-
     async def publish_track(
         self, track: LocalTrack, options: TrackPublishOptions = TrackPublishOptions()
     ) -> LocalTrackPublication:

livekit-rtc/livekit/rtc/room.py

@@ -33,6 +33,12 @@
 from .track import RemoteAudioTrack, RemoteVideoTrack
 from .track_publication import RemoteTrackPublication, TrackPublication
 from .transcription import TranscriptionSegment
+from .data_stream import (
+    TextStreamReader,
+    ByteStreamReader,
+    TextStreamHandler,
+    ByteStreamHandler,
+)
 
 
 EventTypes = Literal[
@@ -140,6 +146,11 @@ def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
         self._first_sid_future = asyncio.Future[str]()
         self._local_participant: LocalParticipant | None = None
 
+        self._text_stream_readers: Dict[str, TextStreamReader] = {}
+        self._byte_stream_readers: Dict[str, ByteStreamReader] = {}
+        self._text_stream_handlers: Dict[str, TextStreamHandler] = {}
+        self._byte_stream_handlers: Dict[str, ByteStreamHandler] = {}
+
     def __del__(self) -> None:
         if self._ffi_handle is not None:
             FfiClient.instance.queue.unsubscribe(self._ffi_queue)
@@ -398,6 +409,28 @@ def on_participant_connected(participant):
         # start listening to room events
         self._task = self._loop.create_task(self._listen_task())
 
+    def register_byte_stream_handler(self, topic: str, handler: ByteStreamHandler):
+        existing_handler = self._byte_stream_handlers.get(topic)
+        if existing_handler is None:
+            self._byte_stream_handlers[topic] = handler
+        else:
+            raise ValueError("byte stream handler for topic '%s' already set" % topic)
+
+    def unregister_byte_stream_handler(self, topic: str):
+        if self._byte_stream_handlers.get(topic):
+            self._byte_stream_handlers.pop(topic)
+
+    def register_text_stream_handler(self, topic: str, handler: TextStreamHandler):
+        existing_handler = self._text_stream_handlers.get(topic)
+        if existing_handler is None:
+            self._text_stream_handlers[topic] = handler
+        else:
+            raise ValueError("text stream handler for topic '%s' already set" % topic)
+
+    def unregister_text_stream_handler(self, topic: str):
+        if self._text_stream_handlers.get(topic):
+            self._text_stream_handlers.pop(topic)
+
     async def disconnect(self) -> None:
         """Disconnects from the room."""
         if not self.isconnected():
@@ -714,28 +747,77 @@ def _on_room_event(self, event: proto_room.RoomEvent):
         elif which == "reconnected":
             self.emit("reconnected")
         elif which == "stream_header_received":
-            self.local_participant._handle_stream_header(
+            self._handle_stream_header(
                 event.stream_header_received.header,
                 event.stream_header_received.participant_identity,
             )
         elif which == "stream_chunk_received":
             task = asyncio.create_task(
-                self.local_participant._handle_stream_chunk(
-                    event.stream_chunk_received.chunk
-                )
+                self._handle_stream_chunk(event.stream_chunk_received.chunk)
             )
             self._data_stream_tasks.add(task)
             task.add_done_callback(self._data_stream_tasks.discard)
 
         elif which == "stream_trailer_received":
             task = asyncio.create_task(
-                self.local_participant._handle_stream_trailer(
-                    event.stream_trailer_received.trailer
-                )
+                self._handle_stream_trailer(event.stream_trailer_received.trailer)
             )
             self._data_stream_tasks.add(task)
             task.add_done_callback(self._data_stream_tasks.discard)
 
+    def _handle_stream_header(
+        self, header: proto_room.DataStream.Header, participant_identity: str
+    ):
+        stream_type = header.WhichOneof("content_header")
+        if stream_type == "text_header":
+            text_stream_handler = self._text_stream_handlers.get(header.topic)
+            if text_stream_handler is None:
+                logging.info(
+                    "ignoring text stream with topic '%s', no callback attached",
+                    header.topic,
+                )
+                return
+
+            text_reader = TextStreamReader(header)
+            self._text_stream_readers[header.stream_id] = text_reader
+            text_stream_handler(text_reader, participant_identity)
+        elif stream_type == "byte_header":
+            logging.warning("received byte header, %s", header.stream_id)
+            byte_stream_handler = self._byte_stream_handlers.get(header.topic)
+            if byte_stream_handler is None:
+                logging.info(
+                    "ignoring byte stream with topic '%s', no callback attached",
+                    header.topic,
+                )
+                return
+
+            byte_reader = ByteStreamReader(header)
+            self._byte_stream_readers[header.stream_id] = byte_reader
+            byte_stream_handler(byte_reader, participant_identity)
+        else:
+            logging.warning("received unknown header type, %s", stream_type)
+        pass
+
+    async def _handle_stream_chunk(self, chunk: proto_room.DataStream.Chunk):
+        text_reader = self._text_stream_readers.get(chunk.stream_id)
+        file_reader = self._byte_stream_readers.get(chunk.stream_id)
+
+        if text_reader:
+            await text_reader._on_chunk_update(chunk)
+        elif file_reader:
+            await file_reader._on_chunk_update(chunk)
+
+    async def _handle_stream_trailer(self, trailer: proto_room.DataStream.Trailer):
+        text_reader = self._text_stream_readers.get(trailer.stream_id)
+        file_reader = self._byte_stream_readers.get(trailer.stream_id)
+
+        if text_reader:
+            await text_reader._on_stream_close(trailer)
+            self._text_stream_readers.pop(trailer.stream_id)
+        elif file_reader:
+            await file_reader._on_stream_close(trailer)
+            self._byte_stream_readers.pop(trailer.stream_id)
+
     async def _drain_rpc_invocation_tasks(self) -> None:
         if self._rpc_invocation_tasks:
             for task in self._rpc_invocation_tasks: