feat(dictate): use local socket-activated Parakeet
This commit is contained in:
@@ -0,0 +1,333 @@
|
||||
from concurrent.futures import Future
|
||||
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.recorder import Capture, RecorderError
|
||||
from pi_dictate_service.server import DictateServiceCore, Outbound
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class PendingWork:
|
||||
future: Future[object]
|
||||
function: object
|
||||
args: tuple[object, ...]
|
||||
|
||||
|
||||
class ControlledExecutor:
|
||||
def __init__(self) -> None:
|
||||
self.pending: list[PendingWork] = []
|
||||
|
||||
def submit(self, function: object, *args: object) -> Future[object]:
|
||||
future: Future[object] = Future()
|
||||
self.pending.append(PendingWork(future, function, args))
|
||||
return future
|
||||
|
||||
def run_next(self) -> None:
|
||||
work = self.pending.pop(0)
|
||||
try:
|
||||
result = work.function(*work.args) # type: ignore[operator]
|
||||
except BaseException as error:
|
||||
work.future.set_exception(error)
|
||||
else:
|
||||
work.future.set_result(result)
|
||||
|
||||
def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self) -> None:
|
||||
self.loaded = False
|
||||
self.transcript = "dictated text"
|
||||
self.load_error: BaseException | None = None
|
||||
self.transcribe_error: BaseException | None = None
|
||||
self.samples: Sequence[float] | None = None
|
||||
|
||||
def load(self) -> None:
|
||||
if self.load_error:
|
||||
raise self.load_error
|
||||
self.loaded = True
|
||||
|
||||
def transcribe(self, samples: Sequence[float]) -> str:
|
||||
self.samples = samples
|
||||
if self.transcribe_error:
|
||||
raise self.transcribe_error
|
||||
return self.transcript
|
||||
|
||||
|
||||
class FakeRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.started = False
|
||||
self.cancelled = False
|
||||
self.capture = Capture(b"\x00\x00\x00\x40", -15, 15, b"")
|
||||
self.stop_error: RecorderError | None = None
|
||||
self.unexpected: RecorderError | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
self.started = True
|
||||
|
||||
def stop(self) -> Capture:
|
||||
if self.stop_error:
|
||||
raise self.stop_error
|
||||
return self.capture
|
||||
|
||||
def cancel(self) -> None:
|
||||
self.cancelled = True
|
||||
|
||||
def poll_unexpected(self) -> RecorderError | None:
|
||||
failure, self.unexpected = self.unexpected, None
|
||||
return failure
|
||||
|
||||
|
||||
class RecorderFactory:
|
||||
def __init__(self) -> None:
|
||||
self.instances: list[FakeRecorder] = []
|
||||
|
||||
def __call__(self) -> FakeRecorder:
|
||||
recorder = FakeRecorder()
|
||||
self.instances.append(recorder)
|
||||
return recorder
|
||||
|
||||
|
||||
class CoreTestCase(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.clock = Clock()
|
||||
self.executor = ControlledExecutor()
|
||||
self.model = FakeModel()
|
||||
self.recorders = RecorderFactory()
|
||||
self.core = DictateServiceCore(
|
||||
model=self.model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
|
||||
def request(
|
||||
self,
|
||||
client: str,
|
||||
request_type: str,
|
||||
request_id: str,
|
||||
generation_id: str | None = None,
|
||||
) -> list[dict[str, object]]:
|
||||
frame: dict[str, object] = {
|
||||
"version": 1,
|
||||
"type": request_type,
|
||||
"requestId": request_id,
|
||||
}
|
||||
if generation_id is not None:
|
||||
frame["generationId"] = generation_id
|
||||
self.core.handle_request(client, parse_request(frame))
|
||||
return self.frames(self.core.drain_events(), client)
|
||||
|
||||
@staticmethod
|
||||
def frames(events: list[Outbound], client: str | None = None) -> list[dict[str, object]]:
|
||||
return [event.frame for event in events if client is None or event.client_id == client]
|
||||
|
||||
def hello(self, client: str = "a") -> None:
|
||||
events = self.request(client, "hello", f"hello-{client}")
|
||||
self.assertEqual(events[0]["type"], "pong")
|
||||
|
||||
def load_to_listening(self, client: str = "a", generation: str = "g1") -> FakeRecorder:
|
||||
self.hello(client)
|
||||
self.request(client, "start", "start", generation)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
events = self.frames(self.core.drain_events(), client)
|
||||
self.assertEqual(events[-1]["state"], "listening")
|
||||
self.assertEqual(events[-1]["requestId"], "start")
|
||||
return self.recorders.instances[-1]
|
||||
|
||||
|
||||
class HandshakeAndOwnershipTests(CoreTestCase):
|
||||
def test_hello_is_required_and_valid_errors_keep_session_open(self) -> None:
|
||||
events = self.request("a", "start", "s1", "g1")
|
||||
self.assertEqual(events[0]["code"], "hello_required")
|
||||
|
||||
self.hello("a")
|
||||
self.request("a", "start", "s2", "g1")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
|
||||
def test_one_generation_is_globally_owned(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.hello("b")
|
||||
|
||||
busy = self.request("b", "start", "s2", "g2")
|
||||
self.assertEqual(busy[0]["code"], "busy")
|
||||
not_owner = self.request("b", "stop", "x", "g1")
|
||||
self.assertEqual(not_owner[0]["code"], "not_owner")
|
||||
stale = self.request("a", "stop", "x2", "old")
|
||||
self.assertEqual(stale[0]["code"], "generation_mismatch")
|
||||
|
||||
stopping = self.request("a", "stop", "x3", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
|
||||
|
||||
class CancellationTests(CoreTestCase):
|
||||
def test_cancel_during_load_is_immediate_and_blocks_until_worker_finishes(self) -> None:
|
||||
self.hello("a")
|
||||
self.request("a", "start", "start", "g1")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.hello("b")
|
||||
self.assertEqual(self.request("b", "start", "start-b", "g2")[0]["code"], "busy")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertEqual(self.recorders.instances, [])
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_cancel_during_inference_suppresses_later_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_disconnect_cancels_capture_and_tombstones_workers(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
|
||||
self.core.disconnect("a")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
self.assertTrue(recorder.cancelled)
|
||||
self.assertEqual(self.core.active_generation_id, None)
|
||||
|
||||
|
||||
class StopAndInferenceTests(CoreTestCase):
|
||||
def test_stop_drains_then_converts_before_inference_and_emits_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
|
||||
stopping = self.request("a", "stop", "stop", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
self.assertEqual(stopping[0]["requestId"], "stop")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(list(self.model.samples or []), [0.0, 0.5])
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "dictated text")
|
||||
self.assertEqual(completed[1]["state"], "idle")
|
||||
|
||||
def test_empty_transcript_is_a_result(self) -> None:
|
||||
self.model.transcript = ""
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "")
|
||||
|
||||
def test_model_and_transcript_failures_are_bounded_terminal_errors(self) -> None:
|
||||
self.model.load_error = RuntimeError("CUDA unavailable")
|
||||
self.hello()
|
||||
self.request("a", "start", "start", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "model_failure")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
model = FakeModel()
|
||||
model.loaded = True
|
||||
model.transcript = "x" * (MAX_TRANSCRIPT_BYTES + 1)
|
||||
core = DictateServiceCore(
|
||||
model=model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
)
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "hello", "requestId": "h"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "start", "requestId": "s", "generationId": "g"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "stop", "requestId": "t", "generationId": "g"}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
core.drain_events()
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
failed = self.frames(core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "transcript_too_large")
|
||||
|
||||
def test_unexpected_recorder_exit_is_terminal(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 3")
|
||||
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
def test_stop_does_not_reclassify_an_already_exited_recorder_as_intentional(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 0")
|
||||
|
||||
failed = self.request("a", "stop", "stop", "g1")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
self.assertEqual(self.executor.pending, [])
|
||||
|
||||
|
||||
class IdleTests(CoreTestCase):
|
||||
def test_only_valid_requests_reset_idle_and_workers_suspend_then_restart_timer(self) -> None:
|
||||
self.core.accepted("a")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "hello", "h")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.request("a", "start", "s", "g")
|
||||
self.clock.advance(901)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "cancel", "c", "g")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.clock.advance(1)
|
||||
self.assertTrue(self.core.should_exit())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user