Mercurial

comparison dictation/server.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
comparison
equal deleted inserted replaced
274:c9be578316a6 275:78699f810817
1 from __future__ import annotations
2
3 from contextlib import asynccontextmanager
4 import asyncio
5 from pathlib import Path
6 import uuid
7
8 from aiortc import RTCPeerConnection, RTCSessionDescription
9 from fastapi import FastAPI, HTTPException, Request
10 from fastapi.responses import FileResponse, JSONResponse
11 from pydantic import BaseModel, Field
12
13 from dictation.config import DictationConfig
14 from dictation.session import DictationSession
15 from dictation.transcriber import FasterWhisperTranscriber, Transcriber
16
17
18 class Offer(BaseModel):
19 sdp: str = Field(min_length=1, max_length=256_000)
20 type: str
21
22
23 class DictationService:
24 def __init__(
25 self,
26 config: DictationConfig,
27 transcriber: Transcriber,
28 ) -> None:
29 self.config = config
30 self.transcriber = transcriber
31 self.sessions: dict[str, DictationSession] = {}
32 self.ready = False
33 self.startup_error: str | None = None
34 self._lock = asyncio.Lock()
35
36 async def start(self) -> None:
37 try:
38 await self.transcriber.warmup()
39 self.ready = True
40 except Exception as error:
41 self.startup_error = str(error)
42 raise
43
44 async def accept_offer(self, offer: Offer) -> dict[str, str]:
45 if offer.type != "offer":
46 raise HTTPException(422, "SDP type must be offer")
47 async with self._lock:
48 if len(self.sessions) >= self.config.max_sessions:
49 raise HTTPException(429, "Dictation session limit reached")
50 peer = RTCPeerConnection()
51 session_id = str(uuid.uuid4())
52 session = DictationSession(
53 session_id,
54 peer,
55 self.transcriber,
56 self.config,
57 )
58 self.sessions[session_id] = session
59
60 @peer.on("datachannel")
61 def on_datachannel(channel) -> None:
62 session.attach_channel(channel)
63
64 @peer.on("track")
65 def on_track(track) -> None:
66 session.attach_track(track)
67
68 @peer.on("connectionstatechange")
69 async def on_connectionstatechange() -> None:
70 if peer.connectionState in {"failed", "closed"}:
71 await self.close_session(session_id)
72
73 try:
74 await peer.setRemoteDescription(
75 RTCSessionDescription(sdp=offer.sdp, type=offer.type)
76 )
77 answer = await peer.createAnswer()
78 await peer.setLocalDescription(answer)
79 return {
80 "sdp": peer.localDescription.sdp,
81 "type": peer.localDescription.type,
82 "sessionId": session_id,
83 }
84 except Exception:
85 await self.close_session(session_id)
86 raise
87
88 async def close_session(self, session_id: str) -> bool:
89 async with self._lock:
90 session = self.sessions.pop(session_id, None)
91 if not session:
92 return False
93 await session.close()
94 return True
95
96 async def close(self) -> None:
97 for session_id in list(self.sessions):
98 await self.close_session(session_id)
99 close = getattr(self.transcriber, "close", None)
100 if close:
101 close()
102
103
104 def create_app(
105 config: DictationConfig | None = None,
106 transcriber: Transcriber | None = None,
107 ) -> FastAPI:
108 resolved_config = config or DictationConfig.from_environment()
109 resolved_transcriber = transcriber or FasterWhisperTranscriber(
110 resolved_config.model_dir,
111 resolved_config.compute_type,
112 )
113 service = DictationService(resolved_config, resolved_transcriber)
114 web_root = Path(__file__).parent / "web"
115
116 @asynccontextmanager
117 async def lifespan(app: FastAPI):
118 app.state.dictation = service
119 try:
120 await service.start()
121 yield
122 finally:
123 await service.close()
124
125 app = FastAPI(
126 title="Zenbu WebRTC Dictation",
127 lifespan=lifespan,
128 )
129
130 @app.get("/")
131 async def index():
132 return FileResponse(web_root / "index.html")
133
134 @app.get("/dictation.js")
135 async def javascript():
136 return FileResponse(web_root / "dictation.js")
137
138 @app.get("/dictation.css")
139 async def stylesheet():
140 return FileResponse(web_root / "dictation.css")
141
142 @app.get("/health")
143 async def health(request: Request):
144 current: DictationService = request.app.state.dictation
145 status = 200 if current.ready else 503
146 return JSONResponse(
147 {
148 "status": "ok" if current.ready else "unavailable",
149 "model": str(current.config.model_dir),
150 "computeType": current.config.compute_type,
151 "activeSessions": len(current.sessions),
152 "maxSessions": current.config.max_sessions,
153 "error": current.startup_error,
154 },
155 status_code=status,
156 )
157
158 @app.post("/api/webrtc/offer")
159 async def offer(body: Offer, request: Request):
160 current: DictationService = request.app.state.dictation
161 if not current.ready:
162 raise HTTPException(503, "Dictation model is unavailable")
163 return await current.accept_offer(body)
164
165 @app.post("/api/webrtc/session/{session_id}/close")
166 async def close_session(session_id: str, request: Request):
167 current: DictationService = request.app.state.dictation
168 await current.close_session(session_id)
169 return {"closed": True}
170
171 return app