Merge pull request #2648 from pipecat-ai/filipi/pcc_small_webrtc

Creating SmallWebRTCRequestHandler for managing peer connections.

src/pipecat/runner/run.py

@@ -70,7 +70,6 @@ import asyncio
 import os
 import sys
 from contextlib import asynccontextmanager
-from typing import Dict
 
 from loguru import logger
 
@@ -79,6 +78,10 @@ from pipecat.runner.types import (
     SmallWebRTCRunnerArguments,
     WebSocketRunnerArguments,
 )
+from pipecat.transports.smallwebrtc.request_handler import (
+    SmallWebRTCRequest,
+    SmallWebRTCRequestHandler,
+)
 
 try:
     import uvicorn
@@ -187,9 +190,6 @@ def _setup_webrtc_routes(app: FastAPI, esp32_mode: bool = False, host: str = "lo
         logger.error(f"WebRTC transport dependencies not installed: {e}")
         return
 
-    # Store connections by pc_id
-    pcs_map: Dict[str, SmallWebRTCConnection] = {}
-
     # Mount the frontend
     app.mount("/client", SmallWebRTCPrebuiltUI)
 
@@ -198,51 +198,33 @@ def _setup_webrtc_routes(app: FastAPI, esp32_mode: bool = False, host: str = "lo
         """Redirect root requests to client interface."""
         return RedirectResponse(url="/client/")
 
+    # Initialize the SmallWebRTC request handler
+    small_webrtc_handler: SmallWebRTCRequestHandler = SmallWebRTCRequestHandler(
+        esp32_mode=esp32_mode, host=host
+    )
+
     @app.post("/api/offer")
-    async def offer(request: dict, background_tasks: BackgroundTasks):
+    async def offer(request: SmallWebRTCRequest, background_tasks: BackgroundTasks):
-        """Handle WebRTC offer requests and manage peer connections."""
+        """Handle WebRTC offer requests via SmallWebRTCRequestHandler."""
-        pc_id = request.get("pc_id")
-
-        if pc_id and pc_id in pcs_map:
-            pipecat_connection = pcs_map[pc_id]
-            logger.info(f"Reusing existing connection for pc_id: {pc_id}")
-            await pipecat_connection.renegotiate(
-                sdp=request["sdp"],
-                type=request["type"],
-                restart_pc=request.get("restart_pc", False),
-            )
-        else:
-            pipecat_connection = SmallWebRTCConnection()
-            await pipecat_connection.initialize(sdp=request["sdp"], type=request["type"])
-
-            @pipecat_connection.event_handler("closed")
-            async def handle_disconnected(webrtc_connection: SmallWebRTCConnection):
-                """Handle WebRTC connection closure and cleanup."""
-                logger.info(f"Discarding peer connection for pc_id: {webrtc_connection.pc_id}")
-                pcs_map.pop(webrtc_connection.pc_id, None)
 
+        # Prepare runner arguments with the callback to run your bot
+        async def webrtc_connection_callback(connection):
             bot_module = _get_bot_module()
-            runner_args = SmallWebRTCRunnerArguments(webrtc_connection=pipecat_connection)
+            runner_args = SmallWebRTCRunnerArguments(webrtc_connection=connection)
             background_tasks.add_task(bot_module.bot, runner_args)
 
-        answer = pipecat_connection.get_answer()
+        # Delegate handling to SmallWebRTCRequestHandler
-
+        answer = await small_webrtc_handler.handle_web_request(
-        # Apply ESP32 SDP munging if enabled
+            request=request,
-        if esp32_mode and host != "localhost":
+            webrtc_connection_callback=webrtc_connection_callback,
-            from pipecat.runner.utils import smallwebrtc_sdp_munging
+        )
-
-            answer["sdp"] = smallwebrtc_sdp_munging(answer["sdp"], host)
-
-        pcs_map[answer["pc_id"]] = pipecat_connection
         return answer
 
     @asynccontextmanager
     async def lifespan(app: FastAPI):
         """Manage FastAPI application lifecycle and cleanup connections."""
         yield
-        coros = [pc.disconnect() for pc in pcs_map.values()]
+        await small_webrtc_handler.close()
-        await asyncio.gather(*coros)
-        pcs_map.clear()
 
     app.router.lifespan_context = lifespan
 

src/pipecat/transports/smallwebrtc/request_handler.py

@@ -0,0 +1,200 @@
+#
+# Copyright (c) 2024–2025, Daily
+#
+# SPDX-License-Identifier: BSD 2-Clause License
+#
+
+"""SmallWebRTC request handler for managing peer connections.
+
+This module provides a client for handling web requests and managing WebRTC connections.
+"""
+
+import asyncio
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from fastapi import HTTPException
+from loguru import logger
+
+from pipecat.transports.smallwebrtc.connection import IceServer, SmallWebRTCConnection
+
+
+@dataclass
+class SmallWebRTCRequest:
+    """Small WebRTC transport session arguments for the runner.
+
+    Parameters:
+        sdp: The SDP string (Session Description Protocol).
+        type: The type of the SDP, either "offer" or "answer".
+        pc_id: Optional identifier for the peer connection.
+        restart_pc: Optional whether to restart the peer connection.
+        request_data: Optional custom data sent by the customer.
+    """
+
+    sdp: str
+    type: str
+    pc_id: Optional[str] = None
+    restart_pc: Optional[bool] = None
+    request_data: Optional[Any] = None
+
+
+class ConnectionMode(Enum):
+    """Enum defining the connection handling modes."""
+
+    SINGLE = "single"  # Only one active connection allowed
+    MULTIPLE = "multiple"  # Multiple simultaneous connections allowed
+
+
+class SmallWebRTCRequestHandler:
+    """SmallWebRTC request handler for managing peer connections.
+
+    This class is responsible for:
+      - Handling incoming SmallWebRTC requests.
+      - Creating and managing WebRTC peer connections.
+      - Supporting ESP32-specific SDP munging if enabled.
+      - Invoking callbacks for newly initialized connections.
+      - Supporting both single and multiple connection modes.
+    """
+
+    def __init__(
+        self,
+        ice_servers: Optional[List[IceServer]] = None,
+        esp32_mode: bool = False,
+        host: Optional[str] = None,
+        connection_mode: ConnectionMode = ConnectionMode.MULTIPLE,
+    ) -> None:
+        """Initialize a SmallWebRTC request handler.
+
+        Args:
+            ice_servers (Optional[List[IceServer]]): List of ICE servers to use for WebRTC
+                connections.
+            esp32_mode (bool): If True, enables ESP32-specific SDP munging.
+            host (Optional[str]): Host address used for SDP munging in ESP32 mode.
+                Ignored if `esp32_mode` is False.
+            connection_mode (ConnectionMode): Mode of operation for handling connections.
+                SINGLE allows only one active connection, MULTIPLE allows several.
+        """
+        self._ice_servers = ice_servers
+        self._esp32_mode = esp32_mode
+        self._host = host
+        self._connection_mode = connection_mode
+
+        # Store connections by pc_id
+        self._pcs_map: Dict[str, SmallWebRTCConnection] = {}
+
+    def _check_single_connection_constraints(self, pc_id: Optional[str]) -> None:
+        """Check if the connection request satisfies single connection mode constraints.
+
+        Args:
+            pc_id: The peer connection ID from the request
+
+        Raises:
+            HTTPException: If constraints are violated in single connection mode
+        """
+        if self._connection_mode != ConnectionMode.SINGLE:
+            return
+
+        if not self._pcs_map:  # No existing connections
+            return
+
+        # Get the existing connection (should be only one in single mode)
+        existing_connection = next(iter(self._pcs_map.values()))
+
+        if existing_connection.pc_id != pc_id and pc_id:
+            logger.warning(
+                f"Connection pc_id mismatch: existing={existing_connection.pc_id}, received={pc_id}"
+            )
+            raise HTTPException(status_code=400, detail="PC ID mismatch with existing connection")
+
+        if not pc_id:
+            logger.warning(
+                "Cannot create new connection: existing connection found but no pc_id received"
+            )
+            raise HTTPException(
+                status_code=400,
+                detail="Cannot create new connection with existing connection active",
+            )
+
+    async def handle_web_request(
+        self,
+        request: SmallWebRTCRequest,
+        webrtc_connection_callback: Callable[[Any], Awaitable[None]],
+    ) -> None:
+        """Handle a SmallWebRTC request and resolve the pending answer.
+
+        This method will:
+          - Reuse an existing WebRTC connection if `pc_id` exists.
+          - Otherwise, create a new `SmallWebRTCConnection`.
+          - Invoke the provided callback with the connection.
+          - Manage ESP32-specific munging if enabled.
+          - Enforce single/multiple connection mode constraints.
+
+        Args:
+            request (SmallWebRTCRequest): The incoming WebRTC request, containing
+                SDP, type, and optionally a `pc_id`.
+            webrtc_connection_callback (Callable[[Any], Awaitable[None]]): An
+                asynchronous callback function that is invoked with the WebRTC connection.
+
+        Raises:
+            HTTPException: If connection mode constraints are violated
+            Exception: Any exception raised during request handling or callback execution
+                will be logged and propagated.
+        """
+        try:
+            pc_id = request.pc_id
+
+            # Check connection mode constraints first
+            self._check_single_connection_constraints(pc_id)
+
+            # After constraints are satisfied, get the existing connection if any
+            existing_connection = self._pcs_map.get(pc_id) if pc_id else None
+
+            if existing_connection:
+                pipecat_connection = existing_connection
+                logger.info(f"Reusing existing connection for pc_id: {pc_id}")
+                await pipecat_connection.renegotiate(
+                    sdp=request.sdp,
+                    type=request.type,
+                    restart_pc=request.restart_pc or False,
+                )
+            else:
+                pipecat_connection = SmallWebRTCConnection(ice_servers=self._ice_servers)
+                await pipecat_connection.initialize(sdp=request.sdp, type=request.type)
+
+                @pipecat_connection.event_handler("closed")
+                async def handle_disconnected(webrtc_connection: SmallWebRTCConnection):
+                    logger.info(f"Discarding peer connection for pc_id: {webrtc_connection.pc_id}")
+                    self._pcs_map.pop(webrtc_connection.pc_id, None)
+
+                # Invoke callback provided in runner arguments
+                try:
+                    await webrtc_connection_callback(pipecat_connection)
+                    logger.debug(
+                        f"webrtc_connection_callback executed successfully for peer: {pipecat_connection.pc_id}"
+                    )
+                except Exception as callback_error:
+                    logger.error(
+                        f"webrtc_connection_callback failed for peer {pipecat_connection.pc_id}: {callback_error}"
+                    )
+
+            answer = pipecat_connection.get_answer()
+
+            if self._esp32_mode and self._host and self._host != "localhost":
+                from pipecat.runner.utils import smallwebrtc_sdp_munging
+
+                answer["sdp"] = smallwebrtc_sdp_munging(answer["sdp"], self._host)
+
+            self._pcs_map[answer["pc_id"]] = pipecat_connection
+
+            return answer
+        except Exception as e:
+            logger.error(f"Error processing SmallWebRTC request: {e}")
+            logger.debug(f"SmallWebRTC request details: {request}")
+            raise
+
+    async def close(self):
+        """Clear the connection map."""
+        coros = [pc.disconnect() for pc in self._pcs_map.values()]
+        await asyncio.gather(*coros)
+        self._pcs_map.clear()