""" LiveKit Stream Manager for real-time video streaming. Uses LiveKit Python SDK for WebRTC video publishing. """ import asyncio import logging import os from typing import Dict, Optional import numpy as np import cv2 from livekit import rtc, api logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # LiveKit server configuration LIVEKIT_URL = os.getenv("LIVEKIT_URL", "ws://localhost:7880") LIVEKIT_API_KEY = os.getenv("LIVEKIT_API_KEY", "devkey") LIVEKIT_API_SECRET = os.getenv("LIVEKIT_API_SECRET", "secret") class LiveKitVideoPublisher: """Publishes video frames to a LiveKit room.""" def __init__(self, room_name: str, participant_identity: str, fps: int = 30): self.room_name = room_name self.participant_identity = participant_identity self.fps = fps self.room: Optional[rtc.Room] = None self.video_source: Optional[rtc.VideoSource] = None self.video_track: Optional[rtc.LocalVideoTrack] = None self._connected = False self._width = 256 self._height = 256 async def connect(self) -> bool: """Connect to LiveKit room and publish video track.""" try: # Generate access token token = api.AccessToken(LIVEKIT_API_KEY, LIVEKIT_API_SECRET) token.with_identity(self.participant_identity) token.with_name("Avatar Bot") token.with_grants(api.VideoGrants( room_join=True, room=self.room_name, can_publish=True, can_subscribe=False, )) jwt_token = token.to_jwt() # Create and connect to room self.room = rtc.Room() @self.room.on("connected") def on_connected(): logger.info(f"[LiveKit] Connected to room: {self.room_name}") self._connected = True @self.room.on("disconnected") def on_disconnected(): logger.info(f"[LiveKit] Disconnected from room: {self.room_name}") self._connected = False await self.room.connect(LIVEKIT_URL, jwt_token) # Create video source and track self.video_source = rtc.VideoSource(self._width, self._height) self.video_track = rtc.LocalVideoTrack.create_video_track( "avatar-video", self.video_source ) # Publish the track options = rtc.TrackPublishOptions() options.source = rtc.TrackSource.SOURCE_CAMERA await self.room.local_participant.publish_track(self.video_track, options) logger.info(f"[LiveKit] Published video track to room: {self.room_name}") self._connected = True return True except Exception as e: logger.error(f"[LiveKit] Connection error: {e}") self._connected = False return False async def send_frame(self, frame: np.ndarray): """Send a video frame to the room.""" if not self._connected or self.video_source is None: return try: # Resize frame if needed h, w = frame.shape[:2] if w != self._width or h != self._height: frame = cv2.resize(frame, (self._width, self._height)) # Convert BGR to RGBA if frame.shape[2] == 3: frame_rgba = cv2.cvtColor(frame, cv2.COLOR_BGR2RGBA) else: frame_rgba = frame # Create VideoFrame and capture video_frame = rtc.VideoFrame( self._width, self._height, rtc.VideoBufferType.RGBA, frame_rgba.tobytes() ) self.video_source.capture_frame(video_frame) except Exception as e: logger.error(f"[LiveKit] Error sending frame: {e}") def is_connected(self) -> bool: """Check if connected to room.""" return self._connected async def disconnect(self): """Disconnect from the room.""" if self.room: await self.room.disconnect() self._connected = False logger.info(f"[LiveKit] Disconnected from room: {self.room_name}") class LiveKitManager: """Manages LiveKit rooms and video publishers.""" def __init__(self): self.publishers: Dict[str, LiveKitVideoPublisher] = {} self._api: Optional[api.LiveKitAPI] = None def _get_api(self) -> api.LiveKitAPI: """Get or create LiveKit API client.""" if self._api is None: self._api = api.LiveKitAPI( LIVEKIT_URL.replace("ws://", "http://").replace("wss://", "https://"), LIVEKIT_API_KEY, LIVEKIT_API_SECRET ) return self._api def generate_viewer_token(self, room_name: str, viewer_identity: str) -> str: """Generate a token for a viewer to join the room.""" token = api.AccessToken(LIVEKIT_API_KEY, LIVEKIT_API_SECRET) token.with_identity(viewer_identity) token.with_name(f"Viewer {viewer_identity}") token.with_grants(api.VideoGrants( room_join=True, room=room_name, can_publish=False, can_subscribe=True, )) return token.to_jwt() async def create_room(self, room_name: str) -> bool: """Create a LiveKit room.""" try: lk_api = self._get_api() await lk_api.room.create_room( api.CreateRoomRequest(name=room_name) ) logger.info(f"[LiveKit] Created room: {room_name}") return True except Exception as e: logger.error(f"[LiveKit] Error creating room: {e}") return False async def get_or_create_publisher( self, session_id: str, fps: int = 30 ) -> LiveKitVideoPublisher: """Get existing publisher or create a new one.""" if session_id in self.publishers: return self.publishers[session_id] room_name = f"avatar-{session_id}" publisher = LiveKitVideoPublisher( room_name=room_name, participant_identity=f"avatar-bot-{session_id}", fps=fps ) # Create room and connect publisher await self.create_room(room_name) await publisher.connect() self.publishers[session_id] = publisher return publisher async def handle_offer( self, session_id: str, sdp: str, type: str, fps: int = 30 ) -> dict: """ Handle WebRTC offer - with LiveKit, we don't use SDP directly. Instead, return room info for the client to connect. """ room_name = f"avatar-{session_id}" # Create publisher if not exists publisher = await self.get_or_create_publisher(session_id, fps) # Generate viewer token for client viewer_token = self.generate_viewer_token(room_name, f"viewer-{session_id}") return { "type": "livekit", "room_name": room_name, "token": viewer_token, "url": LIVEKIT_URL, "sdp": "" # Not used with LiveKit } def is_connected(self, session_id: str) -> bool: """Check if a session has an active publisher.""" publisher = self.publishers.get(session_id) return publisher.is_connected() if publisher else False async def wait_for_connection(self, session_id: str, timeout: float = 5.0) -> bool: """Wait for publisher to be connected.""" publisher = self.publishers.get(session_id) if not publisher: return False start = asyncio.get_event_loop().time() while not publisher.is_connected(): if asyncio.get_event_loop().time() - start > timeout: return False await asyncio.sleep(0.1) return True async def send_frame(self, session_id: str, frame: np.ndarray): """Send a video frame to a session.""" publisher = self.publishers.get(session_id) if publisher and publisher.is_connected(): await publisher.send_frame(frame) async def close_connection(self, session_id: str): """Close a session's connection.""" publisher = self.publishers.pop(session_id, None) if publisher: await publisher.disconnect() logger.info(f"[LiveKit] Closed session: {session_id}") # Global manager instance livekit_manager = LiveKitManager() # Alias for compatibility with existing code webrtc_manager = livekit_manager