| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import json |
| 6 | from collections.abc import AsyncGenerator |
| 7 | from http import HTTPStatus |
| 8 | from uuid import uuid4 |
| 9 | |
| 10 | import numpy as np |
| 11 | import pybase64 as base64 |
| 12 | from fastapi import WebSocket |
| 13 | from starlette.websockets import WebSocketDisconnect |
| 14 | |
| 15 | from vllm import envs |
| 16 | from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo |
| 17 | from vllm.entrypoints.serve.exception_handling.utils import sanitize_message |
| 18 | from vllm.exceptions import VLLMValidationError |
| 19 | from vllm.logger import init_logger |
| 20 | |
| 21 | from .protocol import ( |
| 22 | ErrorEvent, |
| 23 | InputAudioBufferAppend, |
| 24 | InputAudioBufferCommit, |
| 25 | SessionCreated, |
| 26 | TranscriptionDelta, |
| 27 | TranscriptionDone, |
| 28 | ) |
| 29 | from .serving import OpenAIServingRealtime |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class RealtimeConnection: |
| 35 | """Manages WebSocket lifecycle and state for realtime transcription. |
| 36 | |
| 37 | This class handles: |
| 38 | - WebSocket connection lifecycle (accept, receive, send, close) |
| 39 | - Event routing (session.update, append, commit) |
| 40 | - Audio buffering via asyncio.Queue |
| 41 | - Generation task management |
| 42 | - Error handling and cleanup |
| 43 | """ |
| 44 | |
| 45 | def __init__(self, websocket: WebSocket, serving: OpenAIServingRealtime): |
| 46 | self.websocket = websocket |
| 47 | self.connection_id = f"ws-{uuid4()}" |
| 48 | self.serving = serving |
| 49 | self.audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue() |
| 50 | self.generation_task: asyncio.Task | None = None |
| 51 | |
| 52 | self._is_connected = False |
| 53 | self._is_model_validated = False |
| 54 | |
| 55 | self._max_audio_filesize_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB |
| 56 | |
| 57 | async def handle_connection(self): |
| 58 | """Main connection loop.""" |
| 59 | await self.websocket.accept() |
| 60 | logger.debug("WebSocket connection accepted: %s", self.connection_id) |
| 61 | self._is_connected = True |
| 62 | |
| 63 | # Send session created event |
| 64 | await self.send(SessionCreated()) |
| 65 | |
| 66 | try: |
| 67 | while True: |
| 68 | message = await self.websocket.receive_text() |
| 69 | try: |
| 70 | event = json.loads(message) |
| 71 | await self.handle_event(event) |
| 72 | except json.JSONDecodeError: |
| 73 | await self.send_error("Invalid JSON", "invalid_json") |
| 74 | except Exception as e: |
| 75 | logger.exception("Error handling event: %s", e) |
| 76 | await self.send_error(sanitize_message(str(e)), "processing_error") |
| 77 | except WebSocketDisconnect: |
| 78 | logger.debug("WebSocket disconnected: %s", self.connection_id) |
| 79 | self._is_connected = False |
| 80 | except Exception as e: |
| 81 | logger.exception("Unexpected error in connection: %s", e) |
| 82 | finally: |
| 83 | await self.cleanup() |
| 84 | |
| 85 | def _check_model(self, model: str | None) -> None | ErrorResponse: |
| 86 | if self.serving._is_model_supported(model): |
| 87 | return None |
| 88 | |
| 89 | return self.serving.create_error_response( |
| 90 | message=f"The model `{model}` does not exist.", |
| 91 | err_type="NotFoundError", |
| 92 | status_code=HTTPStatus.NOT_FOUND, |
| 93 | param="model", |
| 94 | ) |
| 95 | |
| 96 | async def handle_event(self, event: dict): |
| 97 | """Route events to handlers. |
| 98 | |
| 99 | Supported event types: |
| 100 | - session.update: Configure model |
| 101 | - input_audio_buffer.append: Add audio chunk to queue |
| 102 | - input_audio_buffer.commit: Start transcription generation |
| 103 | """ |
| 104 | event_type = event.get("type") |
| 105 | if event_type == "session.update": |
| 106 | logger.debug("Session updated: %s", event) |
| 107 | model = event.get("model") |
| 108 | if model is None: |
| 109 | await self.send_error("Missing required field: model", "invalid_event") |
| 110 | return |
| 111 | err = self._check_model(model) |
| 112 | if err is not None: |
| 113 | await self.send_error(err.error.message, "model_not_found") |
| 114 | return |
| 115 | self._is_model_validated = True |
| 116 | elif event_type == "input_audio_buffer.append": |
| 117 | append_event = InputAudioBufferAppend(**event) |
| 118 | try: |
| 119 | audio_bytes = base64.b64decode(append_event.audio) |
| 120 | # Convert PCM16 bytes to float32 numpy array |
| 121 | audio_array = ( |
| 122 | np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) |
| 123 | / 32768.0 |
| 124 | ) |
| 125 | |
| 126 | if len(audio_array) / 1024**2 > self._max_audio_filesize_mb: |
| 127 | raise VLLMValidationError( |
| 128 | "Maximum file size exceeded", |
| 129 | parameter="audio_filesize_mb", |
| 130 | value=len(audio_array) / 1024**2, |
| 131 | ) |
| 132 | if len(audio_array) == 0: |
| 133 | raise VLLMValidationError("Can't process empty audio.") |
| 134 | |
| 135 | # Put audio chunk in queue |
| 136 | self.audio_queue.put_nowait(audio_array) |
| 137 | |
| 138 | except Exception as e: |
| 139 | logger.error("Failed to decode audio: %s", e) |
| 140 | await self.send_error("Invalid audio data", "invalid_audio") |
| 141 | |
| 142 | elif event_type == "input_audio_buffer.commit": |
| 143 | if not self._is_model_validated: |
| 144 | err_msg = ( |
| 145 | "Model not validated. Make sure to validate the" |
| 146 | " model by sending a session.update event." |
| 147 | ) |
| 148 | await self.send_error( |
| 149 | err_msg, |
| 150 | "model_not_validated", |
| 151 | ) |
| 152 | return |
| 153 | |
| 154 | commit_event = InputAudioBufferCommit(**event) |
| 155 | # final signals that the audio is finished |
| 156 | if commit_event.final: |
| 157 | self.audio_queue.put_nowait(None) |
| 158 | else: |
| 159 | await self.start_generation() |
| 160 | else: |
| 161 | await self.send_error(f"Unknown event type: {event_type}", "unknown_event") |
| 162 | |
| 163 | async def audio_stream_generator(self) -> AsyncGenerator[np.ndarray, None]: |
| 164 | """Generator that yields audio chunks from the queue.""" |
| 165 | while True: |
| 166 | audio_chunk = await self.audio_queue.get() |
| 167 | if audio_chunk is None: # Sentinel value to stop |
| 168 | break |
| 169 | yield audio_chunk |
| 170 | |
| 171 | async def start_generation(self): |
| 172 | """Start the transcription generation task.""" |
| 173 | if self.generation_task is not None and not self.generation_task.done(): |
| 174 | logger.warning("Generation already in progress, ignoring commit") |
| 175 | return |
| 176 | |
| 177 | # Create audio stream generator |
| 178 | audio_stream = self.audio_stream_generator() |
| 179 | input_stream = asyncio.Queue[list[int]]() |
| 180 | |
| 181 | # Transform to StreamingInput generator |
| 182 | streaming_input_gen = self.serving.transcribe_realtime( |
| 183 | audio_stream, input_stream |
| 184 | ) |
| 185 | |
| 186 | # Start generation task |
| 187 | self.generation_task = asyncio.create_task( |
| 188 | self._run_generation(streaming_input_gen, input_stream) |
| 189 | ) |
| 190 | |
| 191 | async def _run_generation( |
| 192 | self, |
| 193 | streaming_input_gen: AsyncGenerator, |
| 194 | input_stream: asyncio.Queue[list[int]], |
| 195 | ): |
| 196 | """Run the generation and stream results back to the client. |
| 197 | |
| 198 | This method: |
| 199 | 1. Creates sampling parameters from session config |
| 200 | 2. Passes the streaming input generator to engine.generate() |
| 201 | 3. Streams transcription.delta events as text is generated |
| 202 | 4. Sends final transcription.done event with usage stats |
| 203 | 5. Feeds generated token IDs back to input_stream for next iteration |
| 204 | 6. Cleans up the audio queue |
| 205 | """ |
| 206 | request_id = f"rt-{self.connection_id}-{uuid4()}" |
| 207 | full_text = "" |
| 208 | |
| 209 | prompt_token_ids_len: int = 0 |
| 210 | completion_tokens_len: int = 0 |
| 211 | |
| 212 | try: |
| 213 | # Create sampling params |
| 214 | from vllm.sampling_params import RequestOutputKind, SamplingParams |
| 215 | |
| 216 | sampling_params = SamplingParams.from_optional( |
| 217 | temperature=0.0, |
| 218 | max_tokens=self.serving.model_cls.realtime_max_tokens, |
| 219 | output_kind=RequestOutputKind.DELTA, |
| 220 | skip_clone=True, |
| 221 | ) |
| 222 | |
| 223 | # Pass the streaming input generator to the engine |
| 224 | # The engine will consume audio chunks as they arrive and |
| 225 | # stream back transcription results incrementally |
| 226 | result_gen = self.serving.engine_client.generate( |
| 227 | prompt=streaming_input_gen, |
| 228 | sampling_params=sampling_params, |
| 229 | request_id=request_id, |
| 230 | ) |
| 231 | |
| 232 | # Stream results back to client as they're generated |
| 233 | async for output in result_gen: |
| 234 | if output.outputs and len(output.outputs) > 0: |
| 235 | if not prompt_token_ids_len and output.prompt_token_ids: |
| 236 | prompt_token_ids_len = len(output.prompt_token_ids) |
| 237 | |
| 238 | delta = output.outputs[0].text |
| 239 | full_text += delta |
| 240 | |
| 241 | # append output to input |
| 242 | input_stream.put_nowait(list(output.outputs[0].token_ids)) |
| 243 | await self.send(TranscriptionDelta(delta=delta)) |
| 244 | |
| 245 | completion_tokens_len += len(output.outputs[0].token_ids) |
| 246 | |
| 247 | if not self._is_connected: |
| 248 | # finish because websocket connection was killed |
| 249 | break |
| 250 | |
| 251 | usage = UsageInfo( |
| 252 | prompt_tokens=prompt_token_ids_len, |
| 253 | completion_tokens=completion_tokens_len, |
| 254 | total_tokens=prompt_token_ids_len + completion_tokens_len, |
| 255 | ) |
| 256 | |
| 257 | # Send final completion event |
| 258 | await self.send(TranscriptionDone(text=full_text, usage=usage)) |
| 259 | |
| 260 | # Clear queue for next utterance |
| 261 | while not self.audio_queue.empty(): |
| 262 | self.audio_queue.get_nowait() |
| 263 | |
| 264 | except Exception as e: |
| 265 | logger.exception("Error in generation: %s", e) |
| 266 | await self.send_error(sanitize_message(str(e)), "processing_error") |
| 267 | |
| 268 | async def send( |
| 269 | self, event: SessionCreated | TranscriptionDelta | TranscriptionDone |
| 270 | ): |
| 271 | """Send event to client.""" |
| 272 | data = event.model_dump_json() |
| 273 | await self.websocket.send_text(data) |
| 274 | |
| 275 | async def send_error(self, message: str, code: str | None = None): |
| 276 | """Send error event to client.""" |
| 277 | error_event = ErrorEvent(error=message, code=code) |
| 278 | await self.websocket.send_text(error_event.model_dump_json()) |
| 279 | |
| 280 | async def cleanup(self): |
| 281 | """Cleanup resources.""" |
| 282 | # Signal audio stream to stop |
| 283 | self.audio_queue.put_nowait(None) |
| 284 | |
| 285 | # Cancel generation task if running |
| 286 | if self.generation_task and not self.generation_task.done(): |
| 287 | self.generation_task.cancel() |
| 288 | |
| 289 | logger.debug("Connection cleanup complete: %s", self.connection_id) |