vad_mode.py

import base64
import json
import ssl
import time
import uuid
from threading import Lock, Thread

import pyaudio
from loguru import logger
from websockets.sync.client import ClientConnection, connect

from client_events import (
    InputAudioBufferAppend,
    Session,
    SessionUpdate,
    TurnDetection,
    TurnDetectionType,
)
from server_events import ServerEventType

FORMAT = pyaudio.paInt16
CHANNELS = 1
SAMPLE_RATE = 24000
CHUNK_FRAMES = 480
WS_ENDPOINT = "wss://api.stepfun.com/v1/realtime"
MODEL_NAME = "step-1o-audio"
TRACE_ID = str(uuid.uuid4())
API_KEY = "replace with your api key"

# init local input/output aduio
p = pyaudio.PyAudio()
input = p.open(
    format=FORMAT,
    channels=CHANNELS,
    rate=SAMPLE_RATE,
    input=True,
    frames_per_buffer=CHUNK_FRAMES,
)

output = p.open(
    format=FORMAT,
    channels=CHANNELS,
    rate=SAMPLE_RATE,
    output=True,
)


# init audio buffer
# interrupt audio by replacing buffer
# if you want to more precisely, implement AudioBuffer in frame level
class AudioBuffer:

    def __init__(self):
        self._buffer = []
        self._lock = Lock()

    def read(self) -> bytes:
        with self._lock:
            if len(self._buffer):
                return self._buffer.pop(0)
            return b""

    def append(self, data: bytes):
        with self._lock:
            self._buffer.append(data)


buffer = AudioBuffer()


# init websocket
ws: ClientConnection = connect(
    f"{WS_ENDPOINT}?model={MODEL_NAME}",
    ssl=ssl.create_default_context(),
    additional_headers={
        "X-Trace-Id": TRACE_ID,
        "Authorization": f"Bearer {API_KEY}",
    },
)

logger.info(f"connected! trace id: {TRACE_ID}")


def run_sender():
    # init session with vad mode
    ws.send(
        SessionUpdate(
            session=Session(
                turn_detection=TurnDetection(type=TurnDetectionType.ServerVAD)
            )
        ).model_dump_json(exclude_none=True)
    )

    # Keep sending audio
    while True:
        # send small size audio each time to make stream smooth
        data = input.read(CHUNK_FRAMES, exception_on_overflow=False)
        ws.send(
            InputAudioBufferAppend(
                audio=base64.b64encode(data).decode("utf-8")
            ).model_dump_json(exclude_none=True)
        )


sender_thread = Thread(target=run_sender)
sender_thread.start()

audio_playing = False


def play_audio():
    global audio_playing
    while True:
        data = buffer.read()
        if data:
            audio_playing = True
            output.write(data)
        else:
            audio_playing = False
            time.sleep(0.02)


play_thread = Thread(target=play_audio)
play_thread.start()


first_audio_received = False

for msg in ws:
    server_event = json.loads(msg)
    match server_event["type"]:
        case ServerEventType.ResponseAudioDelta:
            if not first_audio_received:
                first_audio_received = True
                logger.info("audio started!")
            logger.info(
                f"received audio delta {len(server_event['delta'])} bytes, event id: {server_event['event_id']}"
            )
            buffer.append(base64.b64decode(server_event["delta"]))
        case ServerEventType.ResponseDone:
            first_audio_received = False
            logger.info("response done!")
        case ServerEventType.InputAudioBufferSpeechStarted:
            if audio_playing:
                logger.info("interrupt detected! resetting audio buffer")
                buffer = AudioBuffer()
        case _:
            logger.info(f"received event {server_event}")