openspeechapi 0.1.0__py3-none-any.whl

openspeech/providers/stt/whisperlivekit.py

@@ -0,0 +1,530 @@
+"""WhisperLiveKit STT provider (local service, Deepgram-compatible WS)."""
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+import io
+import json
+from pathlib import Path
+import re
+import time
+from typing import Any
+from uuid import uuid4
+from urllib.parse import urlencode, urljoin, urlparse, urlunparse
+import wave
+
+from openspeech.core.base import STTProvider
+from openspeech.core.enums import Capability, ExecMode, ProviderType
+from openspeech.core.models import AudioData, STTOptions, Transcription
+from openspeech.core.settings import BaseSettings
+from openspeech.logging_config import logger
+
+@dataclass
+class WhisperLiveKitSTTSettings(BaseSettings):
+    api_url: str = "http://127.0.0.1:12100"
+    ws_path: str = "/asr"
+    ws_mode: str = "diff"
+    http_transcribe_path: str = "/v1/audio/transcriptions"
+    model: str = "openai/whisper-large-v3-turbo"
+    language: str = ""
+    sample_rate: int = 16000
+    encoding: str = "linear16"
+    interim_results: bool = True
+    punctuate: bool = True
+    smart_format: bool = True
+    timeout_s: float = 30.0
+    retries: int = 0
+    # Trace logs for streaming parse/debug (JSONL).
+    trace_enabled: bool = False
+    trace_path: str = ".tmp/logs/stt_wlk_trace.jsonl"
+    trace_max_chars: int = 240
+    # WhisperLiveKit (mlx/simul) may return empty text for very short WAV clips.
+    # To improve UX, pad trailing silence up to this duration before request.
+    min_audio_duration_ms: int = 6000
+
+class WhisperLiveKitSTT(STTProvider):
+    name = "whisperlivekit-stt"
+    provider_type = ProviderType.STT
+    execution_mode = ExecMode.LOCAL
+    settings_cls = WhisperLiveKitSTTSettings
+    capabilities = {
+        Capability.STREAMING,
+        Capability.BATCH,
+        Capability.MULTILINGUAL,
+    }
+    field_options = {"language": ["auto", "en", "zh", "ja", "ko", "es", "fr", "de"]}
+
+    def __init__(self, settings: WhisperLiveKitSTTSettings | None = None) -> None:
+        self.settings = settings or WhisperLiveKitSTTSettings()
+        self._client: Any = None
+        self._owns_client: bool = True
+        self._trace_file: Path | None = None
+
+    def set_http_client(self, client) -> None:
+        self._client = client
+        self._owns_client = False
+
+    async def start(self) -> None:
+        if self._client is None:
+            try:
+                import httpx
+            except ImportError:
+                raise ImportError(
+                    "Install httpx + websockets: pip install openspeech[server] openspeech[whisperlivekit]"
+                )
+            self._client = httpx.AsyncClient(timeout=self.settings.timeout_s, trust_env=False)
+            self._owns_client = True
+        logger.info("{} provider started", self.name)
+        if self.settings.trace_enabled:
+            self._trace_file = Path(self.settings.trace_path).expanduser()
+            self._trace_file.parent.mkdir(parents=True, exist_ok=True)
+            self._trace("start", ws_path=self.settings.ws_path, ws_mode=self.settings.ws_mode)
+
+    async def stop(self) -> None:
+        if self._client and self._owns_client:
+            await self._client.aclose()
+        self._client = None
+        logger.info("{} provider stopped", self.name)
+        self._trace("stop")
+
+    def _truncate(self, value: Any) -> Any:
+        if value is None:
+            return None
+        s = str(value)
+        max_chars = max(40, int(self.settings.trace_max_chars))
+        if len(s) <= max_chars:
+            return s
+        return s[:max_chars] + "..."
+
+    def _trace(self, event: str, **payload: Any) -> None:
+        if not self.settings.trace_enabled or self._trace_file is None:
+            return
+        row = {
+            "ts": time.time(),
+            "event": event,
+            **payload,
+        }
+        try:
+            self._trace_file.parent.mkdir(parents=True, exist_ok=True)
+            with self._trace_file.open("a", encoding="utf-8") as f:
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("Failed to write stt_wlk trace log: {}", exc)
+
+    async def health_check(self) -> bool:
+        if self._client is None:
+            return False
+        try:
+            url = urljoin(self.settings.api_url.rstrip("/") + "/", "health")
+            r = await self._client.get(url)
+            return r.status_code < 500
+        except Exception:
+            return False
+
+    async def transcribe(
+        self, audio: AudioData, opts: STTOptions | None = None
+    ) -> Transcription:
+        if self._client is None:
+            raise RuntimeError("Provider not started — call start() first")
+        logger.info("{}: request received, audio={} bytes", self.name, len(audio.data))
+        opts = opts or STTOptions()
+        url = urljoin(
+            self.settings.api_url.rstrip("/") + "/",
+            self.settings.http_transcribe_path.lstrip("/"),
+        )
+        model = (opts.model or self.settings.model).strip()
+        language = (opts.language or self.settings.language).strip()
+        wav_bytes = self._maybe_pad_wav_for_short_clip(
+            audio.data, min_duration_ms=max(0, int(self.settings.min_audio_duration_ms))
+        )
+        files = {
+            "file": ("audio.wav", wav_bytes, "audio/wav"),
+        }
+        data: dict[str, Any] = {"model": model}
+        if language:
+            data["language"] = language
+        if opts.prompt:
+            data["prompt"] = opts.prompt
+        if opts.temperature is not None:
+            data["temperature"] = opts.temperature
+
+        last_exc: Exception | None = None
+        attempts = max(0, int(self.settings.retries)) + 1
+        resp = None
+        started_at = time.perf_counter()
+        for _ in range(attempts):
+            try:
+                resp = await self._client.post(url, files=files, data=data)
+                break
+            except Exception as exc:  # noqa: BLE001
+                last_exc = exc
+        if resp is None:
+            raise RuntimeError(f"WhisperLiveKit request failed: {last_exc}") from last_exc
+        resp.raise_for_status()
+        payload = resp.json()
+        text = str(payload.get("text", "")).strip()
+        if not text:
+            # Deepgram-like response fallback
+            channel = payload.get("channel", {})
+            alts = channel.get("alternatives", []) if isinstance(channel, dict) else []
+            if alts:
+                text = str(alts[0].get("transcript", "")).strip()
+        elapsed_ms = int((time.perf_counter() - started_at) * 1000)
+        duration_ms = payload.get("duration_ms")
+        if duration_ms is None:
+            metadata = payload.get("metadata", {})
+            if isinstance(metadata, dict):
+                md_dur = metadata.get("duration")
+                if isinstance(md_dur, (int, float)):
+                    duration_ms = int(float(md_dur) * 1000)
+        if duration_ms is None:
+            duration_ms = elapsed_ms
+        result = Transcription(
+            text=text,
+            language=payload.get("language") or language or None,
+            confidence=payload.get("confidence"),
+            duration_ms=int(duration_ms),
+        )
+        logger.info("{}: completed in {:.0f}ms, result={} chars", self.name, elapsed_ms, len(result.text))
+        return result
+
+    @staticmethod
+    def _maybe_pad_wav_for_short_clip(data: bytes, min_duration_ms: int) -> bytes:
+        """Pad trailing silence for short PCM WAV clips to avoid empty short-utterance outputs."""
+        if min_duration_ms <= 0 or len(data) < 44:
+            return data
+        if not (data.startswith(b"RIFF") and data[8:12] == b"WAVE"):
+            return data
+        try:
+            with wave.open(io.BytesIO(data), "rb") as r:
+                params = r.getparams()
+                framerate = int(r.getframerate())
+                sampwidth = int(r.getsampwidth())
+                channels = int(r.getnchannels())
+                frames = r.readframes(r.getnframes())
+        except Exception:
+            return data
+        if framerate <= 0 or sampwidth <= 0 or channels <= 0:
+            return data
+        target_frames = int((min_duration_ms / 1000.0) * framerate)
+        current_frames = len(frames) // max(1, sampwidth * channels)
+        if current_frames >= target_frames:
+            return data
+        add_frames = target_frames - current_frames
+        silence = b"\x00" * (add_frames * sampwidth * channels)
+        out = io.BytesIO()
+        with wave.open(out, "wb") as w:
+            w.setparams(params)
+            w.writeframes(frames + silence)
+        return out.getvalue()
+
+    def _build_ws_url(self, model: str, language: str) -> str:
+        base = self.settings.api_url.rstrip("/")
+        u = urlparse(base)
+        scheme = "wss" if u.scheme == "https" else "ws"
+        path = self.settings.ws_path if self.settings.ws_path.startswith("/") else f"/{self.settings.ws_path}"
+        query: dict[str, str] = {}
+        if path == "/asr":
+            query["mode"] = self.settings.ws_mode or "diff"
+            if language:
+                query["language"] = language
+        else:
+            query = {
+                "model": model,
+                "encoding": self.settings.encoding,
+                "sample_rate": str(int(self.settings.sample_rate)),
+                "interim_results": "true" if self.settings.interim_results else "false",
+                "punctuate": "true" if self.settings.punctuate else "false",
+                "smart_format": "true" if self.settings.smart_format else "false",
+            }
+            if language:
+                query["language"] = language
+        if path != "/asr" and language:
+            query["language"] = language
+        return urlunparse((scheme, u.netloc, path, "", urlencode(query), ""))
+
+    @staticmethod
+    def _merge_text(prev: str, incoming: str) -> str:
+        p = (prev or "").strip()
+        n = (incoming or "").strip()
+        if not n:
+            return p
+        if not p:
+            return n
+        if n == p:
+            return p
+        if n.startswith(p):
+            return n
+        if p.startswith(n):
+            return p
+        if p in n:
+            return n
+        if n in p:
+            return p
+        return f"{p} {n}".strip()
+
+    @staticmethod
+    def _snapshot_text(prev: str, incoming: str) -> str:
+        p = (prev or "").strip()
+        n = (incoming or "").strip()
+        if not n:
+            return p
+        if n == p:
+            return p
+        # WhisperLiveKit /asr (diff/snapshot) messages represent latest hypothesis.
+        # Replace instead of append to avoid repeated growth on revisions.
+        return n
+
+    @staticmethod
+    def _stabilize_asr_text(text: str) -> str:
+        t = str(text or "").strip()
+        if not t:
+            return ""
+        # Collapse obvious repeated Latin chunks, e.g. "HelloHelloHello" -> "Hello".
+        t = re.sub(r"(?i)\b([a-z][a-z0-9']{1,})(?:\1){1,}\b", r"\1", t)
+        # Collapse repeated word with spaces, e.g. "hello hello hello" -> "hello".
+        t = re.sub(r"(?i)\b([a-z][a-z0-9']{1,})\b(?:\s+\1\b){1,}", r"\1", t)
+        # Collapse CJK/punctuation long runs (3+) caused by unstable interim decoding.
+        t = re.sub(r"([\u4e00-\u9fff])\1{2,}", r"\1", t)
+        t = re.sub(r"([,，。.!！?？、;；:：\"'`~\-])\1{2,}", r"\1", t)
+        return t.strip()
+
+    @staticmethod
+    def _extract_asr_text_from_lines(lines: list[Any]) -> str:
+        parts: list[str] = []
+        for line in lines:
+            if not isinstance(line, dict):
+                continue
+            if int(line.get("speaker", 0) or 0) == -2:
+                continue
+            t = str(line.get("text") or "").strip()
+            if t:
+                parts.append(t)
+        return " ".join(parts).strip()
+
+    def _extract_asr_message_text(
+        self,
+        data: dict[str, Any],
+        committed_lines: list[dict[str, Any]],
+    ) -> str:
+        t = str(data.get("type") or "").strip()
+        if t == "snapshot":
+            lines = data.get("lines") if isinstance(data.get("lines"), list) else []
+            committed_lines.clear()
+            committed_lines.extend([x for x in lines if isinstance(x, dict)])
+        elif t == "diff":
+            pruned = int(data.get("lines_pruned") or 0)
+            if pruned > 0:
+                del committed_lines[:pruned]
+            new_lines = data.get("new_lines") if isinstance(data.get("new_lines"), list) else []
+            committed_lines.extend([x for x in new_lines if isinstance(x, dict)])
+        elif t == "ready_to_stop":
+            return ""
+        elif "lines" in data and isinstance(data.get("lines"), list):
+            committed_lines.clear()
+            committed_lines.extend([x for x in data.get("lines", []) if isinstance(x, dict)])
+        base = self._extract_asr_text_from_lines(committed_lines)
+        buffer_t = str(data.get("buffer_transcription") or "").strip()
+        if buffer_t:
+            return self._merge_text(base, buffer_t)
+        return base
+
+    @staticmethod
+    def _extract_deepgram_text(data: dict[str, Any]) -> tuple[str, str | None, float | None]:
+        if data.get("type") != "Results":
+            return "", None, None
+        channel = data.get("channel", {})
+        alts = channel.get("alternatives", []) if isinstance(channel, dict) else []
+        if not alts:
+            return "", None, None
+        transcript = str(alts[0].get("transcript", "")).strip()
+        language = channel.get("detected_language")
+        confidence = alts[0].get("confidence")
+        return transcript, language, confidence
+
+    async def transcribe_stream(
+        self, stream: AsyncIterator[bytes]
+    ) -> AsyncIterator[Any]:
+        if self._client is None:
+            raise RuntimeError("Provider not started — call start() first")
+        try:
+            import websockets
+        except ImportError:
+            raise ImportError(
+                "Install websockets: pip install openspeech[server]"
+            )
+
+        _t0 = time.perf_counter()
+
+        ws_url = self._build_ws_url(self.settings.model, self.settings.language)
+        is_asr = self.settings.ws_path.startswith("/asr")
+        stream_id = uuid4().hex[:12]
+        def trace(ev: str, **kw: Any) -> None:
+            self._trace(ev, stream_id=stream_id, **kw)
+
+        logger.debug("{}: connecting to WhisperLiveKit WebSocket...", self.name)
+        trace("stream_open", ws_url=ws_url, is_asr=is_asr)
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=20) as ws:
+            _t_connected = time.perf_counter()
+            logger.info("{}: WS connected in {:.0f}ms", self.name,
+                        (_t_connected - _t0) * 1000)
+            committed_lines: list[dict[str, Any]] = []
+            latest_text = ""
+            sent_chunks = 0
+            sent_bytes = 0
+            _resp_count = 0
+
+            _sender_stop = asyncio.Event()
+
+            async def send_audio() -> None:
+                nonlocal sent_chunks, sent_bytes
+                try:
+                    async for chunk in stream:
+                        if _sender_stop.is_set():
+                            break
+                        if chunk:
+                            await ws.send(chunk)
+                            sent_chunks += 1
+                            sent_bytes += len(chunk)
+                            if sent_chunks == 1:
+                                logger.debug("{}: first frame sent at {:.0f}ms",
+                                             self.name, (time.perf_counter() - _t0) * 1000)
+                            trace("stream_send_chunk", sent_chunks=sent_chunks, sent_bytes=sent_bytes)
+                    # WhisperLiveKit stop sentinel for both /asr and /v1/listen.
+                    if not _sender_stop.is_set():
+                        await ws.send(b"")
+                        trace("stream_send_stop", sent_chunks=sent_chunks, sent_bytes=sent_bytes)
+                except websockets.exceptions.ConnectionClosed:
+                    pass
+                finally:
+                    logger.debug(
+                        "{}: stream sender done, sent {} frames ({} bytes) in {:.0f}ms",
+                        self.name, sent_chunks, sent_bytes,
+                        (time.perf_counter() - _t0) * 1000,
+                    )
+
+            import asyncio
+            send_task = asyncio.create_task(send_audio())
+            try:
+                while True:
+                    try:
+                        msg = await asyncio.wait_for(ws.recv(), timeout=1.2)
+                    except asyncio.TimeoutError:
+                        if send_task.done():
+                            trace("stream_recv_timeout_done")
+                            break
+                        continue
+                    if isinstance(msg, bytes):
+                        trace("stream_recv_binary", size=len(msg))
+                        continue
+                    data = json.loads(msg)
+                    msg_type = str(data.get("type") or "")
+                    if is_asr:
+                        trace(
+                            "stream_recv_asr",
+                            type=msg_type,
+                            seq=data.get("seq"),
+                            status=data.get("status"),
+                            n_lines=data.get("n_lines"),
+                            lines_pruned=data.get("lines_pruned"),
+                            buffer_transcription=self._truncate(data.get("buffer_transcription")),
+                            new_lines=self._truncate(
+                                " | ".join(
+                                    str(x.get("text", ""))
+                                    for x in (data.get("new_lines") or [])
+                                    if isinstance(x, dict)
+                                )
+                            ),
+                        )
+                    else:
+                        transcript, _lang, _conf = self._extract_deepgram_text(data)
+                        trace(
+                            "stream_recv_deepgram",
+                            type=msg_type,
+                            transcript=self._truncate(transcript),
+                        )
+                    _resp_count += 1
+                    if _resp_count == 1:
+                        logger.debug("{}: first response at {:.0f}ms  type={}",
+                                     self.name, (time.perf_counter() - _t0) * 1000, msg_type)
+                    if msg_type in {"error", "Error"}:
+                        detail = str(data.get("detail") or data.get("message") or data)
+                        trace("stream_error", detail=self._truncate(detail))
+                        raise RuntimeError(f"WhisperLiveKit stream error: {detail}")
+                    if is_asr:
+                        parsed_text = self._extract_asr_message_text(data, committed_lines)
+                        stabilized_text = self._stabilize_asr_text(parsed_text)
+                        trace(
+                            "stream_parse_asr",
+                            parsed_text=self._truncate(parsed_text),
+                            stabilized_text=self._truncate(stabilized_text),
+                            latest_text_before=self._truncate(latest_text),
+                        )
+                        if not stabilized_text:
+                            continue
+                        merged = self._snapshot_text(latest_text, stabilized_text)
+                        trace(
+                            "stream_merge_asr",
+                            merged_text=self._truncate(merged),
+                            changed=(merged != latest_text),
+                        )
+                        if merged == latest_text:
+                            continue
+                        latest_text = merged
+                        trace("stream_emit_asr", emit_text=self._truncate(latest_text))
+                        yield Transcription(
+                            text=latest_text,
+                            language=self.settings.language or None,
+                            confidence=None,
+                            is_partial=True,
+                        )
+                    else:
+                        transcript, language, confidence = self._extract_deepgram_text(data)
+                        trace(
+                            "stream_parse_deepgram",
+                            transcript=self._truncate(transcript),
+                            latest_text_before=self._truncate(latest_text),
+                        )
+                        if not transcript:
+                            continue
+                        merged = self._merge_text(latest_text, transcript)
+                        trace(
+                            "stream_merge_deepgram",
+                            merged_text=self._truncate(merged),
+                            changed=(merged != latest_text),
+                        )
+                        if merged == latest_text:
+                            continue
+                        latest_text = merged
+                        trace("stream_emit_deepgram", emit_text=self._truncate(latest_text))
+                        yield Transcription(
+                            text=latest_text,
+                            language=language or self.settings.language or None,
+                            confidence=confidence,
+                            is_partial=True,
+                        )
+            finally:
+                _sender_stop.set()
+                send_task.cancel()
+                try:
+                    await send_task
+                except asyncio.CancelledError:
+                    pass
+                trace("stream_close", sent_chunks=sent_chunks, sent_bytes=sent_bytes)
+
+            # Emit final snapshot after the receive loop exits
+            if latest_text:
+                logger.info("{}: final result at {:.0f}ms  responses={}  text='{}'",
+                            self.name, (time.perf_counter() - _t0) * 1000,
+                            _resp_count, latest_text[:60])
+                yield Transcription(
+                    text=latest_text,
+                    language=self.settings.language or None,
+                    confidence=None,
+                    is_partial=False,
+                )
+
+            logger.info(
+                "{}: stream completed in {:.0f}ms, frames={}",
+                self.name, (time.perf_counter() - _t0) * 1000, sent_chunks,
+            )