support video input from screenshare

.github/next-release/changeset-ad9e5c3a.md

@@ -0,0 +1,5 @@
+---
+"livekit-agents": patch
+---
+
+support video input from screenshare (#2127)

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

@@ -26,10 +26,17 @@ def __init__(
         self,
         room: rtc.Room,
         *,
-        track_source: rtc.TrackSource.ValueType,
+        track_source: rtc.TrackSource.ValueType | list[rtc.TrackSource.ValueType],
     ) -> None:
-        self._room, self._track_source = room, track_source
+        self._room = room
+        self._accepted_sources = (
+            {track_source}
+            if isinstance(track_source, rtc.TrackSource.ValueType)
+            else set(track_source)
+        )
+
         self._data_ch = utils.aio.Chan[T]()
+        self._publication: rtc.RemoteTrackPublication | None = None
         self._stream: rtc.VideoStream | rtc.AudioStream | None = None
         self._participant_identity: str | None = None
         self._attached = True
@@ -38,6 +45,7 @@ def __init__(
         self._tasks: set[asyncio.Task] = set()
 
         self._room.on("track_subscribed", self._on_track_available)
+        self._room.on("track_unpublished", self._on_track_unavailable)
 
     async def __anext__(self) -> T:
         return await self._data_ch.__anext__()
@@ -50,7 +58,7 @@ def on_attached(self) -> None:
             "input stream attached",
             extra={
                 "participant": self._participant_identity,
-                "source": rtc.TrackSource.Name(self._track_source),
+                "sources": [rtc.TrackSource.Name(source) for source in self._accepted_sources],
             },
         )
         self._attached = True
@@ -60,7 +68,7 @@ def on_detached(self) -> None:
             "input stream detached",
             extra={
                 "participant": self._participant_identity,
-                "source": rtc.TrackSource.Name(self._track_source),
+                "sources": [rtc.TrackSource.Name(source) for source in self._accepted_sources],
             },
         )
         self._attached = False
@@ -94,6 +102,7 @@ async def aclose(self) -> None:
         if self._stream:
             await self._stream.aclose()
             self._stream = None
+        self._publication = None
         if self._forward_atask:
             await utils.aio.cancel_and_wait(self._forward_atask)
 
@@ -102,14 +111,17 @@ async def aclose(self) -> None:
 
     @utils.log_exceptions(logger=logger)
     async def _forward_task(
-        self, old_task: asyncio.Task | None, stream: rtc.VideoStream | rtc.AudioStream
+        self,
+        old_task: asyncio.Task | None,
+        stream: rtc.VideoStream | rtc.AudioStream,
+        track_source: rtc.TrackSource.ValueType,
     ) -> None:
         if old_task:
             await utils.aio.cancel_and_wait(old_task)
 
         extra = {
             "participant": self._participant_identity,
-            "source": rtc.TrackSource.Name(self._track_source),
+            "source": rtc.TrackSource.Name(track_source),
         }
         logger.debug("start reading stream", extra=extra)
         async for event in stream:
@@ -129,24 +141,45 @@ def _close_stream(self) -> None:
             task.add_done_callback(self._tasks.discard)
             self._tasks.add(task)
             self._stream = None
+            self._publication = None
 
     def _on_track_available(
         self,
         track: rtc.RemoteTrack,
         publication: rtc.RemoteTrackPublication,
         participant: rtc.RemoteParticipant,
-    ) -> None:
+    ) -> bool:
         if (
             self._participant_identity != participant.identity
-            or publication.source != self._track_source
+            or publication.source not in self._accepted_sources
+            or (self._publication and self._publication.sid == publication.sid)
         ):
-            return
+            return False
 
         self._close_stream()
         self._stream = self._create_stream(track)
+        self._publication = publication
         self._forward_atask = asyncio.create_task(
-            self._forward_task(self._forward_atask, self._stream)
+            self._forward_task(self._forward_atask, self._stream, publication.source)
         )
+        return True
+
+    def _on_track_unavailable(
+        self, publication: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant
+    ) -> None:
+        if (
+            not self._publication
+            or self._publication.sid != publication.sid
+            or participant.identity != self._participant_identity
+        ):
+            return
+
+        self._close_stream()
+
+        # subscribe to the first available track
+        for publication in participant.track_publications.values():
+            if self._on_track_available(publication.track, publication, participant):
+                return
 
 
 class _ParticipantAudioInputStream(_ParticipantInputStream[rtc.AudioFrame], AudioInput):
@@ -178,7 +211,12 @@ def _create_stream(self, track: rtc.Track) -> rtc.AudioStream:
 class _ParticipantVideoInputStream(_ParticipantInputStream[rtc.VideoFrame], VideoInput):
     def __init__(self, room: rtc.Room) -> None:
         _ParticipantInputStream.__init__(
-            self, room=room, track_source=rtc.TrackSource.SOURCE_CAMERA
+            self,
+            room=room,
+            track_source=[
+                rtc.TrackSource.SOURCE_CAMERA,
+                rtc.TrackSource.SOURCE_SCREENSHARE,
+            ],
         )
 
     @override

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

@@ -176,6 +176,7 @@ def _create_transcription_output(
             self._user_tr_output = _create_transcription_output(
                 is_delta_stream=False, participant=self._participant_identity
             )
+            # TODO(long): add next in the chain for session.output.transcription
             self._agent_tr_output = _create_transcription_output(
                 is_delta_stream=True, participant=self._room.local_participant
             )