diff dictation/session.py @ 275:78699f810817

Add Qwen3-VL and WebRTC dictation services Add Bazel targets for the CUDA-backed Qwen3-VL server and a local WebRTC faster-whisper dictation service. Co-authored-by: Copilot <[email protected]> Copilot-Session: e3d8cb06-6c95-4ae0-9757-651d3796ab00
author MrJuneJune <me@mrjunejune.com>
date Mon, 17 Aug 2026 10:58:47 -0700
parents
children 49e9e591c9bb
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/dictation/session.py	Mon Aug 17 10:58:47 2026 -0700
@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+from aiortc.mediastreams import MediaStreamError
+from av.audio.resampler import AudioResampler
+import numpy as np
+
+from dictation.audio import AudioEventKind, AudioSegmenter
+from dictation.config import DictationConfig
+from dictation.transcriber import Transcriber
+
+
+class DictationSession:
+    def __init__(
+        self,
+        session_id: str,
+        peer: Any,
+        transcriber: Transcriber,
+        config: DictationConfig,
+    ) -> None:
+        self.session_id = session_id
+        self.peer = peer
+        self.transcriber = transcriber
+        self.segmenter = AudioSegmenter(
+            speech_threshold=config.speech_threshold,
+            silence_ms=config.silence_ms,
+            partial_interval_ms=config.partial_interval_ms,
+            max_utterance_seconds=config.max_utterance_seconds,
+        )
+        self.channel = None
+        self.audio_task: asyncio.Task | None = None
+        self.partial_task: asyncio.Task | None = None
+        self.final_tasks: set[asyncio.Task] = set()
+        self.closed = False
+        self.stopping = False
+        self._resampler = AudioResampler(
+            format="s16",
+            layout="mono",
+            rate=16000,
+        )
+
+    def attach_channel(self, channel: Any) -> None:
+        if channel.label != "transcripts":
+            channel.close()
+            return
+        self.channel = channel
+
+        @channel.on("open")
+        def on_open() -> None:
+            self.send("ready", sessionId=self.session_id)
+
+        @channel.on("message")
+        def on_message(message: Any) -> None:
+            if message == "stop":
+                self.stopping = True
+                event = self.segmenter.flush()
+                if event and event.samples is not None:
+                    self._start_final(event.samples)
+                else:
+                    self.send("speech.ended")
+
+        if channel.readyState == "open":
+            self.send("ready", sessionId=self.session_id)
+
+    def attach_track(self, track: Any) -> None:
+        if track.kind != "audio" or self.audio_task is not None:
+            return
+        self.audio_task = asyncio.create_task(self._consume_audio(track))
+
+    async def _consume_audio(self, track: Any) -> None:
+        try:
+            while not self.closed:
+                frame = await track.recv()
+                for converted in self._resampler.resample(frame):
+                    if self.stopping:
+                        continue
+                    samples = converted.to_ndarray().reshape(-1)
+                    normalized = samples.astype(np.float32) / 32768.0
+                    for event in self.segmenter.feed(normalized):
+                        if event.kind == AudioEventKind.SPEECH_STARTED:
+                            self.send("speech.started")
+                        elif (
+                            event.kind == AudioEventKind.PARTIAL_READY
+                            and event.samples is not None
+                        ):
+                            self._start_partial(event.samples)
+                        elif (
+                            event.kind == AudioEventKind.FINAL_READY
+                            and event.samples is not None
+                        ):
+                            self._start_final(event.samples)
+        except MediaStreamError:
+            if not self.stopping:
+                event = self.segmenter.flush()
+                if event and event.samples is not None:
+                    self._start_final(event.samples)
+                else:
+                    self.send("speech.ended")
+        except asyncio.CancelledError:
+            raise
+        except Exception as error:
+            if not self.closed:
+                self.send(
+                    "error",
+                    code="audio_track_failed",
+                    message=str(error),
+                )
+
+    def _start_partial(self, samples: np.ndarray) -> None:
+        if self.partial_task and not self.partial_task.done():
+            return
+        self.partial_task = asyncio.create_task(
+            self._transcribe(samples, final=False)
+        )
+
+    def _start_final(self, samples: np.ndarray) -> None:
+        task = asyncio.create_task(self._transcribe(samples, final=True))
+        self.final_tasks.add(task)
+        task.add_done_callback(self.final_tasks.discard)
+
+    async def _transcribe(self, samples: np.ndarray, *, final: bool) -> None:
+        try:
+            transcript = await self.transcriber.transcribe(
+                samples,
+                final=final,
+            )
+            if transcript.text:
+                self.send(
+                    "transcript.final" if final else "transcript.partial",
+                    text=transcript.text,
+                    language=transcript.language,
+                    probability=transcript.probability,
+                )
+            if final:
+                self.send("speech.ended")
+        except asyncio.CancelledError:
+            raise
+        except Exception as error:
+            self.send(
+                "error",
+                code="transcription_failed",
+                message=str(error),
+            )
+
+    def send(self, event_type: str, **payload: Any) -> None:
+        if not self.channel or self.channel.readyState != "open":
+            return
+        self.channel.send(
+            json.dumps(
+                {"type": event_type, **payload},
+                separators=(",", ":"),
+            )
+        )
+
+    async def close(self) -> None:
+        if self.closed:
+            return
+        self.closed = True
+        tasks = [
+            task
+            for task in [
+                self.audio_task,
+                self.partial_task,
+                *self.final_tasks,
+            ]
+            if task is not None and not task.done()
+        ]
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        if self.channel and self.channel.readyState != "closed":
+            self.channel.close()
+        await self.peer.close()