Mercurial

comparison dictation/session.py @ 275:78699f810817

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
comparison
equal deleted inserted replaced
274:c9be578316a6 275:78699f810817
1 from __future__ import annotations
2
3 import asyncio
4 import json
5 from typing import Any
6
7 from aiortc.mediastreams import MediaStreamError
8 from av.audio.resampler import AudioResampler
9 import numpy as np
10
11 from dictation.audio import AudioEventKind, AudioSegmenter
12 from dictation.config import DictationConfig
13 from dictation.transcriber import Transcriber
14
15
16 class DictationSession:
17 def __init__(
18 self,
19 session_id: str,
20 peer: Any,
21 transcriber: Transcriber,
22 config: DictationConfig,
23 ) -> None:
24 self.session_id = session_id
25 self.peer = peer
26 self.transcriber = transcriber
27 self.segmenter = AudioSegmenter(
28 speech_threshold=config.speech_threshold,
29 silence_ms=config.silence_ms,
30 partial_interval_ms=config.partial_interval_ms,
31 max_utterance_seconds=config.max_utterance_seconds,
32 )
33 self.channel = None
34 self.audio_task: asyncio.Task | None = None
35 self.partial_task: asyncio.Task | None = None
36 self.final_tasks: set[asyncio.Task] = set()
37 self.closed = False
38 self.stopping = False
39 self._resampler = AudioResampler(
40 format="s16",
41 layout="mono",
42 rate=16000,
43 )
44
45 def attach_channel(self, channel: Any) -> None:
46 if channel.label != "transcripts":
47 channel.close()
48 return
49 self.channel = channel
50
51 @channel.on("open")
52 def on_open() -> None:
53 self.send("ready", sessionId=self.session_id)
54
55 @channel.on("message")
56 def on_message(message: Any) -> None:
57 if message == "stop":
58 self.stopping = True
59 event = self.segmenter.flush()
60 if event and event.samples is not None:
61 self._start_final(event.samples)
62 else:
63 self.send("speech.ended")
64
65 if channel.readyState == "open":
66 self.send("ready", sessionId=self.session_id)
67
68 def attach_track(self, track: Any) -> None:
69 if track.kind != "audio" or self.audio_task is not None:
70 return
71 self.audio_task = asyncio.create_task(self._consume_audio(track))
72
73 async def _consume_audio(self, track: Any) -> None:
74 try:
75 while not self.closed:
76 frame = await track.recv()
77 for converted in self._resampler.resample(frame):
78 if self.stopping:
79 continue
80 samples = converted.to_ndarray().reshape(-1)
81 normalized = samples.astype(np.float32) / 32768.0
82 for event in self.segmenter.feed(normalized):
83 if event.kind == AudioEventKind.SPEECH_STARTED:
84 self.send("speech.started")
85 elif (
86 event.kind == AudioEventKind.PARTIAL_READY
87 and event.samples is not None
88 ):
89 self._start_partial(event.samples)
90 elif (
91 event.kind == AudioEventKind.FINAL_READY
92 and event.samples is not None
93 ):
94 self._start_final(event.samples)
95 except MediaStreamError:
96 if not self.stopping:
97 event = self.segmenter.flush()
98 if event and event.samples is not None:
99 self._start_final(event.samples)
100 else:
101 self.send("speech.ended")
102 except asyncio.CancelledError:
103 raise
104 except Exception as error:
105 if not self.closed:
106 self.send(
107 "error",
108 code="audio_track_failed",
109 message=str(error),
110 )
111
112 def _start_partial(self, samples: np.ndarray) -> None:
113 if self.partial_task and not self.partial_task.done():
114 return
115 self.partial_task = asyncio.create_task(
116 self._transcribe(samples, final=False)
117 )
118
119 def _start_final(self, samples: np.ndarray) -> None:
120 task = asyncio.create_task(self._transcribe(samples, final=True))
121 self.final_tasks.add(task)
122 task.add_done_callback(self.final_tasks.discard)
123
124 async def _transcribe(self, samples: np.ndarray, *, final: bool) -> None:
125 try:
126 transcript = await self.transcriber.transcribe(
127 samples,
128 final=final,
129 )
130 if transcript.text:
131 self.send(
132 "transcript.final" if final else "transcript.partial",
133 text=transcript.text,
134 language=transcript.language,
135 probability=transcript.probability,
136 )
137 if final:
138 self.send("speech.ended")
139 except asyncio.CancelledError:
140 raise
141 except Exception as error:
142 self.send(
143 "error",
144 code="transcription_failed",
145 message=str(error),
146 )
147
148 def send(self, event_type: str, **payload: Any) -> None:
149 if not self.channel or self.channel.readyState != "open":
150 return
151 self.channel.send(
152 json.dumps(
153 {"type": event_type, **payload},
154 separators=(",", ":"),
155 )
156 )
157
158 async def close(self) -> None:
159 if self.closed:
160 return
161 self.closed = True
162 tasks = [
163 task
164 for task in [
165 self.audio_task,
166 self.partial_task,
167 *self.final_tasks,
168 ]
169 if task is not None and not task.done()
170 ]
171 for task in tasks:
172 task.cancel()
173 if tasks:
174 await asyncio.gather(*tasks, return_exceptions=True)
175 if self.channel and self.channel.readyState != "closed":
176 self.channel.close()
177 await self.peer.close()