openspeechapi 0.1.0__py3-none-any.whl

openspeech/server/ws/stt_stream.py

@@ -0,0 +1,263 @@
+"""WebSocket STT streaming endpoint."""
+from __future__ import annotations
+import asyncio
+import json
+import uuid
+from openspeech.logging_config import logger
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from openspeech.core.enums import Capability
+from openspeech.logging_config import bind_context
+from openspeech.telemetry.perf import Event, PerfTimer, milestone
+
+router = APIRouter()
+
+
+@router.websocket("/stream")
+async def stt_stream(
+    websocket: WebSocket,
+    provider: str = "faster-whisper",
+    language: str | None = None,
+    sample_rate: int = 16000,
+):
+    # Establish request_id early so all log lines from this connection are
+    # correlated, including auth-failure paths.
+    request_id = (
+        websocket.query_params.get("request_id")
+        or websocket.headers.get("x-request-id")
+        or uuid.uuid4().hex[:12]
+    )
+
+    with bind_context(request_id=request_id, provider=provider, engine=provider):
+        # WebSocket auth via query param
+        server_config = getattr(websocket.app.state, "server_config", None)
+        if server_config is not None and server_config.auth_enabled:
+            token = websocket.query_params.get("token", "")
+            if token not in server_config.api_keys:
+                milestone(Event.WS_ERROR, reason="unauthorized", scope="stt")
+                await websocket.close(code=4001, reason="Unauthorized")
+                return
+
+        with PerfTimer(Event.WS_TOTAL, scope="stt", provider=provider) as ws_timer:
+            await websocket.accept()
+            milestone(
+                Event.WS_ACCEPT,
+                scope="stt",
+                provider=provider,
+                language=language,
+                sample_rate=sample_rate,
+            )
+            dispatcher = websocket.app.state.dispatcher
+
+            # Determine if the provider supports real streaming
+            handle = dispatcher._handles.get(provider)
+            provider_cls = handle.provider_cls if handle else None
+            supports_streaming = (
+                provider_cls is not None
+                and Capability.STREAMING in getattr(provider_cls, "capabilities", set())
+            )
+
+            # Send meta message so frontend knows batch vs streaming mode
+            await websocket.send_json({
+                "type": "meta",
+                "streaming": supports_streaming,
+                "provider": provider,
+                "request_id": request_id,
+            })
+            milestone(
+                Event.WS_META_SENT,
+                level="verbose",
+                scope="stt",
+                streaming=supports_streaming,
+            )
+
+            try:
+                if supports_streaming:
+                    await _run_streaming(
+                        websocket=websocket,
+                        dispatcher=dispatcher,
+                        provider=provider,
+                        language=language,
+                        ws_timer=ws_timer,
+                    )
+                else:
+                    await _run_batch(
+                        websocket=websocket,
+                        dispatcher=dispatcher,
+                        provider=provider,
+                        language=language,
+                        sample_rate=sample_rate,
+                        ws_timer=ws_timer,
+                    )
+            except WebSocketDisconnect:
+                milestone(Event.WS_CLOSED, scope="stt", reason="client_disconnect")
+            except Exception as exc:
+                milestone(
+                    Event.WS_ERROR,
+                    scope="stt",
+                    error_type=type(exc).__name__,
+                    error_message=str(exc),
+                )
+                logger.exception("STT WS error")
+                try:
+                    await websocket.send_json({"type": "error", "detail": str(exc)})
+                except Exception:
+                    pass
+
+
+async def _run_streaming(
+    *,
+    websocket: WebSocket,
+    dispatcher,
+    provider: str,
+    language: str | None,
+    ws_timer: PerfTimer,
+) -> None:
+    """Pipe real-time audio chunks through ``transcribe_stream``."""
+    audio_queue: asyncio.Queue[bytes | None] = asyncio.Queue()
+    first_audio_logged = False
+    first_response_logged = False
+
+    async def audio_source():
+        while True:
+            chunk = await audio_queue.get()
+            if chunk is None:
+                break
+            yield chunk
+
+    async def receive_audio():
+        """Receive WebSocket messages and push audio into the queue."""
+        nonlocal first_audio_logged
+        frame_count = 0
+        total_bytes = 0
+        try:
+            while True:
+                message = await websocket.receive()
+                if message.get("type") == "websocket.disconnect":
+                    break
+                if message.get("bytes") is not None:
+                    if not first_audio_logged:
+                        first_audio_logged = True
+                        ws_timer.emit_milestone(
+                            Event.WS_FIRST_AUDIO_FRAME,
+                            scope="stt",
+                            bytes=len(message["bytes"]),
+                        )
+                    frame_count += 1
+                    total_bytes += len(message["bytes"])
+                    await audio_queue.put(message["bytes"])
+                elif message.get("text") is not None:
+                    data = json.loads(message["text"])
+                    if data.get("type") == "stop":
+                        break
+        finally:
+            ws_timer.add(frames_received=frame_count, bytes_received=total_bytes)
+            await audio_queue.put(None)
+
+    recv_task = asyncio.create_task(receive_audio())
+
+    # Ensure provider is started before streaming (lazy-load)
+    await dispatcher._lifecycle.ensure_ready(provider)
+
+    # Stream via executor.invoke_stream
+    handle = dispatcher._handles[provider]
+    partial_count = 0
+    final_count = 0
+    async for transcription in handle.executor.invoke_stream(
+        "transcribe_stream", stream=audio_source()
+    ):
+        is_partial = getattr(transcription, "is_partial", True)
+        msg_type = "final" if not is_partial else "partial"
+        if not first_response_logged:
+            first_response_logged = True
+            ws_timer.mark_ttfb()
+            ws_timer.emit_milestone(
+                Event.WS_FIRST_RESPONSE,
+                scope="stt",
+                msg_type=msg_type,
+            )
+        if is_partial:
+            partial_count += 1
+        else:
+            final_count += 1
+            ws_timer.emit_milestone(
+                Event.WS_FINAL_SENT,
+                scope="stt",
+                text_preview=(transcription.text or "")[:80],
+            )
+        await websocket.send_json({
+            "type": msg_type,
+            "text": transcription.text,
+            "confidence": transcription.confidence,
+            "language": transcription.language,
+        })
+
+    # Give receive_audio a chance to finish (client may have already
+    # disconnected after receiving "final"); cancel if it doesn't complete
+    # promptly so we don't hang forever.
+    recv_task.cancel()
+    try:
+        await recv_task
+    except asyncio.CancelledError:
+        pass
+    ws_timer.add(partials=partial_count, finals=final_count, mode="streaming")
+    await websocket.send_json({"type": "closed"})
+    milestone(Event.WS_CLOSED, scope="stt", mode="streaming")
+
+
+async def _run_batch(
+    *,
+    websocket: WebSocket,
+    dispatcher,
+    provider: str,
+    language: str | None,
+    sample_rate: int,
+    ws_timer: PerfTimer,
+) -> None:
+    """Accumulate all audio then transcribe in one batch call (fallback)."""
+    audio_buffer = bytearray()
+    logger.info("STT WS batch mode start")
+
+    while True:
+        message = await websocket.receive()
+        if message.get("type") == "websocket.disconnect":
+            break
+        if message.get("bytes") is not None:
+            audio_buffer.extend(message["bytes"])
+        elif message.get("text") is not None:
+            data = json.loads(message["text"])
+            if data.get("type") == "stop":
+                break
+
+    if audio_buffer:
+        from openspeech.core.enums import AudioFormat
+        from openspeech.core.models import AudioData, STTOptions
+
+        audio_data = AudioData(
+            data=bytes(audio_buffer),
+            sample_rate=sample_rate,
+            channels=1,
+            format=AudioFormat.PCM_16K,
+        )
+        opts = STTOptions(language=language)
+        ws_timer.add(bytes_received=len(audio_buffer), mode="batch")
+        with PerfTimer(Event.DISPATCH_TOTAL, scope="stt", provider=provider, method="transcribe"):
+            result = await dispatcher.stt.transcribe(provider, audio_data, opts)
+        ws_timer.mark_ttfb()
+        if result:
+            ws_timer.emit_milestone(
+                Event.WS_FINAL_SENT,
+                scope="stt",
+                text_preview=(result.text or "")[:80],
+            )
+            await websocket.send_json({
+                "type": "final",
+                "text": result.text,
+                "confidence": result.confidence,
+                "language": result.language,
+            })
+    else:
+        logger.warning("STT WS batch: empty audio buffer, skipping transcription")
+
+    await websocket.send_json({"type": "closed"})
+    milestone(Event.WS_CLOSED, scope="stt", mode="batch")

openspeech/server/ws/tts_stream.py

@@ -0,0 +1,207 @@
+"""WebSocket TTS streaming endpoint."""
+from __future__ import annotations
+import json
+import uuid
+from openspeech.logging_config import logger
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from openspeech.core.enums import Capability
+from openspeech.logging_config import bind_context
+from openspeech.telemetry.perf import Event, PerfTimer, milestone
+
+router = APIRouter()
+
+
+@router.websocket("/stream")
+async def tts_stream(websocket: WebSocket):
+    request_id = (
+        websocket.query_params.get("request_id")
+        or websocket.headers.get("x-request-id")
+        or uuid.uuid4().hex[:12]
+    )
+
+    with bind_context(request_id=request_id):
+        # WebSocket auth via query param
+        server_config = getattr(websocket.app.state, "server_config", None)
+        if server_config is not None and server_config.auth_enabled:
+            token = websocket.query_params.get("token", "")
+            if token not in server_config.api_keys:
+                milestone(Event.WS_ERROR, reason="unauthorized", scope="tts")
+                await websocket.close(code=4001, reason="Unauthorized")
+                return
+
+        with PerfTimer(Event.WS_TOTAL, scope="tts") as ws_timer:
+            await websocket.accept()
+            milestone(Event.WS_ACCEPT, scope="tts")
+            dispatcher = websocket.app.state.dispatcher
+
+            try:
+                # Receive synthesis request
+                message = await websocket.receive_text()
+                data = json.loads(message)
+
+                provider = data.get("provider", "openai-tts")
+                text = data.get("text", "")
+                voice = data.get("voice")
+                speed = data.get("speed", 1.0)
+                model = data.get("model")
+                stream_transport = data.get("stream_transport")
+
+                # Bind provider now that we know it.
+                with bind_context(provider=provider, engine=provider):
+                    ws_timer.add(provider=provider, text_len=len(text))
+
+                    from openspeech.core.models import TTSOptions
+                    opts = TTSOptions(
+                        voice=voice,
+                        speed=float(speed),
+                        model=str(model).strip() if model else None,
+                        stream_transport=str(stream_transport).strip() if stream_transport else None,
+                    )
+
+                    # Check if provider supports true streaming
+                    handle = dispatcher._get_handle(provider)
+                    provider_cls = handle.provider_cls
+                    has_streaming = (
+                        Capability.STREAMING in getattr(provider_cls, "capabilities", set())
+                    )
+
+                    # Detect output format from provider settings when possible
+                    # (fallback to wav for PCM-like providers).
+                    audio_format = "wav"
+                    try:
+                        settings = getattr(handle, "settings_dict", {}) or {}
+                        if "output_format" in settings and settings["output_format"]:
+                            # e.g. mp3_44100_128 -> mp3
+                            audio_format = str(settings["output_format"]).split("_", 1)[0]
+                        elif hasattr(provider_cls, "name") and "iflytek" in provider_cls.name:
+                            audio_format = "mp3"
+                    except Exception:
+                        pass
+
+                    if has_streaming:
+                        await _run_streaming(
+                            websocket=websocket,
+                            dispatcher=dispatcher,
+                            provider=provider,
+                            text=text,
+                            opts=opts,
+                            audio_format=audio_format,
+                            ws_timer=ws_timer,
+                        )
+                    else:
+                        await _run_batch(
+                            websocket=websocket,
+                            dispatcher=dispatcher,
+                            provider=provider,
+                            text=text,
+                            opts=opts,
+                            ws_timer=ws_timer,
+                        )
+            except WebSocketDisconnect:
+                milestone(Event.WS_CLOSED, scope="tts", reason="client_disconnect")
+            except Exception as exc:
+                milestone(
+                    Event.WS_ERROR,
+                    scope="tts",
+                    error_type=type(exc).__name__,
+                    error_message=str(exc),
+                )
+                logger.exception("TTS WS error")
+                try:
+                    await websocket.send_json({"type": "error", "detail": str(exc)})
+                except Exception:
+                    pass
+
+
+async def _run_streaming(
+    *,
+    websocket: WebSocket,
+    dispatcher,
+    provider: str,
+    text: str,
+    opts,
+    audio_format: str,
+    ws_timer: PerfTimer,
+) -> None:
+    """True streaming: yield chunks as they're produced."""
+    await websocket.send_json({
+        "type": "meta",
+        "streaming": True,
+        "audio_format": audio_format,
+    })
+    milestone(Event.WS_META_SENT, level="verbose", scope="tts", streaming=True)
+
+    total_bytes = 0
+    chunk_seq = 0
+    async for chunk in dispatcher.tts.synthesize_stream(provider, text, opts):
+        if chunk.data:
+            if chunk_seq == 0:
+                ws_timer.mark_ttfb()
+                ws_timer.emit_milestone(
+                    Event.WS_FIRST_RESPONSE,
+                    scope="tts",
+                    chunk_bytes=len(chunk.data),
+                )
+            await websocket.send_bytes(chunk.data)
+            total_bytes += len(chunk.data)
+            chunk_seq += 1
+        if chunk.is_final:
+            break
+
+    ws_timer.add(chunks_sent=chunk_seq, bytes_sent=total_bytes, mode="streaming")
+    await websocket.send_json({
+        "type": "done",
+        "total_bytes": total_bytes,
+        "streaming": True,
+        "audio_format": audio_format,
+    })
+    milestone(Event.WS_CLOSED, scope="tts", mode="streaming")
+
+
+async def _run_batch(
+    *,
+    websocket: WebSocket,
+    dispatcher,
+    provider: str,
+    text: str,
+    opts,
+    ws_timer: PerfTimer,
+) -> None:
+    """Batch synthesize, then stream chunks to client."""
+    await dispatcher._lifecycle.ensure_ready(provider)
+    with PerfTimer(Event.DISPATCH_TOTAL, scope="tts", method="synthesize"):
+        result = await dispatcher.tts.synthesize(provider, text, opts)
+
+    if not result:
+        return
+
+    ws_timer.mark_ttfb()
+    fmt = getattr(result, "format", "wav") or "wav"
+    await websocket.send_json({
+        "type": "meta",
+        "sample_rate": result.sample_rate,
+        "channels": result.channels,
+        "duration_ms": result.duration_ms or 0,
+        "total_bytes": len(result.data),
+        "audio_format": str(fmt),
+    })
+    milestone(Event.WS_META_SENT, level="verbose", scope="tts", streaming=False)
+
+    chunk_size = 4096
+    total_bytes = len(result.data)
+    chunk_seq = 0
+    for i in range(0, total_bytes, chunk_size):
+        chunk_data = result.data[i : i + chunk_size]
+        await websocket.send_bytes(chunk_data)
+        chunk_seq += 1
+
+    ws_timer.add(chunks_sent=chunk_seq, bytes_sent=total_bytes, mode="batch")
+    await websocket.send_json({
+        "type": "done",
+        "total_bytes": total_bytes,
+        "sample_rate": result.sample_rate,
+        "streaming": False,
+        "audio_format": str(fmt),
+    })
+    milestone(Event.WS_CLOSED, scope="tts", mode="batch")

openspeech/telemetry/__init__.py

@@ -0,0 +1,21 @@
+"""Telemetry helpers for structured performance logging."""
+
+from openspeech.telemetry.perf import (
+    PerfTimer,
+    Event,
+    milestone,
+    perf_event,
+    perf_enabled,
+    timed,
+    timed_async,
+)
+
+__all__ = [
+    "Event",
+    "PerfTimer",
+    "milestone",
+    "perf_enabled",
+    "perf_event",
+    "timed",
+    "timed_async",
+]