diff --git a/packages/sdk-python/pyproject.toml b/packages/sdk-python/pyproject.toml index fb54de7c00..d7519a1a7d 100644 --- a/packages/sdk-python/pyproject.toml +++ b/packages/sdk-python/pyproject.toml @@ -14,6 +14,16 @@ dependencies = ["pydantic>=2.6,<3"] [project.optional-dependencies] gen = ["datamodel-code-generator==0.68.0", "black==26.5.1", "ruff==0.15.21"] dev = ["pytest==9.1.1", "pytest-asyncio==1.4.0", "ruff==0.15.21", "mypy==2.3.0"] +# QVAC-21806 PoC only -- neither package is on PyPI yet (both v0.0.1, git-only). +# Pinned to their current HEAD commits for reproducibility while unreleased. +bare-rpc = [ + "bare-rpc-python @ git+https://github.com/holepunchto/bare-rpc-python.git@5219c528866678553cb9716f80a199c9accfe8c1", + "compact-encoding-python @ git+https://github.com/holepunchto/compact-encoding-python.git@e7e1f080eb413c6ea06f8c6400a9716857c67266", +] + +[tool.hatch.metadata] +# QVAC-21806 PoC: bare-rpc-python/compact-encoding-python aren't on PyPI yet. +allow-direct-references = true [tool.hatch.build.targets.wheel] packages = ["src/qvac"] @@ -49,3 +59,10 @@ module = "qvac._generated.models._internal" # enum member on construction (verified: `Model().field == SomeEnum.x` # is True). Scoped to this one generated module, not project-wide. disable_error_code = ["assignment"] + +[[tool.mypy.overrides]] +module = ["bare_rpc", "compact_encoding"] +# The bare-rpc extra is optional and git-only (no PyPI release, no py.typed +# marker), so it's absent from the default install CI type-checks against. +# bare_rpc_transport.py already guards the import at runtime. +ignore_missing_imports = true diff --git a/packages/sdk-python/src/qvac/bare_rpc_transport.py b/packages/sdk-python/src/qvac/bare_rpc_transport.py new file mode 100644 index 0000000000..b2cee750fe --- /dev/null +++ b/packages/sdk-python/src/qvac/bare_rpc_transport.py @@ -0,0 +1,211 @@ +"""Production `qvac._transport.Transport` implementation, backed by +`bare_rpc.RPC` (github.com/holepunchto/bare-rpc-python). + +Mirrors the JS SDK's `client/rpc/node-rpc-client.ts`: create a Unix +socket / Windows named pipe, spawn the worker with that path (plus +`HOME_DIR`) as its JSON config argument, and wire `bare_rpc.RPC` to the +connection the worker dials back in on. + +Locating the Bare binary and the worker's JS entry point is not this +module's job — `command` is supplied by the caller. This is a thin +client that expects the worker to be available separately; bundling +those artifacts into the installed package is a separate concern. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import tempfile +from collections.abc import AsyncIterable, AsyncIterator, Sequence +from typing import Any + +try: + import bare_rpc +except ImportError: + bare_rpc = None + +BARE_RPC_AVAILABLE = bare_rpc is not None + + +class BareRpcNotInstalledError(ImportError): + def __init__(self) -> None: + super().__init__( + "bare_rpc is not installed -- install the 'bare-rpc' extra " + "(`pip install qvac[bare-rpc]`) to use BareRpcTransport" + ) + + +def _json_or_raise(data: bytes) -> Any: + """Parse a JSON payload; the SDK reports failures in-band as {"type":"error"}.""" + obj = json.loads(data.decode("utf-8")) + if isinstance(obj, dict) and obj.get("type") == "error": + raise RuntimeError("worker: " + str(obj.get("message", "unknown error"))) + return obj + + +class BareRpcTransport: + """Spawns a QVAC SDK worker and speaks to it via `bare_rpc.RPC`, + satisfying `qvac._transport.Transport`'s async call/call_stream/call_duplex + shape directly. + + `command` is the Bare invocation up to (not including) the worker's + JSON config argument, e.g. `["bare", "/path/to/worker.js"]` — this + class appends `{"QVAC_IPC_SOCKET_PATH": ..., "HOME_DIR": ...}` itself, + since that handshake is protocol, not caller, concern. + """ + + def __init__(self, command: Sequence[str], *, home_dir: str | None = None) -> None: + if bare_rpc is None: + raise BareRpcNotInstalledError() + self._command = list(command) + self._home_dir = home_dir or os.path.expanduser("~") + self._sock_path = os.path.join( + tempfile.gettempdir(), f"qvac-worker-{os.getpid()}-{id(self)}.sock" + ) + self._server: asyncio.AbstractServer | None = None + self._proc: asyncio.subprocess.Process | None = None + self._writer: asyncio.StreamWriter | None = None + self._read_task: asyncio.Task | None = None + self.rpc: bare_rpc.RPC | None = None + + async def connect(self, *, timeout: float = 30) -> BareRpcTransport: + if os.path.exists(self._sock_path): + os.unlink(self._sock_path) + + connected: asyncio.Future[None] = asyncio.get_running_loop().create_future() + + async def on_client( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + self._writer = writer + self._read_task = asyncio.current_task() + if not connected.done(): + connected.set_result(None) + # rpc is created right after the worker is spawned below, long + # before the worker can start and dial back into this callback. + assert self.rpc is not None + try: + while True: + chunk = await reader.read(65536) + if not chunk: + break + await self.rpc.receive(chunk) + except asyncio.CancelledError: + pass + + self._server = await asyncio.start_unix_server(on_client, path=self._sock_path) + + config = json.dumps( + {"QVAC_IPC_SOCKET_PATH": self._sock_path, "HOME_DIR": self._home_dir} + ) + self._proc = await asyncio.create_subprocess_exec(*self._command, config) + + def send(frame: bytes) -> None: + # bare_rpc only calls send once it has a connection, at which + # point on_client above has already set the writer. + assert self._writer is not None + self._writer.write(frame) + + self.rpc = bare_rpc.RPC(send=send) + try: + await asyncio.wait_for(connected, timeout=timeout) + except asyncio.TimeoutError: + await self.close() + raise + return self + + async def close(self) -> None: + if self._read_task: + self._read_task.cancel() + if self.rpc: + self.rpc.close() + if self._proc: + self._proc.terminate() + try: + await asyncio.wait_for(self._proc.wait(), timeout=5) + except asyncio.TimeoutError: + self._proc.kill() + if self._writer: + self._writer.close() + if self._server: + self._server.close() + try: + await asyncio.wait_for(self._server.wait_closed(), timeout=5) + except asyncio.TimeoutError: + pass + if os.path.exists(self._sock_path): + os.unlink(self._sock_path) + + async def __aenter__(self) -> BareRpcTransport: + return await self.connect() + + async def __aexit__(self, *exc: object) -> None: + await self.close() + + # ---- Transport protocol ---------------------------------------------- + + def _require_rpc(self) -> Any: + if self.rpc is None: + raise RuntimeError("BareRpcTransport used before connect()") + return self.rpc + + async def call(self, payload: dict) -> dict: + """Unary, via bare_rpc.RPC.request -- no hand-rolled framing at all.""" + data = await self._require_rpc().request( + command=0, data=json.dumps(payload).encode("utf-8") + ) + return _json_or_raise(data) + + async def call_stream(self, payload: dict) -> AsyncIterator[dict]: + """Server-stream, via bare_rpc.RPC.request_with_response_stream.""" + stream = await self._require_rpc().request_with_response_stream( + command=0, data=json.dumps(payload).encode("utf-8") + ) + buffer = "" + async for chunk in stream: + buffer += chunk.decode("utf-8") + lines = buffer.split("\n") + buffer = lines.pop() + for line in lines: + if line.strip(): + yield _json_or_raise(line.encode("utf-8")) + if buffer.strip(): + yield _json_or_raise(buffer.encode("utf-8")) + + async def call_duplex( + self, payload: dict, up: AsyncIterable[bytes] + ) -> AsyncIterator[dict]: + """Duplex, via bare_rpc.RPC.create_bidirectional_stream -- first outgoing + chunk is the JSON payload, then `up`'s chunks; yields parsed response + chunks with the same buffer-and-split-on-newline handling as call_stream.""" + outgoing, incoming = await self._require_rpc().create_bidirectional_stream( + command=0 + ) + await outgoing.write(json.dumps(payload).encode("utf-8")) + + async def _pump_up() -> None: + async for chunk in up: + await outgoing.write(chunk) + await outgoing.end() + + pump_task = asyncio.ensure_future(_pump_up()) + try: + buffer = "" + async for chunk in incoming: + buffer += chunk.decode("utf-8") + lines = buffer.split("\n") + buffer = lines.pop() + for line in lines: + if line.strip(): + yield _json_or_raise(line.encode("utf-8")) + if buffer.strip(): + yield _json_or_raise(buffer.encode("utf-8")) + finally: + if not pump_task.done(): + pump_task.cancel() + try: + await pump_task + except asyncio.CancelledError: + pass diff --git a/packages/sdk-python/tests/test_bare_rpc_transport.py b/packages/sdk-python/tests/test_bare_rpc_transport.py new file mode 100644 index 0000000000..0ebac275e7 --- /dev/null +++ b/packages/sdk-python/tests/test_bare_rpc_transport.py @@ -0,0 +1,266 @@ +"""Same real-worker rigor as test_poc_smoke.py/test_poc_progress.py, but +exercising the production `qvac.bare_rpc_transport.BareRpcTransport` +against a real spawned SDK worker for all three wire call shapes, +including both duplex-shaped methods beyond text-to-speech: parakeet's +`transcribeStream` and the BCI addon's `bciTranscribeStream`. +""" + +from __future__ import annotations + +import array +import asyncio +import os +import wave +from pathlib import Path + +import pytest +import pytest_asyncio + +from qvac.bare_rpc_transport import BARE_RPC_AVAILABLE, BareRpcTransport +from qvac.methods import ( + bci_transcribe_stream, + completion_stream, + heartbeat, + load_model, + text_to_speech_stream, + transcribe_stream, +) +from qvac.models import ( + BCI_WINDOWED, + PARAKEET_CTC_0_6B_Q4_0, + QWEN3_600M_INST_Q4, + TTS_EN_SUPERTONIC_Q4_0, +) +from qvac.schemas import ( + BciTranscribeStreamRequest, + CompletionStreamRequest, + HeartbeatRequest, + LoadModelRequest, + ModelType, + TextToSpeechStreamRequest, + TranscribeStreamRequest, +) + +SDK_DIR = os.environ.get( + "QVAC_POC_SDK_DIR", + str(Path(__file__).resolve().parent.parent.parent / "sdk"), +) +BARE_BIN = f"{SDK_DIR}/node_modules/bare-runtime-darwin-arm64/bin/bare" +WORKER_PATH = f"{SDK_DIR}/dist/server/worker.js" +AUDIO_FIXTURE = f"{SDK_DIR}/e2e/assets/audio/transcription-short-wav.wav" +NEURAL_FIXTURE = f"{SDK_DIR}/e2e/assets/neural/neural-not-too-controversial.bin" + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.skipif( + not BARE_RPC_AVAILABLE, + reason="bare_rpc not installed -- install the 'bare-rpc' extra " + "(`pip install -e '.[bare-rpc]'`) to run these tests", + ), + pytest.mark.skipif( + not os.path.exists(WORKER_PATH), + reason=f"no built SDK worker found at {WORKER_PATH!r} -- run `bun run build` in packages/sdk, or set QVAC_POC_SDK_DIR", + ), +] + + +@pytest_asyncio.fixture +async def transport(): + async with BareRpcTransport([BARE_BIN, WORKER_PATH]) as t: + yield t + + +async def test_heartbeat_unary(transport) -> None: + response = await heartbeat(transport, HeartbeatRequest(type="heartbeat")) + assert response.type == "heartbeat" + assert isinstance(response.number, float) + + +async def test_load_model_and_completion_stream(transport) -> None: + load_request = LoadModelRequest.model_validate( + { + "type": "loadModel", + "modelSrc": QWEN3_600M_INST_Q4.src, + "modelType": "llamacpp-completion", + # Qwen3 is a thinking model: the worker reserves context for the + # reasoning trace, so the metadata-default budget overflows even a + # tiny prompt. Give it an explicit window (matches the SDK e2e). + "modelConfig": {"n_ctx": 2048}, + } + ) + load_response = await load_model(transport, load_request) + assert load_response.success, load_response.error + model_id = load_response.model_id + + completion_request = CompletionStreamRequest.model_validate( + { + "type": "completionStream", + "modelId": model_id, + "history": [{"role": "user", "content": "Say hello in five words."}], + "stream": True, + } + ) + + text = "" + async for chunk in completion_stream(transport, completion_request): + for event in chunk.events: + if event.type == "contentDelta": + text += event.text + assert text.strip(), "expected real completion text via the bare_rpc server-stream" + + +async def _as_async_iter(items): + for item in items: + yield item + + +async def _paced_chunks(chunks, delay_s): + for i, chunk in enumerate(chunks): + yield chunk + if delay_s > 0 and i < len(chunks) - 1: + await asyncio.sleep(delay_s) + + +async def test_load_model_and_tts_stream_duplex(transport) -> None: + """Exercises call_duplex end to end: text goes up the request stream while + synthesized audio comes down the response stream, concurrently.""" + load_request = LoadModelRequest.model_validate( + { + "type": "loadModel", + "modelSrc": TTS_EN_SUPERTONIC_Q4_0.src, + "modelType": ModelType.tts_ggml, + "modelConfig": {"ttsEngine": "supertonic", "language": "en"}, + } + ) + load_response = await load_model(transport, load_request) + assert load_response.success, load_response.error + model_id = load_response.model_id + + tts_request = TextToSpeechStreamRequest.model_validate( + {"type": "textToSpeechStream", "modelId": model_id} + ) + text = b"Hello from QVAC. This is streaming text to speech." + + samples = [] + saw_done = False + async for chunk in text_to_speech_stream( + transport, tts_request, _as_async_iter([text]) + ): + samples.extend(chunk.buffer) + saw_done = saw_done or chunk.done + assert samples, "expected real synthesized audio via the bare_rpc duplex stream" + assert saw_done, "expected a terminal done=True event on the response stream" + + +def _wav_to_s16le_mono_16k(path: str) -> bytes: + """Decode a 16-bit PCM wav to 16 kHz mono s16le -- the wire format parakeet's + duplex `transcribeStream` expects (see the SDK e2e `parakeet-stream-runner.ts`, + which converts its f32 fixture samples to s16le bytes before writing them).""" + with wave.open(path, "rb") as wf: + channels, width, rate = wf.getnchannels(), wf.getsampwidth(), wf.getframerate() + raw = wf.readframes(wf.getnframes()) + if width != 2: + raise RuntimeError(f"expected 16-bit PCM wav, got sampwidth={width}") + samples = array.array("h") + samples.frombytes(raw) + mono = array.array("h", samples[0::channels]) if channels > 1 else samples + target = 16000 + if rate != target: + if rate % target != 0: + raise RuntimeError(f"can't cleanly decimate {rate}Hz to {target}Hz") + mono = array.array("h", mono[0 :: rate // target]) + return mono.tobytes() + + +async def test_transcribe_stream_duplex(transport) -> None: + """Parakeet's streaming session only decodes when fed at roughly real-time + cadence (see transcription-parakeet's live-stream-simulation.test.js / + duplex-streaming tests) -- chunks are paced with a real `asyncio.sleep` + between writes, matching the SDK e2e runner's `writeInChunks(delayMs)`.""" + load_request = LoadModelRequest.model_validate( + { + "type": "loadModel", + "modelSrc": PARAKEET_CTC_0_6B_Q4_0.src, + "modelType": ModelType.parakeet_transcription, + "modelConfig": {}, + } + ) + load_response = await load_model(transport, load_request) + assert load_response.success, load_response.error + model_id = load_response.model_id + + chunk_ms = 1000 + pcm = _wav_to_s16le_mono_16k(AUDIO_FIXTURE) + bytes_per_chunk = int(16000 * chunk_ms / 1000) * 2 + trailing_silence = bytes( + int(16000 * 1.5) * 2 + ) # settle time so the stream finalizes + chunks = [pcm[i : i + bytes_per_chunk] for i in range(0, len(pcm), bytes_per_chunk)] + chunks += [ + trailing_silence[i : i + bytes_per_chunk] + for i in range(0, len(trailing_silence), bytes_per_chunk) + ] + + transcribe_request = TranscribeStreamRequest.model_validate( + { + "type": "transcribeStream", + "modelId": model_id, + "parakeetStreamingConfig": {"chunkMs": chunk_ms, "emitPartials": True}, + } + ) + + text = "" + async for response in transcribe_stream( + transport, transcribe_request, _paced_chunks(chunks, chunk_ms / 1000) + ): + piece = response.text or (response.segment.text if response.segment else None) + if piece: + text += piece + assert text.strip(), "expected real transcript text via the bare_rpc duplex stream" + + +async def test_bci_transcribe_stream_duplex(transport) -> None: + """The BCI addon's sliding-window driver is fed arbitrary-size chunks with + no real-time pacing requirement (see the SDK e2e `bci-executor.ts`, which + writes 64 KiB slices back-to-back) -- unlike parakeet's transcribeStream.""" + load_request = LoadModelRequest.model_validate( + { + "type": "loadModel", + "modelSrc": BCI_WINDOWED.src, + "modelType": ModelType.bci_whispercpp_transcription, + "modelConfig": { + "whisperConfig": {"language": "en", "temperature": 0.0}, + "miscConfig": {"caption_enabled": False}, + # neural-not-too-controversial.bin was recorded on session day 1. + "bciConfig": {"day_idx": 1}, + }, + } + ) + load_response = await load_model(transport, load_request) + assert load_response.success, load_response.error + model_id = load_response.model_id + + bci_request = BciTranscribeStreamRequest.model_validate( + { + "type": "bciTranscribeStream", + "modelId": model_id, + "streamOpts": {"emit": "delta"}, + } + ) + + with open(NEURAL_FIXTURE, "rb") as f: + neural_bytes = f.read() + chunk_size = 64 * 1024 + chunks = [ + neural_bytes[i : i + chunk_size] + for i in range(0, len(neural_bytes), chunk_size) + ] + + text = "" + async for response in bci_transcribe_stream( + transport, bci_request, _as_async_iter(chunks) + ): + piece = response.text or (response.segment.text if response.segment else None) + if piece: + text += piece + assert "controversial" in text.lower(), f"unexpected BCI transcript: {text!r}"