534 lines
21 KiB
Python
534 lines
21 KiB
Python
"""Backend-neutral transcription and realtime turn abstractions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
from queue import Empty, Full, Queue
|
|
import threading
|
|
import time
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
from .recorder import Capture
|
|
|
|
MAX_PCM_QUEUE_BYTES = 240_000 # five seconds of 24 kHz mono PCM16
|
|
MAX_PCM_QUEUE_CHUNKS = 64
|
|
MAX_PROVIDER_EVENT_BYTES = 64 * 1024
|
|
MAX_PROVIDER_EVENT_QUEUE_BYTES = 512 * 1024
|
|
MAX_PROVIDER_EVENT_QUEUE_ITEMS = 128
|
|
MAX_PROVIDER_TEXT_BYTES = 32 * 1024
|
|
RECEIVE_TIMEOUT = object()
|
|
|
|
|
|
class ProviderError(RuntimeError):
|
|
"""A bounded, redacted provider failure suitable for the service boundary."""
|
|
|
|
def __init__(self, code: str, detail: str = "transcription provider failed") -> None:
|
|
self.code = code
|
|
try:
|
|
self.detail = detail.encode("utf-8", errors="replace")[:2048].decode("utf-8", errors="ignore")
|
|
except (AttributeError, UnicodeError):
|
|
self.detail = "transcription provider failed"
|
|
super().__init__(self.detail)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BackendDescriptor:
|
|
backend: str
|
|
model: str
|
|
|
|
|
|
@runtime_checkable
|
|
class TranscriptionTurn(Protocol):
|
|
sample_rate_hz: int
|
|
retain_pcm: bool
|
|
|
|
def enqueue_pcm(self, chunk: bytes) -> None: ...
|
|
def finish(self, capture: Capture | None = None) -> str: ...
|
|
def poll_error(self) -> ProviderError | None: ...
|
|
def cancel(self) -> None: ...
|
|
def close(self) -> None: ...
|
|
|
|
@property
|
|
def cleanup_complete(self) -> bool: ...
|
|
|
|
|
|
@runtime_checkable
|
|
class TranscriptionBackend(Protocol):
|
|
descriptor: BackendDescriptor
|
|
sample_rate_hz: int
|
|
|
|
def start_turn(self) -> TranscriptionTurn: ...
|
|
|
|
def cancel_prepare(self) -> None: ...
|
|
|
|
|
|
class PcmQueue:
|
|
"""A byte- and item-bounded nonblocking queue for recorder drain threads."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
max_bytes: int = MAX_PCM_QUEUE_BYTES,
|
|
max_chunks: int = MAX_PCM_QUEUE_CHUNKS,
|
|
) -> None:
|
|
if max_bytes <= 0 or max_chunks <= 0:
|
|
raise ValueError("PCM queue bounds must be positive")
|
|
self._queue: Queue[bytes | None] = Queue(maxsize=max_chunks)
|
|
self._max_bytes = max_bytes
|
|
self._queued_bytes = 0
|
|
self._lock = threading.Lock()
|
|
self._close_lock = threading.Lock()
|
|
self._closed = False
|
|
self._end_queued = False
|
|
|
|
@property
|
|
def queued_bytes(self) -> int:
|
|
with self._lock:
|
|
return self._queued_bytes
|
|
|
|
def put(self, chunk: bytes) -> None:
|
|
if not isinstance(chunk, bytes) or not chunk:
|
|
raise ProviderError("queue_invalid", "provider received invalid audio")
|
|
with self._lock:
|
|
if self._closed:
|
|
raise ProviderError("queue_closed", "provider audio queue is closed")
|
|
if len(chunk) > self._max_bytes or self._queued_bytes + len(chunk) > self._max_bytes:
|
|
raise ProviderError("queue_overflow", "provider audio queue is full")
|
|
self._queued_bytes += len(chunk)
|
|
try:
|
|
self._queue.put_nowait(chunk)
|
|
except Full as error:
|
|
with self._lock:
|
|
self._queued_bytes -= len(chunk)
|
|
raise ProviderError("queue_overflow", "provider audio queue is full") from error
|
|
|
|
def close(self, *, timeout: float = 15.0) -> None:
|
|
if not self._close_lock.acquire(timeout=max(0.0, timeout)):
|
|
raise ProviderError("queue_overflow", "provider audio queue could not be finalized")
|
|
try:
|
|
with self._lock:
|
|
if self._end_queued:
|
|
return
|
|
self._closed = True
|
|
# A full queue still has valid audio that must be drained. The
|
|
# closed flag is a side-channel EOF marker when the sentinel cannot
|
|
# be inserted yet, rather than dropping the queued tail.
|
|
try:
|
|
self._queue.put_nowait(None)
|
|
except Full:
|
|
pass
|
|
with self._lock:
|
|
self._end_queued = True
|
|
finally:
|
|
self._close_lock.release()
|
|
|
|
def get(self, timeout: float | None = None) -> bytes | None:
|
|
deadline = None if timeout is None else time.monotonic() + timeout
|
|
while True:
|
|
poll_timeout = 0.1 if deadline is None else max(0.0, min(0.1, deadline - time.monotonic()))
|
|
try:
|
|
chunk = self._queue.get(timeout=poll_timeout)
|
|
except Empty:
|
|
with self._lock:
|
|
closed = self._closed
|
|
if closed:
|
|
return None
|
|
if deadline is not None and time.monotonic() >= deadline:
|
|
raise
|
|
continue
|
|
if chunk is not None:
|
|
with self._lock:
|
|
self._queued_bytes -= len(chunk)
|
|
return chunk
|
|
|
|
|
|
class LocalNemoTurn:
|
|
"""Post-stop local inference turn retaining audio only in the recorder."""
|
|
|
|
sample_rate_hz = 16_000
|
|
retain_pcm = True
|
|
|
|
def __init__(self, model: Any) -> None:
|
|
self._model = model
|
|
self._cancelled = False
|
|
self._closed = False
|
|
self._finished = False
|
|
|
|
@property
|
|
def cleanup_complete(self) -> bool:
|
|
return self._closed
|
|
|
|
def enqueue_pcm(self, chunk: bytes) -> None:
|
|
# Local turns use Capture.pcm so no second copy is made here.
|
|
if not self.retain_pcm:
|
|
raise ProviderError("local_invalid", "local turn cannot stream audio")
|
|
|
|
def finish(self, capture: Capture | None = None) -> str:
|
|
if self._cancelled:
|
|
raise ProviderError("local_cancelled", "local turn was cancelled")
|
|
if self._closed or self._finished:
|
|
raise ProviderError("local_invalid", "local turn was already finalized")
|
|
self._finished = True
|
|
if capture is None:
|
|
raise ProviderError("local_invalid", "local capture is missing")
|
|
if capture.sample_rate_hz != self.sample_rate_hz:
|
|
raise ProviderError("local_invalid", "local capture has an unsupported sample rate")
|
|
from .audio import pcm_s16le_to_floats
|
|
|
|
max_bytes = getattr(self._model, "max_pcm_bytes", 5 * 60 * 16_000 * 2)
|
|
samples = pcm_s16le_to_floats(capture.pcm, max_bytes=max_bytes)
|
|
result = self._model.transcribe(samples)
|
|
if not isinstance(result, str):
|
|
raise ProviderError("local_failure", "local model returned an invalid transcript")
|
|
return result
|
|
|
|
def poll_error(self) -> ProviderError | None:
|
|
return None
|
|
|
|
def cancel(self) -> None:
|
|
self._cancelled = True
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
|
|
|
|
class LocalNemoBackend:
|
|
"""Backend wrapper preserving the lazy, warm NeMo model behavior."""
|
|
|
|
sample_rate_hz = 16_000
|
|
|
|
def __init__(self, model: Any) -> None:
|
|
self._model = model
|
|
self.descriptor = BackendDescriptor("local", "nemo-parakeet")
|
|
|
|
@property
|
|
def ready(self) -> bool:
|
|
return bool(getattr(self._model, "loaded", False))
|
|
|
|
def start_turn(self) -> LocalNemoTurn:
|
|
self._model.load()
|
|
return LocalNemoTurn(self._model)
|
|
|
|
def cancel_prepare(self) -> None:
|
|
# NeMo model loading is not interruptible, but the service tombstone
|
|
# prevents a cancelled preparation from starting capture.
|
|
return None
|
|
|
|
|
|
class QueuedCloudTurn:
|
|
"""Common ordered sender/receiver for cloud adapters.
|
|
|
|
Provider handshakes happen before this object is created. The recorder only
|
|
calls ``enqueue_pcm``; it never performs network I/O. A sender thread drains
|
|
the bounded queue and sends the end marker after every queued PCM chunk. A
|
|
receiver thread continuously consumes provider events during capture so a
|
|
websocket's receive buffer cannot stall the sender.
|
|
"""
|
|
|
|
sample_rate_hz = 24_000
|
|
retain_pcm = False
|
|
|
|
def __init__(
|
|
self,
|
|
transport: Any,
|
|
*,
|
|
timeout: float = 15.0,
|
|
max_queue_bytes: int = MAX_PCM_QUEUE_BYTES,
|
|
max_queue_chunks: int = MAX_PCM_QUEUE_CHUNKS,
|
|
) -> None:
|
|
if max_queue_bytes <= 0 or max_queue_bytes > MAX_PCM_QUEUE_BYTES:
|
|
raise ValueError("cloud PCM backlog must not exceed five seconds")
|
|
self._transport = transport
|
|
self._timeout = timeout
|
|
self._queue = PcmQueue(max_bytes=max_queue_bytes, max_chunks=max_queue_chunks)
|
|
self._provider_events: Queue[tuple[dict[str, Any], int]] = Queue(maxsize=MAX_PROVIDER_EVENT_QUEUE_ITEMS)
|
|
self._provider_event_bytes = 0
|
|
self._provider_closed = False
|
|
self._sender_error: ProviderError | None = None
|
|
self._receiver_error: ProviderError | None = None
|
|
self._event_lock = threading.Lock()
|
|
self._cancel_event = threading.Event()
|
|
self._transport_close_lock = threading.Lock()
|
|
self._transport_close_thread: threading.Thread | None = None
|
|
self._cancelled = False
|
|
self._finished = False
|
|
self._closed = False
|
|
self._sent_bytes = 0
|
|
self._stream_started_at = time.monotonic()
|
|
self._sender = threading.Thread(target=self._send_loop, daemon=True, name="dictate-provider-send")
|
|
self._receiver = threading.Thread(target=self._receive_loop, daemon=True, name="dictate-provider-recv")
|
|
self._sender.start()
|
|
self._receiver.start()
|
|
|
|
@property
|
|
def cleanup_complete(self) -> bool:
|
|
close_thread = self._transport_close_thread
|
|
return (
|
|
not self._sender.is_alive()
|
|
and not self._receiver.is_alive()
|
|
and (close_thread is None or not close_thread.is_alive())
|
|
)
|
|
|
|
def enqueue_pcm(self, chunk: bytes) -> None:
|
|
if self._cancelled or self._finished:
|
|
raise ProviderError(f"{self._provider_prefix}_queue_closed", "provider turn is no longer accepting audio")
|
|
try:
|
|
self._queue.put(chunk)
|
|
except ProviderError as error:
|
|
raise ProviderError(f"{self._provider_prefix}_{error.code}", error.detail) from error
|
|
|
|
def close_input(self) -> None:
|
|
"""Signal EOF after the recorder has drained stdout to EOF."""
|
|
self._queue.close()
|
|
|
|
def poll_error(self) -> ProviderError | None:
|
|
error = self._sender_error or self._receiver_error
|
|
if error is not None:
|
|
return error
|
|
with self._event_lock:
|
|
provider_closed = self._provider_closed
|
|
if provider_closed and not self._finished and not self._cancelled:
|
|
return ProviderError(
|
|
self._provider_prefix + "_premature_close",
|
|
"provider closed before recording was finalized",
|
|
)
|
|
return None
|
|
|
|
def finish(self, capture: Capture | None = None) -> str:
|
|
if self._cancelled:
|
|
raise ProviderError("cancelled", "provider turn was cancelled")
|
|
if self._finished:
|
|
raise ProviderError("invalid_state", "provider turn was already finalized")
|
|
self._finished = True
|
|
try:
|
|
self._queue.close(timeout=self._timeout)
|
|
except ProviderError:
|
|
self.cancel()
|
|
raise
|
|
self._sender.join(timeout=self._timeout)
|
|
if self._sender.is_alive():
|
|
self.cancel()
|
|
raise ProviderError(self._provider_prefix + "_timeout", "provider audio sender timed out")
|
|
if self._sender_error is not None:
|
|
error = self._sender_error
|
|
self.close()
|
|
raise error
|
|
try:
|
|
transcript = self._finish_provider()
|
|
# A correlated final isn't enough to publish a result while the
|
|
# receiver can still report a queued provider failure. Stop and
|
|
# join it first, then make the receiver error check authoritative.
|
|
self._cancel_event.set()
|
|
self._request_transport_close()
|
|
self._receiver.join(timeout=self._timeout)
|
|
if self._receiver.is_alive():
|
|
raise ProviderError(
|
|
self._provider_prefix + "_timeout",
|
|
"provider receiver did not stop after the final transcript",
|
|
)
|
|
if self._receiver_error is not None:
|
|
raise self._receiver_error
|
|
return transcript
|
|
finally:
|
|
self.close()
|
|
|
|
def cancel(self) -> None:
|
|
if self._cancelled and self.cleanup_complete:
|
|
return
|
|
self._cancelled = True
|
|
self._cancel_event.set()
|
|
self._request_transport_close()
|
|
try:
|
|
self._queue.close(timeout=0.1)
|
|
except ProviderError:
|
|
pass
|
|
# Closing the transport is the cancellation signal. Do not block the
|
|
# protocol loop waiting for provider threads; ``cleanup_complete``
|
|
# keeps the tombstoned generation busy until they actually exit.
|
|
|
|
def close(self) -> None:
|
|
if self._closed and self.cleanup_complete:
|
|
return
|
|
self._cancel_event.set()
|
|
self._request_transport_close()
|
|
# Keep close bounded and nonblocking from the service loop. A later
|
|
# poll retries the joins while cleanup_complete prevents a new turn.
|
|
for thread in (self._sender, self._receiver, self._transport_close_thread):
|
|
if thread is not None:
|
|
thread.join(timeout=min(self._timeout, 0.05))
|
|
self._closed = self.cleanup_complete
|
|
|
|
def _send_loop(self) -> None:
|
|
try:
|
|
while True:
|
|
chunk = self._queue.get(timeout=self._timeout)
|
|
if chunk is None:
|
|
# Queue closure means commit only for a successful recorder
|
|
# EOF. Cancellation closes the same bounded queue merely to
|
|
# wake this worker and must never finalize partial audio.
|
|
if not self._cancel_event.is_set():
|
|
self._send_end()
|
|
return
|
|
# Pace the stream against the declared PCM rate. This prevents
|
|
# a fast recorder from building more than five seconds of
|
|
# provider-side backlog while preserving ordered sends.
|
|
target_elapsed = self._sent_bytes / (self.sample_rate_hz * 2)
|
|
elapsed = time.monotonic() - self._stream_started_at
|
|
if target_elapsed > elapsed and self._cancel_event.wait(target_elapsed - elapsed):
|
|
return
|
|
if self._cancel_event.is_set():
|
|
return
|
|
self._send_audio(chunk)
|
|
self._sent_bytes += len(chunk)
|
|
except ProviderError as error:
|
|
self._sender_error = error
|
|
self._abort_transport()
|
|
except Empty:
|
|
self._sender_error = ProviderError(
|
|
self._provider_prefix + "_timeout", "provider audio sender timed out"
|
|
)
|
|
self._abort_transport()
|
|
except Exception:
|
|
self._sender_error = ProviderError(
|
|
self._provider_prefix + "_send", "provider audio upload failed"
|
|
)
|
|
self._abort_transport()
|
|
|
|
def _receive_loop(self) -> None:
|
|
try:
|
|
while not self._cancel_event.is_set():
|
|
event = self._receive_provider_event(timeout=0.1)
|
|
if event is RECEIVE_TIMEOUT:
|
|
continue
|
|
if event is None:
|
|
with self._event_lock:
|
|
self._provider_closed = True
|
|
return
|
|
if not isinstance(event, dict):
|
|
raise ProviderError(self._provider_prefix + "_malformed_event", "provider event is invalid")
|
|
try:
|
|
size = len(json.dumps(event, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
|
|
except (TypeError, UnicodeError, ValueError, RecursionError, OverflowError) as error:
|
|
raise ProviderError(self._provider_prefix + "_malformed_event", "provider event is invalid") from error
|
|
with self._event_lock:
|
|
if self._provider_event_bytes + size > MAX_PROVIDER_EVENT_QUEUE_BYTES:
|
|
raise ProviderError(self._provider_prefix + "_event_overflow", "provider event queue is full")
|
|
self._provider_event_bytes += size
|
|
try:
|
|
self._provider_events.put_nowait((event, size))
|
|
except Full as error:
|
|
with self._event_lock:
|
|
self._provider_event_bytes -= size
|
|
raise ProviderError(self._provider_prefix + "_event_overflow", "provider event queue is full") from error
|
|
except ProviderError as error:
|
|
# Preserve a provider failure that races with successful
|
|
# finalization. Cancellation does not consume this error, but it
|
|
# avoids recursively closing a transport we are already closing.
|
|
self._receiver_error = error
|
|
if not self._cancel_event.is_set():
|
|
self._abort_transport()
|
|
except Exception:
|
|
if not self._cancel_event.is_set():
|
|
self._receiver_error = ProviderError(
|
|
self._provider_prefix + "_receive", "provider event stream failed"
|
|
)
|
|
self._abort_transport()
|
|
|
|
def _next_provider_event(self, timeout: float) -> dict[str, Any] | None:
|
|
try:
|
|
event, size = self._provider_events.get(timeout=max(0.0, timeout))
|
|
except Empty as error:
|
|
error_from_receiver = self._receiver_error
|
|
if error_from_receiver is not None:
|
|
raise error_from_receiver
|
|
with self._event_lock:
|
|
closed = self._provider_closed
|
|
if closed:
|
|
return None
|
|
raise ProviderError(self._provider_prefix + "_timeout", "timed out waiting for provider event") from error
|
|
with self._event_lock:
|
|
self._provider_event_bytes -= size
|
|
return event
|
|
|
|
@property
|
|
def _provider_prefix(self) -> str:
|
|
return "provider"
|
|
|
|
def _send_audio(self, chunk: bytes) -> None:
|
|
raise NotImplementedError
|
|
|
|
def _send_end(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
def _finish_provider(self) -> str:
|
|
raise NotImplementedError
|
|
|
|
def _receive_provider_event(self, *, timeout: float) -> dict[str, Any] | None | object:
|
|
raise NotImplementedError
|
|
|
|
def _request_transport_close(self) -> None:
|
|
with self._transport_close_lock:
|
|
if self._transport_close_thread is not None:
|
|
return
|
|
self._transport_close_thread = threading.Thread(
|
|
target=self._abort_transport,
|
|
daemon=True,
|
|
name="dictate-provider-close",
|
|
)
|
|
self._transport_close_thread.start()
|
|
|
|
def _abort_transport(self) -> None:
|
|
try:
|
|
self._transport.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def bounded_provider_text(value: Any, *, field: str = "transcript") -> str:
|
|
if not isinstance(value, str):
|
|
raise ProviderError("provider_malformed", f"provider {field} is invalid")
|
|
try:
|
|
if len(value.encode("utf-8")) > MAX_PROVIDER_TEXT_BYTES:
|
|
raise ProviderError("provider_malformed", f"provider {field} is too large")
|
|
except (UnicodeEncodeError, AttributeError) as error:
|
|
raise ProviderError("provider_malformed", f"provider {field} is invalid") from error
|
|
return value
|
|
|
|
|
|
def provider_event(value: Any) -> dict[str, Any]:
|
|
"""Validate one already decoded event and apply a serialized-size bound."""
|
|
|
|
if not isinstance(value, dict) or any(not isinstance(key, str) for key in value):
|
|
raise ProviderError("provider_malformed", "provider event is not an object")
|
|
try:
|
|
encoded = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
|
except (TypeError, UnicodeError, ValueError, RecursionError, OverflowError) as error:
|
|
raise ProviderError("provider_malformed", "provider event is invalid") from error
|
|
if len(encoded) > MAX_PROVIDER_EVENT_BYTES:
|
|
raise ProviderError("provider_malformed", "provider event is too large")
|
|
return value
|
|
|
|
|
|
def backend_from_config(config: Any) -> TranscriptionBackend:
|
|
"""Construct only the selected backend, keeping cloud imports lazy."""
|
|
if config.backend == "local":
|
|
from .backends.local import LocalNemoBackend
|
|
return LocalNemoBackend(config.local)
|
|
if config.backend == "meta":
|
|
from .backends.meta import MetaBackend
|
|
return MetaBackend(config.meta)
|
|
if config.backend == "openai":
|
|
from .backends.openai import OpenAIBackend
|
|
return OpenAIBackend(config.openai)
|
|
raise ProviderError("configuration", "unsupported transcription backend")
|
|
|
|
|
|
__all__ = [
|
|
"BackendDescriptor", "LocalNemoBackend", "LocalNemoTurn", "MAX_PCM_QUEUE_BYTES",
|
|
"MAX_PROVIDER_EVENT_BYTES", "MAX_PROVIDER_EVENT_QUEUE_BYTES", "MAX_PROVIDER_TEXT_BYTES",
|
|
"PcmQueue", "ProviderError", "QueuedCloudTurn", "RECEIVE_TIMEOUT",
|
|
"TranscriptionBackend", "TranscriptionTurn", "backend_from_config",
|
|
"bounded_provider_text", "provider_event",
|
|
]
|