Spaces:

achin42
/

channel

Build error

channel / socket_api.py

Upload folder using huggingface_hub

bc1e6dc verified almost 2 years ago

4.87 kB

from fastapi import APIRouter, WebSocket, Depends, BackgroundTasks, WebSocketDisconnect
from uuid import UUID
from typing import Optional, Dict, Any
import json

from channel.socket_service import SocketService

router = APIRouter(prefix="/ws", tags=["websocket"])

# Singleton instance of SocketService
_socket_service: Optional[SocketService] = None

def get_socket_service() -> SocketService:
    global _socket_service
    if _socket_service is None:
        _socket_service = SocketService()
    return _socket_service

@router.websocket("/{user_id}/{device_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    user_id: UUID,
    device_id: str,
    background_tasks: BackgroundTasks,
    socket_service: SocketService = Depends(get_socket_service)
):
    """
    WebSocket endpoint for real-time communication
    
    Args:
        websocket: The WebSocket connection
        user_id: The ID of the user
        device_id: The ID of the device
    """
    try:
        print(f"New connection request - User: {user_id}, Device: {device_id}")
        # Store background tasks reference
        socket_service.background_tasks = background_tasks
        
        # Start cleanup task if not already running
        background_tasks.add_task(socket_service.cleanup_expired_connections)
        
        # Get initial metadata from client
        metadata = None
        try:
            initial_message = await websocket.receive_json()
            print(f"Initial message received: {initial_message}")
            if isinstance(initial_message, dict):
                metadata = initial_message.get("metadata", {})
        except Exception:
            # If no initial message or invalid, proceed without metadata
            pass
            
        # Connect the websocket
        await socket_service.connect(websocket, user_id, device_id, metadata)
        
        # Handle incoming messages
        while True:
            try:
                message = await websocket.receive_json()
                
                # Handle heartbeat responses
                if isinstance(message, dict) and message.get("type") == "heartbeat_response":
                    await socket_service.update_heartbeat(user_id, device_id)
                    continue
                    
                # For other messages, broadcast to all other devices of this user
                if isinstance(message, dict):
                    # Extract the content from the message
                    if message.get("type") == "message" and "content" in message:
                        # Create a new message with the content
                        broadcast_message = {
                            "type": "message",
                            "content": message["content"]
                        }
                        # Broadcast the content to all other devices of this user
                        broadcast_count = await socket_service.broadcast_message(
                            user_id=user_id,
                            message=broadcast_message,
                            exclude_device_id=device_id  # Exclude the sender's device
                        )
                        
                        # Send acknowledgment to the sender
                        await websocket.send_json({
                            "type": "message_status",
                            "status": "delivered",
                            "broadcast_count": broadcast_count
                        })
                
            except WebSocketDisconnect:
                break
            except Exception as e:
                # Send error message to client
                try:
                    await websocket.send_json({
                        "type": "error",
                        "message": str(e)
                    })
                except Exception:
                    break
                    
    finally:
        # Clean up the connection
        await socket_service.disconnect(user_id, device_id)

@router.post("/{user_id}/broadcast")
async def broadcast_to_user(
    user_id: UUID,
    message: Dict[str, Any],
    sender_device_id: Optional[str] = None,
    socket_service: SocketService = Depends(get_socket_service)
):
    """
    Broadcast a message to all active connections of a user, optionally excluding a sender device
    
    Args:
        user_id: The ID of the user to broadcast to
        message: The message to broadcast
        sender_device_id: Optional device ID to exclude from broadcast
        
    Returns:
        Dict containing broadcast status and count of successful broadcasts
    """
    broadcast_count = await socket_service.broadcast_message(user_id, message, exclude_device_id=sender_device_id)
    return {
        "status": "message_broadcast",
        "broadcast_count": broadcast_count
    } 

	from fastapi import APIRouter, WebSocket, Depends, BackgroundTasks, WebSocketDisconnect
	from uuid import UUID
	from typing import Optional, Dict, Any
	import json

	from channel.socket_service import SocketService

	router = APIRouter(prefix="/ws", tags=["websocket"])

	# Singleton instance of SocketService
	_socket_service: Optional[SocketService] = None

	def get_socket_service() -> SocketService:
	global _socket_service
	if _socket_service is None:
	_socket_service = SocketService()
	return _socket_service

	@router.websocket("/{user_id}/{device_id}")
	async def websocket_endpoint(
	websocket: WebSocket,
	user_id: UUID,
	device_id: str,
	background_tasks: BackgroundTasks,
	socket_service: SocketService = Depends(get_socket_service)
	):
	"""
	WebSocket endpoint for real-time communication

	Args:
	websocket: The WebSocket connection
	user_id: The ID of the user
	device_id: The ID of the device
	"""
	try:
	print(f"New connection request - User: {user_id}, Device: {device_id}")
	# Store background tasks reference
	socket_service.background_tasks = background_tasks

	# Start cleanup task if not already running
	background_tasks.add_task(socket_service.cleanup_expired_connections)

	# Get initial metadata from client
	metadata = None
	try:
	initial_message = await websocket.receive_json()
	print(f"Initial message received: {initial_message}")
	if isinstance(initial_message, dict):
	metadata = initial_message.get("metadata", {})
	except Exception:
	# If no initial message or invalid, proceed without metadata
	pass

	# Connect the websocket
	await socket_service.connect(websocket, user_id, device_id, metadata)

	# Handle incoming messages
	while True:
	try:
	message = await websocket.receive_json()

	# Handle heartbeat responses
	if isinstance(message, dict) and message.get("type") == "heartbeat_response":
	await socket_service.update_heartbeat(user_id, device_id)
	continue

	# For other messages, broadcast to all other devices of this user
	if isinstance(message, dict):
	# Extract the content from the message
	if message.get("type") == "message" and "content" in message:
	# Create a new message with the content
	broadcast_message = {
	"type": "message",
	"content": message["content"]
	}
	# Broadcast the content to all other devices of this user
	broadcast_count = await socket_service.broadcast_message(
	user_id=user_id,
	message=broadcast_message,
	exclude_device_id=device_id # Exclude the sender's device
	)

	# Send acknowledgment to the sender
	await websocket.send_json({
	"type": "message_status",
	"status": "delivered",
	"broadcast_count": broadcast_count
	})

	except WebSocketDisconnect:
	break
	except Exception as e:
	# Send error message to client
	try:
	await websocket.send_json({
	"type": "error",
	"message": str(e)
	})
	except Exception:
	break

	finally:
	# Clean up the connection
	await socket_service.disconnect(user_id, device_id)

	@router.post("/{user_id}/broadcast")
	async def broadcast_to_user(
	user_id: UUID,
	message: Dict[str, Any],
	sender_device_id: Optional[str] = None,
	socket_service: SocketService = Depends(get_socket_service)
	):
	"""
	Broadcast a message to all active connections of a user, optionally excluding a sender device

	Args:
	user_id: The ID of the user to broadcast to
	message: The message to broadcast
	sender_device_id: Optional device ID to exclude from broadcast

	Returns:
	Dict containing broadcast status and count of successful broadcasts
	"""
	broadcast_count = await socket_service.broadcast_message(user_id, message, exclude_device_id=sender_device_id)
	return {
	"status": "message_broadcast",
	"broadcast_count": broadcast_count
	}