stream handlers and address comments

Author: lukasIO

examples/data_streams.py

@@ -15,15 +15,19 @@ async def main(room: rtc.Room):
 
     async def greetParticipant(identity: str):
         text_writer = await room.local_participant.stream_text(
-            destination_identities=[identity]
+            destination_identities=[identity], topic="chat"
         )
         for char in "Hi! Just a friendly message":
             await text_writer.write(char)
-        await text_writer.close()
+        await text_writer.aclose()
 
-    async def on_text_received(reader: rtc.TextStreamReader):
+    async def on_chat_message_received(
+        reader: rtc.TextStreamReader, participant_identity: str
+    ):
         full_text = await reader.read_all()
-        logger.info(full_text)
+        logger.info(
+            "Received chat message from %s: '%s'", participant_identity, full_text
+        )
 
     @room.on("participant_connected")
     def on_participant_connected(participant: rtc.RemoteParticipant):
@@ -32,26 +36,22 @@ def on_participant_connected(participant: rtc.RemoteParticipant):
         )
         asyncio.create_task(greetParticipant(participant.identity))
 
-    # track_subscribed is emitted whenever the local participant is subscribed to a new track
-    @room.on("text_stream_received")
-    def on_text_stream_received(
-        reader: rtc.TextStreamReader,
-        participant_identity: str,
-    ):
-        logger.info("text stream received from: %s", participant_identity)
-        asyncio.create_task(on_text_received(reader=reader))
+    room.set_text_stream_handler(
+        lambda reader, participant_identity: asyncio.create_task(
+            on_chat_message_received(reader, participant_identity)
+        ),
+        "chat",
+    )
 
     # By default, autosubscribe is enabled. The participant will be subscribed to
     # all published tracks in the room
     await room.connect(URL, TOKEN)
     logger.info("connected to room %s", room.name)
 
     for identity, participant in room.remote_participants.items():
-        print("Sending a welcome message to %s", participant.identity)
+        logger.info("Sending a welcome message to %s", identity)
         await greetParticipant(participant.identity)
 
-    logger.info("exiting")
-
 
 if __name__ == "__main__":
     logging.basicConfig(

livekit-rtc/livekit/rtc/data_stream.py

@@ -17,6 +17,7 @@
 import asyncio
 import uuid
 import datetime
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import AsyncIterator, Optional, TypedDict, Dict, List
 from ._proto.room_pb2 import DataStream as proto_DataStream
@@ -241,7 +242,7 @@ async def _send_trailer(self, trailer: proto_DataStream.Trailer):
         if cb.send_stream_chunk.error:
             raise ConnectionError(cb.send_stream_trailer.error)
 
-    async def close(self):
+    async def aclose(self):
         await self._send_trailer(
             trailer=proto_DataStream.Trailer(
                 stream_id=self._header.stream_id, reason=""
@@ -355,3 +356,7 @@ async def write(self, data: bytes):
     @property
     def info(self) -> ByteStreamInfo:
         return self._info
+
+
+TextStreamHandler = Callable[[TextStreamReader, str], None]
+ByteStreamHandler = Callable[[ByteStreamReader, str], None]

livekit-rtc/livekit/rtc/participant.py

@@ -595,7 +595,7 @@ async def send_text(
 
         for chunk in split_utf8(text, STREAM_CHUNK_SIZE):
             await writer.write(chunk)
-        await writer.close()
+        await writer.aclose()
 
         return writer.info
 
@@ -658,7 +658,7 @@ async def send_file(
         async with aiofiles.open(file_path, "rb") as f:
             while bytes := await f.read(STREAM_CHUNK_SIZE):
                 await writer.write(bytes)
-        await writer.close()
+        await writer.aclose()
 
         return writer.info
 

livekit-rtc/livekit/rtc/room.py

@@ -33,7 +33,12 @@
 from .track import RemoteAudioTrack, RemoteVideoTrack
 from .track_publication import RemoteTrackPublication, TrackPublication
 from .transcription import TranscriptionSegment
-from .data_stream import TextStreamReader, ByteStreamReader
+from .data_stream import (
+    TextStreamReader,
+    ByteStreamReader,
+    TextStreamHandler,
+    ByteStreamHandler,
+)
 
 EventTypes = Literal[
     "participant_connected",
@@ -63,8 +68,6 @@
     "disconnected",
     "reconnecting",
     "reconnected",
-    "text_stream_received",
-    "file_stream_received",
 ]
 
 
@@ -142,6 +145,8 @@ def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
         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:
@@ -421,6 +426,22 @@ async def disconnect(self) -> None:
         await self._task
         FfiClient.instance.queue.unsubscribe(self._ffi_queue)
 
+    def set_byte_stream_handler(
+        self, handler: ByteStreamHandler | None, topic: str = ""
+    ):
+        if handler is None:
+            self._byte_stream_handlers.pop(topic)
+        else:
+            self._byte_stream_handlers[topic] = handler
+
+    def set_text_stream_handler(
+        self, handler: TextStreamHandler | None, topic: str = ""
+    ):
+        if handler is None:
+            self._text_stream_handlers.pop(topic)
+        else:
+            self._text_stream_handlers[topic] = handler
+
     async def _listen_task(self) -> None:
         # listen to incoming room events
         while True:
@@ -757,13 +778,29 @@ def _handle_stream_header(
         self, header: proto_room.DataStream.Header, participant_identity: str
     ):
         if header.text_header:
+            stream_handler = self._text_stream_handlers.get(header.topic)
+            if stream_handler is None:
+                logging.debug(
+                    "ignoring text stream with topic '%s', no callback attached",
+                    header.topic,
+                )
+                return
+
             reader = TextStreamReader(header)
             self._text_stream_readers[header.stream_id] = reader
-            self.emit("text_stream_received", reader, participant_identity)
+            stream_handler(reader, participant_identity)
         elif header.byte_header:
+            stream_handler = self._byte_stream_handlers.get(header.topic)
+            if stream_handler is None:
+                logging.debug(
+                    "ignoring byte stream with topic '%s', no callback attached",
+                    header.topic,
+                )
+                return
+
             reader = ByteStreamReader(header)
             self._byte_stream_readers[header.stream_id] = reader
-            self.emit("file_stream_received", reader, participant_identity)
+            stream_handler(reader, participant_identity)
         pass
 
     def _handle_stream_chunk(self, chunk: proto_room.DataStream.Chunk):