add websocket stream

Author: chenghao-mou

examples/bargein_server.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import base64
+import json
 import logging
 import os
+import time
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from typing import Any
 
 import numpy as np
 import onnxruntime as ort
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 from pydantic import BaseModel, Field
 
 logging.basicConfig(level=logging.INFO)
@@ -153,7 +155,6 @@ async def detect_bargein(request: BargeinRequest) -> BargeinResponse:
     try:
         # Decode the waveform
         waveform = decode_waveform(request.waveform)
-        logger.info(f"Decoded waveform shape: {waveform.shape}")
 
         # Run inference
         is_bargein = run_inference(waveform, request.threshold, request.min_frames)
@@ -174,6 +175,120 @@ async def detect_bargein(request: BargeinRequest) -> BargeinResponse:
         ) from None
 
 
+@app.websocket("/bargein")
+async def websocket_bargein(websocket: WebSocket) -> None:
+    """
+    WebSocket endpoint for bargein detection.
+
+    Protocol:
+    - Client sends: {"type": "session.create", "settings": {"sample_rate": "16000"}}
+    - Server sends: {"type": "session.created"}
+    - Client sends: {"type": "input_audio", "audio": "<base64>", "sample_rate": 16000, "num_channels": 1, "threshold": 0.95, "min_frames": 2}
+    - Server sends: {"type": "bargein_detected"} (when detected)
+    - Client sends: {"type": "session.finalize"}
+    - Server sends: {"type": "session.finalized"}
+    """
+    await websocket.accept()
+    logger.info("WebSocket connection established")
+
+    if onnx_session is None:
+        await websocket.send_json({"type": "error", "message": "ONNX model not loaded"})
+        await websocket.close()
+        return
+
+    try:
+        # Wait for session.create message
+        while True:
+            try:
+                data = await websocket.receive_text()
+                msg = json.loads(data)
+                msg_type = msg.get("type")
+
+                if msg_type == "session.create":
+                    logger.info("Session created")
+                    await websocket.send_json({"type": "session.created"})
+                    break
+                else:
+                    await websocket.send_json(
+                        {"type": "error", "message": f"Expected session.create, got {msg_type}"}
+                    )
+                    await websocket.close()
+                    return
+            except json.JSONDecodeError as e:
+                await websocket.send_json({"type": "error", "message": f"Invalid JSON: {str(e)}"})
+                await websocket.close()
+                return
+
+        # Process audio frames
+        while True:
+            try:
+                data = await websocket.receive_text()
+                msg = json.loads(data)
+                msg_type = msg.get("type")
+
+                if msg_type == "input_audio":
+                    # Decode and process audio
+
+                    audio_b64 = msg.get("audio")
+                    threshold = msg.get("threshold", 0.95)
+                    min_frames = msg.get("min_frames", 2)
+                    created_at = msg.get("created_at", time.time())
+
+                    if not audio_b64:
+                        await websocket.send_json(
+                            {"type": "error", "message": "Missing audio data"}
+                        )
+                        continue
+
+                    try:
+                        waveform = decode_waveform(audio_b64)
+                        is_bargein = run_inference(waveform, threshold, min_frames)
+                        delta = time.time() - created_at
+
+                        await websocket.send_json(
+                            {"type": "inference_done", "delta": delta, "is_bargein": is_bargein}
+                        )
+
+                        if is_bargein:
+                            logger.info("Bargein detected via WebSocket")
+                            await websocket.send_json({"type": "bargein_detected"})
+
+                    except Exception as e:
+                        logger.error(f"Error processing audio: {e}", exc_info=True)
+                        await websocket.send_json(
+                            {"type": "error", "message": f"Error processing audio: {str(e)}"}
+                        )
+
+                elif msg_type == "session.finalize":
+                    logger.info("Session finalized")
+                    await websocket.send_json({"type": "session.finalized"})
+                    break
+
+                else:
+                    logger.warning(f"Unknown message type: {msg_type}")
+                    await websocket.send_json(
+                        {"type": "error", "message": f"Unknown message type: {msg_type}"}
+                    )
+
+            except json.JSONDecodeError as e:
+                await websocket.send_json({"type": "error", "message": f"Invalid JSON: {str(e)}"})
+                continue
+
+    except WebSocketDisconnect:
+        logger.info("WebSocket disconnected")
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}", exc_info=True)
+        try:
+            await websocket.send_json({"type": "error", "message": f"Internal error: {str(e)}"})
+        except Exception:
+            pass
+    finally:
+        try:
+            await websocket.close()
+        except Exception:
+            pass
+
+
 @app.get("/health")
 async def health() -> dict[str, Any]:
     """Detailed health check endpoint."""

examples/voice_agents/basic_agent.py

@@ -14,9 +14,9 @@
     metrics,
     room_io,
 )
+from livekit.agents.inference.bargein import BargeinDetector
 from livekit.agents.llm import function_tool
 from livekit.plugins import silero
-from livekit.plugins.bargein_detector import BargeinDetector
 from livekit.plugins.turn_detector.multilingual import MultilingualModel
 
 # uncomment to enable Krisp background voice/noise cancellation
@@ -94,7 +94,7 @@ async def entrypoint(ctx: JobContext):
         # See more at https://docs.livekit.io/agents/build/turns
         turn_detection=MultilingualModel(),
         vad=ctx.proc.userdata["vad"],
-        bargein_detector=BargeinDetector(),
+        bargein_detector=BargeinDetector(use_proxy=True),
         # allow the LLM to generate a response while waiting for the end of turn
         # See more at https://docs.livekit.io/agents/build/audio/#preemptive-generation
         preemptive_generation=True,

livekit-agents/livekit/agents/bargein.py

@@ -6,11 +6,16 @@
 from collections.abc import AsyncIterator
 from dataclasses import dataclass
 from enum import Enum, unique
-from typing import Literal, Union
+from typing import Generic, Literal, TypeVar, Union
+
+from pydantic import BaseModel, ConfigDict, Field
 
 from livekit import rtc
 
-from .utils import aio
+from ._exceptions import APIConnectionError, APIError
+from .log import logger
+from .types import APIConnectOptions
+from .utils import aio, log_exceptions
 
 
 @unique
@@ -38,10 +43,25 @@ class BargeinEvent:
     """Time taken to perform the inference, in seconds."""
 
 
-class BargeinDetector(ABC, rtc.EventEmitter[Literal["bargein_detected"]]):
-    def __init__(self) -> None:
+class BargeinError(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    type: Literal["bargein_error"] = "bargein_error"
+    timestamp: float
+    label: str
+    error: Exception = Field(..., exclude=True)
+    recoverable: bool
+
+
+TEvent = TypeVar("TEvent")
+
+
+class BargeinDetector(
+    ABC, rtc.EventEmitter[Union[Literal["bargein_detected", "error"], TEvent]], Generic[TEvent]
+):
+    def __init__(self, *, sample_rate: int) -> None:
         super().__init__()
         self._label = f"{type(self).__module__}.{type(self).__name__}"
+        self._sample_rate = sample_rate
 
     @property
     def model(self) -> str:
@@ -55,6 +75,17 @@ def provider(self) -> str:
     def label(self) -> str:
         return self._label
 
+    def _emit_error(self, api_error: Exception, recoverable: bool) -> None:
+        self.emit(
+            "error",
+            BargeinError(
+                timestamp=time.time(),
+                label=self._label,
+                error=api_error,
+                recoverable=recoverable,
+            ),
+        )
+
     @abstractmethod
     def stream(self) -> BargeinStream: ...
 
@@ -75,7 +106,7 @@ class _OverlapSpeechEndedSentinel:
     class _FlushSentinel:
         pass
 
-    def __init__(self, bargein_detector: BargeinDetector) -> None:
+    def __init__(self, bargein_detector: BargeinDetector, conn_options: APIConnectOptions) -> None:
         self._bargein_detector = bargein_detector
         self._last_activity_time = time.perf_counter()
         self._input_ch = aio.Chan[
@@ -91,9 +122,60 @@ def __init__(self, bargein_detector: BargeinDetector) -> None:
         self._event_ch = aio.Chan[BargeinEvent]()
         self._task = asyncio.create_task(self._main_task())
         self._task.add_done_callback(lambda _: self._event_ch.close())
+        self._num_retries = 0
+        self._conn_options = conn_options
+        self._sample_rate = bargein_detector._sample_rate
+        self._resampler: rtc.AudioResampler | None = None
 
     @abstractmethod
-    async def _main_task(self) -> None: ...
+    async def _run(self) -> None: ...
+
+    @log_exceptions(logger=logger)
+    async def _main_task(self) -> None:
+        max_retries = self._conn_options.max_retry
+
+        while self._num_retries <= max_retries:
+            try:
+                return await self._run()
+            except APIError as e:
+                if max_retries == 0:
+                    self._emit_error(e, recoverable=False)
+                    raise
+                elif self._num_retries == max_retries:
+                    self._emit_error(e, recoverable=False)
+                    raise APIConnectionError(
+                        f"failed to detect bargein after {self._num_retries} attempts",
+                    ) from e
+                else:
+                    self._emit_error(e, recoverable=True)
+
+                    retry_interval = self._conn_options._interval_for_retry(self._num_retries)
+                    logger.warning(
+                        f"failed to detect bargein, retrying in {retry_interval}s",
+                        exc_info=e,
+                        extra={
+                            "bargein_detector": self._bargein_detector._label,
+                            "attempt": self._num_retries,
+                        },
+                    )
+                    await asyncio.sleep(retry_interval)
+
+                self._num_retries += 1
+
+            except Exception as e:
+                self._emit_error(e, recoverable=False)
+                raise
+
+    def _emit_error(self, api_error: Exception, recoverable: bool) -> None:
+        self._bargein_detector.emit(
+            "error",
+            BargeinError(
+                timestamp=time.time(),
+                label=self._bargein_detector._label,
+                error=api_error,
+                recoverable=recoverable,
+            ),
+        )
 
     def start_agent_speech(self) -> None:
         """Mark the start of the agent's speech"""
@@ -130,7 +212,28 @@ def push_frame(
         """Push some audio frame to be analyzed"""
         self._check_input_not_ended()
         self._check_not_closed()
-        self._input_ch.send_nowait(frame)
+
+        if not isinstance(frame, rtc.AudioFrame):
+            self._input_ch.send_nowait(frame)
+            return
+
+        if self._sample_rate != frame.sample_rate:
+            if not self._resampler:
+                self._resampler = rtc.AudioResampler(
+                    input_rate=frame.sample_rate,
+                    output_rate=self._sample_rate,
+                    num_channels=1,
+                    quality=rtc.AudioResamplerQuality.LOW,
+                )
+            elif self._resampler._input_rate != frame.sample_rate:
+                raise ValueError("the sample rate of the input frames must be consistent")
+
+        if self._resampler:
+            frames = self._resampler.push(frame)
+            for frame in frames:
+                self._input_ch.send_nowait(frame)
+        else:
+            self._input_ch.send_nowait(frame)
 
     def flush(self) -> None:
         """Mark the end of the current segment"""