from __future__ import annotations import asyncio import json import time from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from typing import TypeVar from fastapi import FastAPI, WebSocket, WebSocketDisconnect from pydantic import ValidationError from .models import ASRControl, ASRStart, error_event, final_event, partial_event from .providers import ASRProvider, ProviderError, Transcript, build_provider from .settings import ASRSettings T = TypeVar("T") class CapacityExceeded(RuntimeError): pass class _ConnectionGate: def __init__(self, limit: int) -> None: self.limit = limit self.active = 0 self._lock = asyncio.Lock() async def try_enter(self) -> bool: async with self._lock: if self.active >= self.limit: return False self.active += 1 return True async def leave(self) -> None: async with self._lock: self.active = max(0, self.active - 1) class _ProviderGate: def __init__(self, limit: int, queue_timeout_s: float) -> None: self._semaphore = asyncio.Semaphore(limit) self._queue_timeout_s = queue_timeout_s async def run(self, factory: Callable[[], Awaitable[T]], timeout_s: float) -> T: try: await asyncio.wait_for(self._semaphore.acquire(), self._queue_timeout_s) except TimeoutError as exc: raise CapacityExceeded("ASR provider is busy") from exc try: return await asyncio.wait_for(factory(), timeout_s) finally: self._semaphore.release() def create_app( settings: ASRSettings | None = None, *, provider: ASRProvider | None = None, ) -> FastAPI: settings = (settings or ASRSettings.from_env()).validated() provider = provider or build_provider(settings) connections = _ConnectionGate(settings.max_connections) provider_gate = _ProviderGate( settings.max_provider_concurrency, settings.provider_queue_timeout_s ) @asynccontextmanager async def lifespan(_: FastAPI): yield await provider.close() app = FastAPI( title="OMNI local streaming ASR", version="1.0", docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan, ) app.state.settings = settings app.state.provider = provider app.state.connections = connections @app.get("/healthz") async def health() -> dict[str, object]: return { "status": "ready", "provider": provider.name, "streaming_partials": provider.supports_partials, "model_downloads": "disabled", "active_connections": connections.active, "max_connections": connections.limit, } async def send_error(websocket: WebSocket, detail: str) -> None: try: await websocket.send_json(error_event(detail).model_dump(exclude_none=True)) except (RuntimeError, WebSocketDisconnect): pass @app.websocket("/v1/stream") async def stream(websocket: WebSocket) -> None: if not await connections.try_enter(): await websocket.close(code=1013, reason="ASR connection capacity reached") return accepted = False try: await websocket.accept() accepted = True try: first = await asyncio.wait_for( websocket.receive(), timeout=settings.start_timeout_s ) except TimeoutError: await send_error(websocket, "ASR_START_TIMEOUT: input.start was not received") await websocket.close(code=1008, reason="input.start timeout") return if first.get("type") == "websocket.disconnect": return if first.get("text") is None: await send_error(websocket, "ASR_PROTOCOL: first frame must be JSON input.start") await websocket.close(code=1008, reason="input.start required") return try: start = ASRStart.model_validate_json(first["text"]) except (ValidationError, ValueError): await send_error(websocket, "ASR_PROTOCOL: invalid input.start frame") await websocket.close(code=1008, reason="invalid input.start") return frame_width = 2 * start.channels duration_limit = int( start.sample_rate_hz * frame_width * settings.max_audio_seconds ) audio_limit = min(settings.max_audio_bytes, duration_limit) partial_stride = max( frame_width, int( start.sample_rate_hz * frame_width * settings.partial_interval_ms / 1000 ), ) audio = bytearray() last_partial_at = 0 last_partial_text = "" session_started = time.monotonic() while True: remaining = settings.session_timeout_s - (time.monotonic() - session_started) if remaining <= 0: await send_error(websocket, "ASR_SESSION_TIMEOUT: session exceeded its limit") return try: message = await asyncio.wait_for( websocket.receive(), timeout=min(settings.idle_timeout_s, remaining) ) except TimeoutError: await send_error(websocket, "ASR_IDLE_TIMEOUT: no audio or control frame") return if message.get("type") == "websocket.disconnect": return chunk = message.get("bytes") if chunk is not None: if not chunk: continue if len(chunk) > settings.max_frame_bytes: await send_error(websocket, "ASR_FRAME_TOO_LARGE: audio frame rejected") return if len(chunk) % frame_width: await send_error(websocket, "ASR_PCM_ALIGNMENT: incomplete PCM sample frame") return if len(audio) + len(chunk) > audio_limit: await send_error(websocket, "ASR_BUFFER_FULL: audio limit reached") return audio.extend(chunk) if ( provider.supports_partials and len(audio) - last_partial_at >= partial_stride ): last_partial_at = len(audio) try: transcript = await provider_gate.run( lambda: provider.partial(bytes(audio), start), min(settings.transcribe_timeout_s, 2.0), ) except (CapacityExceeded, ProviderError, TimeoutError): # Partials are opportunistic; final transcription remains available. transcript = None if transcript and transcript.text and transcript.text != last_partial_text: last_partial_text = transcript.text await websocket.send_json( partial_event( transcript.text, language=transcript.language or start.language, confidence=transcript.confidence, ).model_dump(exclude_none=True) ) continue text = message.get("text") if text is None: await send_error(websocket, "ASR_PROTOCOL: unsupported WebSocket frame") return try: control = ASRControl.model_validate(json.loads(text)) except (ValidationError, ValueError, TypeError): await send_error(websocket, "ASR_PROTOCOL: invalid control frame") return if control.type == "input.cancel": audio.clear() return if not audio: await send_error(websocket, "ASR_EMPTY_AUDIO: nothing to transcribe") return try: result: Transcript = await provider_gate.run( lambda: provider.final(bytes(audio), start), settings.transcribe_timeout_s, ) except CapacityExceeded: await send_error(websocket, "ASR_BUSY: provider capacity reached") return except TimeoutError: await send_error(websocket, "ASR_PROVIDER_TIMEOUT: transcription timed out") return except ProviderError as exc: await send_error(websocket, f"ASR_PROVIDER_ERROR: {exc}") return if not result.text.strip(): await send_error(websocket, "ASR_PROVIDER_ERROR: empty transcript") return await websocket.send_json( final_event( result.text.strip(), language=result.language or start.language, confidence=result.confidence, ).model_dump(exclude_none=True) ) return except WebSocketDisconnect: return finally: await connections.leave() if accepted: try: await websocket.close(code=1000) except (RuntimeError, WebSocketDisconnect): pass return app