pepperplus-cb/src/control_backend/agents/vad_agent.py

import logging

import numpy as np
import torch
import zmq
import zmq.asyncio as azmq
from spade.agent import Agent
from spade.behaviour import CyclicBehaviour

from control_backend.agents.transcription import TranscriptionAgent
from control_backend.core.config import settings
from control_backend.core.zmq_context import context as zmq_context

logger = logging.getLogger(__name__)


class SocketPoller[T]:
    """
    Convenience class for polling a socket for data with a timeout, persisting a zmq.Poller for
    multiple usages.
    """

    def __init__(
        self,
        socket: azmq.Socket,
        timeout_ms: int = settings.behaviour_settings.socket_poller_timeout_ms,
    ):
        """
        :param socket: The socket to poll and get data from.
        :param timeout_ms: A timeout in milliseconds to wait for data.
        """
        self.socket = socket
        self.poller = zmq.Poller()
        self.poller.register(self.socket, zmq.POLLIN)
        self.timeout_ms = timeout_ms

    async def poll(self, timeout_ms: int | None = None) -> T | None:
        """
        Get data from the socket, or None if the timeout is reached.

        :param timeout_ms: If given, the timeout. Otherwise, `self.timeout_ms` is used.
        :return: Data from the socket or None.
        """
        timeout_ms = timeout_ms or self.timeout_ms
        socks = dict(self.poller.poll(timeout_ms))
        if socks.get(self.socket) == zmq.POLLIN:
            return await self.socket.recv()
        return None


class Streaming(CyclicBehaviour):
    def __init__(self, audio_in_socket: azmq.Socket, audio_out_socket: azmq.Socket):
        super().__init__()
        self.audio_in_poller = SocketPoller[bytes](audio_in_socket)
        self.model, _ = torch.hub.load(
            repo_or_dir=settings.vad_settings.repo_or_dir,
            model=settings.vad_settings.model_name,
            force_reload=False,
        )
        self.audio_out_socket = audio_out_socket

        self.audio_buffer = np.array([], dtype=np.float32)
        self.i_since_speech = (
            settings.behaviour_settings.vad_initial_since_speech
        )  # Used to allow small pauses in speech

    async def run(self) -> None:
        data = await self.audio_in_poller.poll()
        if data is None:
            if len(self.audio_buffer) > 0:
                logger.debug("No audio data received. Discarding buffer until new data arrives.")
                self.audio_buffer = np.array([], dtype=np.float32)
                self.i_since_speech = settings.behaviour_settings.vad_initial_since_speech
            return

        # copy otherwise Torch will be sad that it's immutable
        chunk = np.frombuffer(data, dtype=np.float32).copy()
        prob = self.model(torch.from_numpy(chunk), settings.vad_settings.sample_rate_hz).item()
        non_speech_patience = settings.behaviour_settings.vad_non_speech_patience_chunks
        prob_threshold = settings.behaviour_settings.vad_prob_threshold

        if prob > prob_threshold:
            if self.i_since_speech > non_speech_patience:
                logger.debug("Speech started.")
            self.audio_buffer = np.append(self.audio_buffer, chunk)
            self.i_since_speech = 0
            return
        self.i_since_speech += 1

        # prob < 0.5, so speech maybe ended. Wait a bit more before to be more certain
        if self.i_since_speech <= non_speech_patience:
            self.audio_buffer = np.append(self.audio_buffer, chunk)
            return

        # Speech probably ended. Make sure we have a usable amount of data.
        if len(self.audio_buffer) >= 3 * len(chunk):
            logger.debug("Speech ended.")
            await self.audio_out_socket.send(self.audio_buffer[: -2 * len(chunk)].tobytes())

        # At this point, we know that the speech has ended.
        # Prepend the last chunk that had no speech, for a more fluent boundary
        self.audio_buffer = chunk


class VADAgent(Agent):
    """
    An agent which listens to an audio stream, does Voice Activity Detection (VAD), and sends
    fragments with detected speech to other agents over ZeroMQ.
    """

    def __init__(self, audio_in_address: str, audio_in_bind: bool):
        jid = settings.agent_settings.vad_agent_name + "@" + settings.agent_settings.host
        super().__init__(jid, settings.agent_settings.vad_agent_name)

        self.audio_in_address = audio_in_address
        self.audio_in_bind = audio_in_bind

        self.audio_in_socket: azmq.Socket | None = None
        self.audio_out_socket: azmq.Socket | None = None

    async def stop(self):
        """
        Stop listening to audio, stop publishing audio, close sockets.
        """
        if self.audio_in_socket is not None:
            self.audio_in_socket.close()
            self.audio_in_socket = None
        if self.audio_out_socket is not None:
            self.audio_out_socket.close()
            self.audio_out_socket = None
        return await super().stop()

    def _connect_audio_in_socket(self):
        self.audio_in_socket = zmq_context.socket(zmq.SUB)
        self.audio_in_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        if self.audio_in_bind:
            self.audio_in_socket.bind(self.audio_in_address)
        else:
            self.audio_in_socket.connect(self.audio_in_address)
        self.audio_in_poller = SocketPoller[bytes](self.audio_in_socket)

    def _connect_audio_out_socket(self) -> int | None:
        """Returns the port bound, or None if binding failed."""
        try:
            self.audio_out_socket = zmq_context.socket(zmq.PUB)
            return self.audio_out_socket.bind_to_random_port("tcp://*", max_tries=100)
        except zmq.ZMQBindError:
            logger.error("Failed to bind an audio output socket after 100 tries.")
            self.audio_out_socket = None
            return None

    async def setup(self):
        logger.info("Setting up %s", self.jid)

        self._connect_audio_in_socket()

        audio_out_port = self._connect_audio_out_socket()
        if audio_out_port is None:
            await self.stop()
            return
        audio_out_address = f"tcp://localhost:{audio_out_port}"

        streaming = Streaming(self.audio_in_socket, self.audio_out_socket)
        self.add_behaviour(streaming)

        # Start agents dependent on the output audio fragments here
        transcriber = TranscriptionAgent(audio_out_address)
        await transcriber.start()

        logger.info("Finished setting up %s", self.jid)