Building a Voice Pipeline by Hand

package com.example.docsassistant.voice;

/** What a voice session needs from the wire. Implementations must be safe to call from any thread. */
public interface AudioTransport {

    /** Send synthesised audio to the caller, already in the transport's format (see {@link #outputFormat()}). */
    void send(byte[] audio);

    /** Drop every audio buffer not yet played (browser: a control message; Twilio: the clear message). */
    void clear();

    /** Hang up or close the socket. */
    void close();

    /** Describes the PCM the session should produce; the adapter converts if the wire differs. */
    AudioFormat outputFormat();

    record AudioFormat(int sampleRateHz, String encoding) { }
}

package com.example.docsassistant.voice;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;
import org.springframework.web.socket.handler.ConcurrentWebSocketSessionDecorator;

public class WebSocketVoiceHandler extends BinaryWebSocketHandler {

    private final VoiceSessionFactory sessions;
    private final Map<String, VoiceSession> bySocket = new ConcurrentHashMap<>();

    public WebSocketVoiceHandler(VoiceSessionFactory sessions) {
        this.sessions = sessions;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession socket) {
        WebSocketSession safe = new ConcurrentWebSocketSessionDecorator(socket, 5_000, 512 * 1024);
        this.bySocket.put(socket.getId(), this.sessions.create(socket.getId(), new Transport(safe)));
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession socket, BinaryMessage message) {
        byte[] frame = new byte[message.getPayloadLength()];
        message.getPayload().get(frame);
        this.bySocket.get(socket.getId()).onInboundAudio(frame);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession socket, CloseStatus status) {
        VoiceSession session = this.bySocket.remove(socket.getId());
        if (session != null) {
            session.close();   // cancels the current turn
        }
    }

    /** Adapter: PCM16 at 16 kHz both ways; the browser client plays binary messages and obeys {"type":"clear"}. */
    private record Transport(WebSocketSession socket) implements AudioTransport {

        @Override
        public void send(byte[] audio) {
            try {
                this.socket.sendMessage(new BinaryMessage(audio));
            } catch (IOException e) {
                throw new IllegalStateException(e);
            }
        }

        @Override
        public void clear() {
            try {
                this.socket.sendMessage(new TextMessage("{\"type\":\"clear\"}"));
            } catch (IOException e) {
                throw new IllegalStateException(e);
            }
        }

        @Override
        public void close() {
            try {
                this.socket.close();
            } catch (IOException ignored) {
                // already closed
            }
        }

        @Override
        public AudioFormat outputFormat() {
            return new AudioFormat(16_000, "pcm_s16le");
        }
    }
}

package com.example.docsassistant.voice;

import java.io.ByteArrayOutputStream;
import java.util.concurrent.atomic.AtomicBoolean;

import org.springframework.ai.audio.transcription.TranscriptionModel;
import org.springframework.ai.audio.tts.StreamingTextToSpeechModel;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.core.io.ByteArrayResource;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

public class VoiceSession {

    private final String conversationId;
    private final AudioTransport transport;
    private final VoiceActivityDetector vad;
    private final TranscriptionModel stt;
    private final StreamingTextToSpeechModel tts;
    private final ChatClient chat;

    private final ByteArrayOutputStream utterance = new ByteArrayOutputStream();
    private final AtomicBoolean speaking = new AtomicBoolean(false);
    private volatile Disposable turn = () -> { };

    public VoiceSession(String conversationId, AudioTransport transport, VoiceActivityDetector vad,
            TranscriptionModel stt, StreamingTextToSpeechModel tts, ChatClient chat) {
        this.conversationId = conversationId;
        this.transport = transport;
        this.vad = vad;
        this.stt = stt;
        this.tts = tts;
        this.chat = chat;
    }

    /** Called by the transport for every inbound PCM16 16 kHz frame (512 samples). */
    public synchronized void onInboundAudio(byte[] frame) {
        switch (this.vad.accept(frame)) {
            case SPEECH_START -> {
                if (this.speaking.get()) {
                    bargeIn();
                }
                this.utterance.reset();
                this.utterance.writeBytes(frame);
            }
            case SPEECH_END -> {
                byte[] pcm = this.utterance.toByteArray();
                this.utterance.reset();
                startTurn(pcm);
            }
            case NONE -> {
                if (this.vad.inSpeech()) {
                    this.utterance.writeBytes(frame);
                }
            }
        }
    }

    private void startTurn(byte[] pcm) {
        this.turn = Mono.fromCallable(() -> this.stt.transcribe(new WavResource(pcm, 16_000)))   // blocking call
                .subscribeOn(Schedulers.boundedElastic())
                .filter(text -> !text.isBlank())
                .flatMapMany(this::answer)
                .doOnSubscribe(s -> this.speaking.set(true))
                .doOnNext(this.transport::send)
                .doFinally(signal -> this.speaking.set(false))
                .subscribe(audio -> { }, error -> this.transport.clear());
    }

    private Flux<byte[]> answer(String question) {
        Flux<String> tokens = this.chat.prompt()
                .user(question)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, this.conversationId))
                .stream()
                .content();
        return Flux.defer(() -> {
            SentenceChunker chunker = new SentenceChunker();
            return tokens
                    .concatMapIterable(chunker::accept)
                    .concatWith(Mono.fromSupplier(chunker::flush).filter(s -> !s.isBlank()))
                    .flatMapSequential(sentence -> this.tts.stream(sentence), 2);
        });
    }

    /** The caller spoke over the agent: stop everything, then empty the wire. */
    private void bargeIn() {
        this.turn.dispose();          // cancels the model stream, the synthesis requests and the sending
        this.transport.clear();       // audio already queued downstream would otherwise keep playing
        this.speaking.set(false);
    }

    public void close() {
        this.turn.dispose();
        this.transport.close();
    }
}

import asyncio
import json

import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from silero_vad import VADIterator, load_silero_vad

from docsassistant.voice.turn import run_turn   # shown below
from docsassistant.voice.stt import transcribe  # faster-whisper in a worker thread, see the STT page

app = FastAPI()
vad_model = load_silero_vad()            # loaded once, shared weights; each call gets its own VADIterator state


class WebSocketTransport:
    """Python twin of AudioTransport: the session code only calls these three coroutines."""

    def __init__(self, websocket: WebSocket) -> None:
        self._ws = websocket

    async def send(self, audio: bytes) -> None:
        await self._ws.send_bytes(audio)

    async def clear(self) -> None:
        await self._ws.send_text(json.dumps({"type": "clear"}))

    async def close(self) -> None:
        await self._ws.close()


@app.websocket("/voice")
async def voice(websocket: WebSocket) -> None:
    await websocket.accept()
    transport = WebSocketTransport(websocket)
    vad = VADIterator(vad_model, threshold=0.5, sampling_rate=16000, min_silence_duration_ms=500)
    thread_id = websocket.query_params.get("conversation", "anonymous")
    utterance = bytearray()
    in_speech = False
    turn: asyncio.Task | None = None

    try:
        while True:
            frame = await websocket.receive_bytes()          # 512 samples of PCM16 = 1024 bytes
            samples = torch.from_numpy(np.frombuffer(frame, dtype=np.int16).astype(np.float32) / 32768.0)
            event = vad(samples)

            if event and "start" in event:
                if turn and not turn.done():                 # barge-in: the agent is still talking or thinking
                    turn.cancel()
                    await asyncio.gather(turn, return_exceptions=True)
                    await transport.clear()
                in_speech = True
                utterance.clear()
            if in_speech:
                utterance.extend(frame)
            if event and "end" in event:
                in_speech = False
                pcm = bytes(utterance)
                utterance.clear()
                turn = asyncio.create_task(run_turn(pcm, thread_id, transport))
    except WebSocketDisconnect:
        pass
    finally:
        if turn and not turn.done():
            turn.cancel()
            await asyncio.gather(turn, return_exceptions=True)
        vad.reset_states()

import asyncio
from collections.abc import AsyncIterator

from docsassistant.agent import agent                    # create_agent(...) with the retriever tool and checkpointer
from docsassistant.voice.sentences import sentences      # the chunker from the architecture page
from docsassistant.voice.stt import transcribe
from docsassistant.voice.tts import synthesize           # async generator of audio bytes for one sentence

_END = object()


async def llm_tokens(text: str, thread_id: str) -> AsyncIterator[str]:
    config = {"configurable": {"thread_id": thread_id}}
    async for part in agent.astream(
        {"messages": [{"role": "user", "content": text}]},
        config=config,
        stream_mode="messages",
        version="v2",
    ):
        if part["type"] == "messages":
            token, metadata = part["data"]
            if metadata.get("langgraph_node") == "model" and isinstance(token.content, str):
                yield token.content


async def run_turn(pcm: bytes, thread_id: str, transport) -> None:
    text = await asyncio.to_thread(transcribe, pcm)          # blocking recogniser off the event loop
    if not text.strip():
        return

    queue: asyncio.Queue = asyncio.Queue(maxsize=2)          # at most two sentences ahead of the voice

    async def produce() -> None:
        async for sentence in sentences(llm_tokens(text, thread_id)):
            await queue.put(sentence)
        await queue.put(_END)

    async def consume() -> None:
        while (sentence := await queue.get()) is not _END:
            async for audio in synthesize(sentence):
                await transport.send(audio)

    async with asyncio.TaskGroup() as group:                 # cancelling run_turn cancels both children
        group.create_task(produce())
        group.create_task(consume())