jb/harness-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(dictate): add voice-reactive listening waveform

Browse Source
This commit is contained in:
jb
2026-08-26 13:40:47 -04:00
parent e234468a73
commit d1a811a5c5
16 changed files with 383 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/dictate/tests/test_model.py
+17 -11
View File
@@ -35,13 +35,15 @@ class FakeNemoModel:
return [type("Hypothesis", (), {"text": "hello locally"})()]
class FakeASRModel:
class FakeConcreteParakeetModel:
restored_path: str | None = None
restored_location: object | None = None
restored: FakeNemoModel | None = None
@classmethod
def restore_from(cls, restore_path: str, map_location: str) -> FakeNemoModel:
def restore_from(cls, restore_path: str, map_location: object) -> FakeNemoModel:
cls.restored_path = restore_path
cls.restored_location = map_location
self = FakeNemoModel()
cls.restored = self
return self
@@ -69,28 +71,29 @@ class NemoParakeetModelTests(unittest.TestCase):
boost_alpha=0.5,
expected_sha256=hashlib.sha256(b"model").hexdigest(),
)
self.assertIsNone(FakeASRModel.restored_path)
self.assertIsNone(FakeConcreteParakeetModel.restored_path)
adapter.load()
transcript = adapter.transcribe([0.0, 0.5])
self.assertEqual(os.environ["HF_HUB_OFFLINE"], "1")
self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], "1")
self.assertEqual(FakeASRModel.restored_path, str(model_path))
self.assertTrue(FakeASRModel.restored.eval_called)
self.assertEqual(FakeASRModel.restored.to_device, "cuda")
self.assertEqual(FakeConcreteParakeetModel.restored_path, str(model_path))
self.assertEqual(FakeConcreteParakeetModel.restored_location, "device:cuda")
self.assertTrue(FakeConcreteParakeetModel.restored.eval_called)
self.assertEqual(FakeConcreteParakeetModel.restored.to_device, "cuda")
self.assertEqual(transcript, "hello locally")
self.assertEqual(FakeASRModel.restored.audio, [([0.0, 0.5], "float32")])
self.assertEqual(FakeConcreteParakeetModel.restored.audio, [([0.0, 0.5], "float32")])
self.assertEqual(
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
["JSON", "npm test"],
)
self.assertEqual(
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
0.5,
)
self.assertEqual(
FakeASRModel.restored.transcribe_options,
FakeConcreteParakeetModel.restored.transcribe_options,
{"batch_size": 1, "num_workers": 0, "verbose": False},
)
@@ -123,12 +126,14 @@ class NemoParakeetModelTests(unittest.TestCase):
collections = ModuleType("nemo.collections")
asr = ModuleType("nemo.collections.asr")
models = ModuleType("nemo.collections.asr.models")
models.ASRModel = FakeASRModel
models.EncDecRNNTBPEModel = FakeConcreteParakeetModel
asr.models = models
collections.asr = asr
nemo.collections = collections
omegaconf = ModuleType("omegaconf")
omegaconf.open_dict = lambda _config: nullcontext()
torch = ModuleType("torch")
torch.device = lambda name: f"device:{name}"
return {
"nemo": nemo,
"nemo.collections": collections,
@@ -136,6 +141,7 @@ class NemoParakeetModelTests(unittest.TestCase):
"nemo.collections.asr.models": models,
"numpy": FakeNumpy("numpy"),
"omegaconf": omegaconf,
"torch": torch,
}
services/dictate/tests/test_protocol_audio.py
+13 -12
View File
@@ -7,6 +7,7 @@ from pi_dictate_service.protocol import (
MAX_FRAME_BYTES,
MAX_ID_BYTES,
MAX_TRANSCRIPT_BYTES,
PROTOCOL_VERSION,
FrameDecoder,
ProtocolError,
encode_frame,
@@ -18,10 +19,10 @@ class FrameDecoderTests(unittest.TestCase):
def test_decodes_fragmented_and_multiple_frames(self) -> None:
decoder = FrameDecoder()
self.assertEqual(decoder.feed(b'{"version":1'), [])
self.assertEqual(decoder.feed(f'{{"version":{PROTOCOL_VERSION}'.encode()), [])
frames = decoder.feed(
b',"type":"ping","requestId":"one"}\n'
b'{"version":1,"type":"ping","requestId":"two"}\n'
+ f'{{"version":{PROTOCOL_VERSION},"type":"ping","requestId":"two"}}\n'.encode()
)
self.assertEqual([frame["requestId"] for frame in frames], ["one", "two"])
@@ -41,14 +42,14 @@ class FrameDecoderTests(unittest.TestCase):
class RequestValidationTests(unittest.TestCase):
def test_requires_exact_fields_bounded_ids_and_supported_version(self) -> None:
request = parse_request(
{"version": 1, "type": "start", "requestId": "r1", "generationId": "g1"}
{"version": PROTOCOL_VERSION, "type": "start", "requestId": "r1", "generationId": "g1"}
)
self.assertEqual(request.generation_id, "g1")
with self.assertRaises(ProtocolError) as unknown:
parse_request(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "start",
"requestId": "r1",
"generationId": "g1",
@@ -61,7 +62,7 @@ class RequestValidationTests(unittest.TestCase):
with self.assertRaises(ProtocolError) as long_id:
parse_request(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "ping",
"requestId": "é" * ((MAX_ID_BYTES // 2) + 1),
}
@@ -69,7 +70,7 @@ class RequestValidationTests(unittest.TestCase):
self.assertEqual(long_id.exception.code, "invalid_request")
with self.assertRaises(ProtocolError) as unsupported:
parse_request({"version": 2, "type": "hello", "requestId": "r2"})
parse_request({"version": PROTOCOL_VERSION + 1, "type": "hello", "requestId": "r2"})
self.assertEqual(unsupported.exception.code, "unsupported_version")
self.assertTrue(unsupported.exception.close_connection)
@@ -78,12 +79,12 @@ class RequestValidationTests(unittest.TestCase):
with self.subTest(request_type=request_type):
with self.assertRaises(ProtocolError) as failed:
parse_request(
{"version": 1, "type": request_type, "requestId": "request"}
{"version": PROTOCOL_VERSION, "type": request_type, "requestId": "request"}
)
self.assertEqual(failed.exception.code, "invalid_request")
with self.assertRaises(ProtocolError) as surrogate:
parse_request({"version": 1, "type": "ping", "requestId": "\ud800"})
parse_request({"version": PROTOCOL_VERSION, "type": "ping", "requestId": "\ud800"})
self.assertEqual(surrogate.exception.request_id, "")
huge_integer = b'{"version":' + (b"9" * 5000) + b',"type":"hello","requestId":"r"}\n'
@@ -96,13 +97,13 @@ class RequestValidationTests(unittest.TestCase):
oversized = "x" * (MAX_ID_BYTES + 1)
with self.assertRaises(ProtocolError) as unsupported:
parse_request(
{"version": 2, "type": "hello", "requestId": oversized}
{"version": PROTOCOL_VERSION + 1, "type": "hello", "requestId": oversized}
)
self.assertEqual(unsupported.exception.request_id, "")
encoded = encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "error",
"requestId": oversized,
"generationId": "\ud800",
@@ -118,7 +119,7 @@ class RequestValidationTests(unittest.TestCase):
with self.assertRaises(ProtocolError) as transcript:
encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "result",
"generationId": "g1",
"transcript": "é" * ((MAX_TRANSCRIPT_BYTES // 2) + 1),
@@ -128,7 +129,7 @@ class RequestValidationTests(unittest.TestCase):
encoded = encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "error",
"requestId": "r1",
"code": "failed",
services/dictate/tests/test_recorder.py
+69 -3
View File
@@ -1,19 +1,34 @@
import io
import math
import subprocess
import threading
import unittest
from pi_dictate_service.recorder import PW_RECORD_COMMAND, PwRecorder, RecorderError
from pi_dictate_service.recorder import (
AUDIO_LEVEL_WINDOW_BYTES,
PW_RECORD_COMMAND,
PwRecorder,
RecorderError,
)
class TrackingStream(io.BytesIO):
def __init__(self, value: bytes, events: list[str], name: str) -> None:
def __init__(
self,
value: bytes,
events: list[str],
name: str,
max_read_size: int | None = None,
) -> None:
super().__init__(value)
self._events = events
self._name = name
self._max_read_size = max_read_size
self.eof = threading.Event()
def read(self, size: int = -1) -> bytes:
if self._max_read_size is not None:
size = min(size, self._max_read_size)
value = super().read(size)
if not value and not self.eof.is_set():
self._events.append(f"{self._name}-eof")
@@ -91,10 +106,11 @@ class PwRecorderTests(unittest.TestCase):
stderr_limit: int = 64 * 1024,
clock: Clock | None = None,
terminate_returncode: int = -15,
stdout_read_size: int | None = None,
) -> tuple[PwRecorder, FakeProcess, ProcessFactory, list[str]]:
events: list[str] = []
process = FakeProcess(
TrackingStream(pcm, events, "stdout"),
TrackingStream(pcm, events, "stdout", stdout_read_size),
TrackingStream(stderr, events, "stderr"),
returncode=returncode,
timeout=timeout,
@@ -138,6 +154,56 @@ class PwRecorderTests(unittest.TestCase):
self.assertIn("stdout-eof", events)
self.assertIn("wait", events)
def test_tracks_normalized_audio_level_from_pcm_without_exposing_audio(self) -> None:
recorder, process, _factory, _events = self.make_recorder(
pcm=b"\xff\x7f" * (AUDIO_LEVEL_WINDOW_BYTES // 2),
max_bytes=AUDIO_LEVEL_WINDOW_BYTES,
)
recorder.start()
self.assertTrue(process.stdout.eof.wait(1))
self.assertGreater(recorder.audio_level, 0.99)
self.assertLessEqual(recorder.audio_level, 1.0)
recorder.cancel()
def test_audio_level_accumulates_fragmented_reads_into_100ms_windows(self) -> None:
loud_sample = (32767).to_bytes(2, "little", signed=True)
pcm = (
loud_sample * (AUDIO_LEVEL_WINDOW_BYTES // 4)
+ b"\x00" * (AUDIO_LEVEL_WINDOW_BYTES // 2)
+ loud_sample * 400
)
recorder, process, _factory, _events = self.make_recorder(
pcm=pcm,
max_bytes=len(pcm),
stdout_read_size=800,
)
recorder.start()
self.assertTrue(process.stdout.eof.wait(1))
expected_rms = (0.5**0.5) * 32767 / 32768
expected_level = (20 * math.log10(expected_rms) + 60) / 60
self.assertAlmostEqual(recorder.audio_level, expected_level, places=3)
recorder.cancel()
def test_audio_level_maps_silence_known_rms_decay_and_partial_samples(self) -> None:
recorder = PwRecorder()
recorder._update_audio_level(b"\x00\x00" * 16)
self.assertEqual(recorder.audio_level, 0.0)
amplitude = 3277 # approximately -20 dBFS
recorder._update_audio_level(amplitude.to_bytes(2, "little", signed=True) * 16)
self.assertAlmostEqual(recorder.audio_level, 2 / 3, places=3)
recorder._update_audio_level(b"\x00\x00" * 16)
self.assertAlmostEqual(recorder.audio_level, (2 / 3) * 0.65, places=3)
previous = recorder.audio_level
recorder._update_audio_level(b"\xff")
self.assertEqual(recorder.audio_level, previous)
def test_accepts_pw_record_code_one_after_intentional_sigterm(self) -> None:
recorder, process, _factory, _events = self.make_recorder(terminate_returncode=1)
recorder.start()
services/dictate/tests/test_server_core.py
+38 -5
View File
@@ -3,7 +3,7 @@ from collections.abc import Sequence
from dataclasses import dataclass
import unittest
from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, parse_request
from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, PROTOCOL_VERSION, parse_request
from pi_dictate_service.recorder import Capture, RecorderError
from pi_dictate_service.server import DictateServiceCore, Outbound
@@ -73,6 +73,7 @@ class FakeRecorder:
self.started = False
self.cancelled = False
self.capture = Capture(b"\x00\x00\x00\x40", -15, 15, b"")
self.audio_level = 0.0
self.stop_error: RecorderError | None = None
self.unexpected: RecorderError | None = None
@@ -124,7 +125,7 @@ class CoreTestCase(unittest.TestCase):
generation_id: str | None = None,
) -> list[dict[str, object]]:
frame: dict[str, object] = {
"version": 1,
"version": PROTOCOL_VERSION,
"type": request_type,
"requestId": request_id,
}
@@ -217,6 +218,38 @@ class CancellationTests(CoreTestCase):
self.assertEqual(self.core.active_generation_id, None)
class AudioLevelTests(CoreTestCase):
def test_emits_owned_normalized_audio_levels_at_no_more_than_ten_hertz(self) -> None:
recorder = self.load_to_listening()
recorder.audio_level = 0.625
events: list[Outbound] = []
for _ in range(40):
self.clock.advance(0.025)
self.core.poll()
events.extend(self.core.drain_events())
self.assertGreaterEqual(len(events), 9)
self.assertLessEqual(len(events), 10)
self.assertEqual({event.client_id for event in events}, {"a"})
self.assertEqual({event.frame["type"] for event in events}, {"level"})
self.assertEqual({event.frame["level"] for event in events}, {0.625})
def test_clamps_invalid_levels_and_stops_emitting_after_listening(self) -> None:
recorder = self.load_to_listening()
for raw_level, expected in ((2.0, 1.0), (-1.0, 0.0), (float("nan"), 0.0)):
recorder.audio_level = raw_level
self.clock.advance(0.1)
self.core.poll()
level = self.frames(self.core.drain_events(), "a")
self.assertEqual(level[0]["level"], expected)
self.request("a", "stop", "stop", "g1")
self.clock.advance(0.2)
self.core.poll()
self.assertEqual(self.core.drain_events(), [])
class StopAndInferenceTests(CoreTestCase):
def test_stop_drains_then_converts_before_inference_and_emits_result(self) -> None:
self.load_to_listening()
@@ -271,11 +304,11 @@ class StopAndInferenceTests(CoreTestCase):
executor=self.executor,
clock=self.clock,
)
core.handle_request("a", parse_request({"version": 1, "type": "hello", "requestId": "h"}))
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "hello", "requestId": "h"}))
core.drain_events()
core.handle_request("a", parse_request({"version": 1, "type": "start", "requestId": "s", "generationId": "g"}))
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "start", "requestId": "s", "generationId": "g"}))
core.drain_events()
core.handle_request("a", parse_request({"version": 1, "type": "stop", "requestId": "t", "generationId": "g"}))
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "stop", "requestId": "t", "generationId": "g"}))
self.executor.run_next()
core.poll()
core.drain_events()
services/dictate/tests/test_unix_server.py
+14 -14
View File
@@ -73,24 +73,24 @@ class UnixLoopTests(unittest.TestCase):
def test_accepts_fragmented_requests_and_writes_persistent_responses(self) -> None:
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello",')
client.sendall(b'{"version":2,"type":"hello",')
client.sendall(b'"requestId":"hello"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
client.sendall(b'{"version":2,"type":"ping","requestId":"ping"}\n')
pong = self.receive_frame(client)
self.assertEqual(pong["requestId"], "ping")
self.assertEqual(pong["service"], "pi-dictate")
def test_valid_state_error_keeps_open_but_malformed_frame_errors_once_then_closes(self) -> None:
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.receive_frame(client)
client.sendall(
b'{"version":1,"type":"stop","requestId":"stop","generationId":"g"}\n'
b'{"version":2,"type":"stop","requestId":"stop","generationId":"g"}\n'
)
self.assertEqual(self.receive_frame(client)["code"], "invalid_state")
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
client.sendall(b'{"version":2,"type":"ping","requestId":"ping"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")
client.sendall(b'{broken}\n')
@@ -101,9 +101,9 @@ class UnixLoopTests(unittest.TestCase):
def test_malicious_json_values_do_not_crash_the_service(self) -> None:
with self.connect() as client:
client.sendall(b'{"version":1,"type":[],"requestId":"bad-type"}\n')
client.sendall(b'{"version":2,"type":[],"requestId":"bad-type"}\n')
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
client.sendall(b'{"version":1,"type":"ping","requestId":"\\ud800"}\n')
client.sendall(b'{"version":2,"type":"ping","requestId":"\\ud800"}\n')
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
with self.connect() as client:
@@ -115,14 +115,14 @@ class UnixLoopTests(unittest.TestCase):
self.assertEqual(client.recv(1), b"")
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"healthy"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")
def test_closes_non_reading_client_before_output_queue_can_grow_unbounded(self) -> None:
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.receive_frame(client)
ping = b'{"version":1,"type":"ping","requestId":"p"}\n'
ping = b'{"version":2,"type":"ping","requestId":"p"}\n'
try:
client.sendall(ping * 1200)
except BrokenPipeError:
@@ -139,7 +139,7 @@ class UnixLoopTests(unittest.TestCase):
self.assertLessEqual(len(received), 64 * 1024)
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"healthy"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")
def test_rejects_connections_above_the_fixed_limit(self) -> None:
@@ -148,7 +148,7 @@ class UnixLoopTests(unittest.TestCase):
time.sleep(0.1)
overflow = self.connect()
try:
overflow.sendall(b'{"version":1,"type":"hello","requestId":"overflow"}\n')
overflow.sendall(b'{"version":2,"type":"hello","requestId":"overflow"}\n')
with self.assertRaises((ConnectionResetError, BrokenPipeError)):
if overflow.recv(1) == b"":
raise ConnectionResetError
@@ -161,12 +161,12 @@ class UnixLoopTests(unittest.TestCase):
def test_client_reset_does_not_crash_the_service(self) -> None:
reset_client = self.connect()
reset_client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
reset_client.sendall(b'{"version":1,"type":"hello","requestId":"abandoned"}\n')
reset_client.sendall(b'{"version":2,"type":"hello","requestId":"abandoned"}\n')
reset_client.close()
time.sleep(0.15)
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")