diff --git a/py/src/braintrust/_audio/__init__.py b/py/src/braintrust/_audio/__init__.py new file mode 100644 index 00000000..d4d03155 --- /dev/null +++ b/py/src/braintrust/_audio/__init__.py @@ -0,0 +1 @@ +"""Internal shared audio infrastructure. No framework or codec dependencies.""" diff --git a/py/src/braintrust/_audio/config.py b/py/src/braintrust/_audio/config.py new file mode 100644 index 00000000..4b9d615b --- /dev/null +++ b/py/src/braintrust/_audio/config.py @@ -0,0 +1,53 @@ +"""Lightweight, framework-independent audio recording configuration. + +Configuration never enables capture. Integrations require their dedicated audio +opt-ins and enforce their framework's privacy policy independently. +""" + +import math +from dataclasses import dataclass +from typing import Literal + + +@dataclass(frozen=True, kw_only=True) +class AudioEncoding: + codec: Literal["opus", "pcm"] = "opus" + container: Literal["ogg", "wav"] | None = "ogg" + bitrate_bps: int = 64000 + complexity: int = 10 + application: Literal["audio", "voip", "lowdelay"] = "audio" + + def __post_init__(self): + if (self.codec, self.container) not in (("opus", "ogg"), ("pcm", None), ("pcm", "wav")): + raise ValueError("Supported audio formats are opus/ogg, pcm, and pcm/wav") + if ( + type(self.bitrate_bps) is not int + or type(self.complexity) is not int + or not 500 <= self.bitrate_bps <= 512000 + or not 0 <= self.complexity <= 10 + ): + raise ValueError("Invalid Opus bitrate or complexity") + if self.application not in ("audio", "voip", "lowdelay"): + raise ValueError("Invalid Opus application") + if self.codec == "pcm" and (self.bitrate_bps, self.complexity, self.application) != (64000, 10, "audio"): + raise ValueError("Opus options cannot be used with PCM") + + +@dataclass(frozen=True, kw_only=True) +class AudioRecordingOptions: + segment_duration_seconds: float = 60 + flush_threshold_bytes: int = 32 * 1024 * 1024 + encoding: AudioEncoding | None = None + max_duration_seconds: float | None = None + shutdown_timeout_seconds: float = 30 + + def __post_init__(self): + for value in (self.segment_duration_seconds, self.shutdown_timeout_seconds, self.max_duration_seconds): + if value is not None and (not math.isfinite(value) or value <= 0): + raise ValueError("Audio durations must be positive and finite") + if self.segment_duration_seconds < 1 / 48000: + raise ValueError("An audio segment must contain at least one sample") + if type(self.flush_threshold_bytes) is not int or self.flush_threshold_bytes <= 0: + raise ValueError("Audio flush threshold must be positive") + if self.encoding is not None and not isinstance(self.encoding, AudioEncoding): + raise TypeError("Audio encoding must be an AudioEncoding instance") diff --git a/py/src/braintrust/_audio/middleware.py b/py/src/braintrust/_audio/middleware.py new file mode 100644 index 00000000..e964be2a --- /dev/null +++ b/py/src/braintrust/_audio/middleware.py @@ -0,0 +1,110 @@ +import importlib +import io +import logging +import wave +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import replace +from fractions import Fraction +from typing import Any, Protocol + +from braintrust._audio.config import AudioEncoding + +from .types import AudioResult, SegmentMeta +from .worker import run_native + + +AudioNext = Callable[[Any, SegmentMeta], Awaitable[AudioResult]] + + +class AudioMiddleware(Protocol): + async def __call__(self, payload: Any, meta: SegmentMeta, next: AudioNext) -> AudioResult: ... + + +def _validate(meta: SegmentMeta, original: SegmentMeta): + if not isinstance(meta, SegmentMeta) or ( + meta.index != original.index + or Fraction(meta.start_sample, meta.sample_rate_hz) != Fraction(original.start_sample, original.sample_rate_hz) + or Fraction(meta.sample_count, meta.sample_rate_hz) != Fraction(original.sample_count, original.sample_rate_hz) + or meta.channels != original.channels + ): + raise ValueError("Audio middleware changed segment identity, placement, duration, or sources") + if ( + meta.encoding == "pcm" + and meta.mime_type == "audio/pcm" + and ( + meta.attributes.get("pcm.sample_format") != "s16le" + or meta.attributes.get("pcm.channel_layout") != "interleaved" + ) + ): + raise ValueError("Raw PCM requires sample format and channel layout") + + +async def run_middleware(stack: Sequence[AudioMiddleware], payload: bytes, meta: SegmentMeta) -> AudioResult: + async def run(index, data, current): + _validate(current, meta) + if index == len(stack): + return AudioResult(data, current) + called = False + + def next(data, new_meta): + nonlocal called + if called: + raise ValueError("Audio middleware called next more than once") + called = True + return run(index + 1, data, new_meta) + + return await stack[index](data, current, next) + + result = await run(0, payload, meta) + if not isinstance(result, AudioResult) or not isinstance(result.data, bytes): + raise ValueError("Audio middleware must return final bytes") + _validate(result.meta, meta) + if result.meta.encoding == "pcm" and result.meta.mime_type == "audio/pcm": + if len(result.data) != result.meta.sample_count * len(result.meta.channels) * 2: + raise ValueError("PCM byte length does not match playable duration") + return result + + +def _wav(data, meta): + buffer = io.BytesIO() + with wave.Wave_write(buffer) as output: + output.setnchannels(len(meta.channels)) + output.setsampwidth(2) + output.setframerate(meta.sample_rate_hz) + output.writeframes(data) + return buffer.getvalue() + + +async def wav_middleware(data, meta, next): + data = await run_native(_wav, data, meta) + return await next(data, replace(meta, mime_type="audio/wav")) + + +_warned_missing_extension = False + + +def encoding_stack(config: AudioEncoding | None, *, channels: int | None = None) -> tuple[AudioMiddleware, ...]: + global _warned_missing_extension + selected = config or AudioEncoding() + if selected.codec == "pcm": + return (wav_middleware,) if selected.container == "wav" else () + if channels is not None: + if channels not in (1, 2): + raise ValueError("Opus/Ogg supports mono or stereo recordings") + if selected.bitrate_bps > 256000 * channels: + raise ValueError("Opus bitrate exceeds 256000 bits/s per output channel") + try: + extension = importlib.import_module("braintrust_audio") + stack = tuple(extension.create_stack(selected, contract_version=1)) + if not all(callable(layer) for layer in stack): + raise ValueError("Invalid audio extension middleware contract") + return stack + except Exception as error: + if config is not None: + raise ValueError("Opus/Ogg requires a compatible braintrust-audio installation") from error + if not _warned_missing_extension: + _warned_missing_extension = True + logging.getLogger("braintrust").warning( + "braintrust-audio is unavailable; recording raw PCM. Install braintrust-audio for Opus/Ogg." + ) + return () diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py new file mode 100644 index 00000000..44d80e39 --- /dev/null +++ b/py/src/braintrust/_audio/recorder.py @@ -0,0 +1,315 @@ +import asyncio +import logging +import math +import sys +from array import array +from bisect import insort_right +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from fractions import Fraction + +from braintrust._audio.config import AudioRecordingOptions + +from .middleware import AudioMiddleware, run_middleware +from .types import AudioChannel, AudioResult, RecordingResult, SegmentMeta +from .worker import run_native + + +SAMPLE_RATE = 48000 +NS = 1_000_000_000 + + +@dataclass(frozen=True) +class Packet: + source: str + data: bytes + rate: int + channels: int + start: Fraction # Recording samples; retaining the fraction preserves run phase. + + @property + def count(self): + return len(self.data) // (2 * self.channels) + + @property + def end(self): + return self.start + Fraction(self.count * SAMPLE_RATE, self.rate) + + +def _samples(data): + result = array("h") + result.frombytes(data) + if sys.byteorder != "little": + result.byteswap() + return result + + +def _continues(packet: Packet, following: Packet | None): + return following is not None and ( + following.start == packet.end and following.rate == packet.rate and following.channels == packet.channels + ) + + +def render(packets: tuple[Packet, ...], sources: tuple[str, ...], start: int, end: int): + """Sparse placement and phase-continuous linear resampling, in a worker. + + Phase is derived from absolute rational positions, never rounded per packet. + A following packet supplies the interpolation sample at a packet boundary. + Silence is materialized for only the segment being processed. + """ + output = array("h", [0]) * ((end - start) * len(sources)) + represented = [] + for channel, source in enumerate(sources): + stream = [packet for packet in packets if packet.source == source] + present = False + for index, packet in enumerate(stream): + first = max(start, math.ceil(packet.start)) + last = min(end, math.ceil(packet.end)) + if first >= last: + continue + present = True + samples = _samples(packet.data) + # Downmix before interpolation. The next packet belongs to the same + # run only when its position, rate and channel count agree exactly. + mono = [ + sum(samples[i : i + packet.channels]) // packet.channels + for i in range(0, len(samples), packet.channels) + ] + following = mono[-1] + if index + 1 < len(stream): + other = stream[index + 1] + if _continues(packet, other): + frame = _samples(other.data[: other.channels * 2]) + following = sum(frame) // other.channels + phase = (Fraction(first) - packet.start) * packet.rate / SAMPLE_RATE + step = Fraction(packet.rate, SAMPLE_RATE) + denominator = math.lcm(phase.denominator, step.denominator) + numerator = phase.numerator * (denominator // phase.denominator) + increment = step.numerator * (denominator // step.denominator) + for position in range(first, last): + sample, remainder = divmod(numerator, denominator) + value = mono[sample] + next_value = mono[sample + 1] if sample + 1 < len(mono) else following + output[(position - start) * len(sources) + channel] = ( + value + (next_value - value) * remainder // denominator + ) + numerator += increment + represented.append(AudioChannel((source,) if present else ())) + if sys.byteorder != "little": + output.byteswap() + return output.tobytes(), tuple(represented) + + +class AudioRecorder: + """Accept positioned PCM without waiting for normalization, encoding or IO. + + ``advance`` is a promise by the adapter: no source will add audio before its + watermark. Thresholds request flushing; they never reject admitted samples. + """ + + def __init__( + self, + *, + sources: tuple[str, ...], + options: AudioRecordingOptions, + middleware: tuple[AudioMiddleware, ...], + on_segment: Callable[[AudioResult], Awaitable[None]], + on_omitted: Callable[[SegmentMeta, str], None], + ): + if not sources or len(set(sources)) != len(sources): + raise ValueError("Recording sources must be nonempty and unique") + self.sources = tuple(sources) + self.options = options + self.middleware = middleware + self._limit = ( + round(options.max_duration_seconds * SAMPLE_RATE) if options.max_duration_seconds is not None else None + ) + self.on_segment = on_segment + self.on_omitted = on_omitted + self._packets: list[Packet] = [] + self._watermark = 0 + self._sealed = 0 + self._end = 0 + self._index = 0 + self._task: asyncio.Task | None = None + self._finishing = False + self._discarded = False + self.reason: str | None = None + self._warned_backlog = False + self._processing: SegmentMeta | None = None + self._gaps: list[tuple[int, int, str]] = [] + + @property + def buffered_bytes(self): + return sum(len(packet.data) for packet in self._packets) + + def record( + self, source_id: str, pcm: bytes, *, sample_rate_hz: int, channels: int, offset_ns: int | Fraction + ) -> bool: + if self._finishing or self.reason: + return False + if source_id not in self.sources or sample_rate_hz <= 0 or channels <= 0 or len(pcm) % (channels * 2): + raise ValueError("Invalid positioned PCM") + if not pcm: + return True + start = Fraction(offset_ns) * SAMPLE_RATE / NS + if start < self._sealed or start < 0: + raise ValueError("Audio arrived before the sealed recording watermark") + if self._limit is not None and start >= self._limit: + return False + packet = Packet(source_id, bytes(pcm), sample_rate_hz, channels, start) + # Rendering and flush eligibility both need each source's lookahead in + # timeline order, even when packets arrive out of order. + insort_right(self._packets, packet, key=lambda packet: packet.start) + self._end = max(self._end, math.ceil(packet.end)) + if self._limit is not None: + self._end = min(self._end, self._limit) + buffered = self.buffered_bytes + if buffered > self.options.flush_threshold_bytes * 4 and not self._warned_backlog: + self._warned_backlog = True + logging.getLogger("braintrust").warning( + "Audio recording backlog exceeds its flush target; capture continues" + ) + self._pump() + return True + + def advance(self, through_ns: int | Fraction): + self._watermark = max(self._watermark, math.floor(Fraction(through_ns) * SAMPLE_RATE / NS)) + # Only the adapter's watermark establishes that every source has had a + # chance to deliver its in-window audio. + if self._limit is not None and self._watermark >= self._limit: + self.stop("duration_limit") + self._pump() + + def unavailable(self, start_ns, end_ns, reason): + """Known bounds of unavailable audio; never encode it as real silence.""" + start = max(0, math.floor(Fraction(start_ns) * SAMPLE_RATE / NS)) + end = math.ceil(Fraction(end_ns) * SAMPLE_RATE / NS) + if self._limit is not None: + end = min(end, self._limit) + if end <= start or self._discarded: + return + if start < self._sealed: + raise ValueError("Unavailable audio arrived before the sealed watermark") + self._gaps.append((start, end, reason)) + self._end = max(self._end, end) + self._pump() + + def stop(self, reason: str): + self.reason = self.reason or reason + + def discard(self, reason: str): + """Stop and erase queued PCM; in-flight workers retain their own copies.""" + if self._discarded: + return + self.stop(reason) + self._discarded = True + if self._processing is not None: + self.on_omitted(self._processing, reason) + if self._end > self._sealed: + self.on_omitted( + SegmentMeta( + self._index, + self._sealed, + self._end - self._sealed, + SAMPLE_RATE, + tuple(AudioChannel(()) for _ in self.sources), + ), + reason, + ) + self._index += 1 + self._packets.clear() + + def _cut(self): + eligible = min(self._end, self._watermark) + if not self._finishing and self.reason != "duration_limit": + following: dict[str, Packet] = {} + for packet in reversed(self._packets): + if self._watermark <= packet.end and not _continues(packet, following.get(packet.source)): + # A packet starting exactly at the watermark may still + # arrive. Keep its predecessor's interpolation tail until + # that lookahead arrives, the watermark passes the boundary, + # or capture is finalized. Partial flushes retain the packet + # because its end remains beyond the sealed position. + last_sample = packet.start + Fraction((packet.count - 1) * SAMPLE_RATE, packet.rate) + eligible = min(eligible, math.floor(last_sample) + 1) + following[packet.source] = packet + target = self._sealed + max(1, round(self.options.segment_duration_seconds * SAMPLE_RATE)) + if eligible <= self._sealed: + return None + if eligible >= target: + return target + if self._finishing or self.buffered_bytes >= self.options.flush_threshold_bytes: + return eligible + return None + + def _pump(self): + if self._task is None or self._task.done(): + if self._cut() is not None: + self._task = asyncio.create_task(self._drain()) + + async def _drain(self): + while not self._discarded and (end := self._cut()) is not None: + start = self._sealed + meta = SegmentMeta( + self._index, start, end - start, SAMPLE_RATE, tuple(AudioChannel(()) for _ in self.sources) + ) + self._index += 1 + self._processing = meta + # Snapshot references before detaching accepted buffers. Native work + # owns them even if the awaiting coroutine is cancelled. + packets = tuple(self._packets) + self._sealed = end + gap = next((gap for gap in self._gaps if gap[0] < end and gap[1] > start), None) + if gap is not None: + self.on_omitted(meta, gap[2]) + self._packets = [packet for packet in self._packets if packet.end > end] + self._gaps = [gap for gap in self._gaps if gap[1] > end] + self._processing = None + continue + try: + data, channels = await run_native(render, packets, self.sources, start, end) + meta = SegmentMeta(meta.index, start, end - start, SAMPLE_RATE, channels) + result = await run_middleware(self.middleware, data, meta) + if not self._discarded: + await self.on_segment(result) + else: + self.on_omitted(meta, self.reason or "recording_stopped") + except asyncio.CancelledError: + self.on_omitted(meta, "shutdown_timeout") + raise + except Exception: + self.stop("encoding_or_upload_failed") + self.on_omitted(meta, self.reason) + # Already accepted audio is unavailable; retain honest placement + # rather than claim the remainder of the call was recorded. + if self._end > end: + self.on_omitted( + SegmentMeta(self._index, end, self._end - end, SAMPLE_RATE, meta.channels), self.reason + ) + self._index += 1 + self._packets.clear() + self._sealed = self._end + self._processing = None + self._discarded = True + return + self._packets = [packet for packet in self._packets if packet.end > end] + self._processing = None + + async def abort(self, reason: str) -> None: + """Discard queued audio and wait for the background task to cancel.""" + self.discard(reason) + if self._task is not None: + self._task.cancel() + await asyncio.gather(self._task, return_exceptions=True) + + async def finish(self) -> RecordingResult: + self._finishing = True + # The adapter settles pending playback before calling finish. + self._watermark = max(self._watermark, self._end) + if self._limit is not None and self._end >= self._limit: + self.stop("duration_limit") + self._pump() + if self._task is not None: + await asyncio.shield(self._task) + return RecordingResult(self._index, self.reason) diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py new file mode 100644 index 00000000..970b1045 --- /dev/null +++ b/py/src/braintrust/_audio/session.py @@ -0,0 +1,202 @@ +import asyncio +import logging +import time +from collections.abc import Callable + +from braintrust._audio.config import AudioRecordingOptions +from braintrust.env import BraintrustEnv +from braintrust.logger import Attachment + +from .middleware import encoding_stack +from .recorder import AudioRecorder +from .types import AudioResult, RecordingResult, SegmentMeta +from .worker import run_native + + +class AudioSessionRecorder: + """Schema registries and recorder-local direct attachment export. + + ``allowed`` must include capture opt-ins and effective framework privacy. It + is checked again before upload. Trace customizers cannot veto audio uploads. + """ + + def __init__( + self, + span, + *, + sources: tuple[str, ...] = (), + options: AudioRecordingOptions | None = None, + allowed: Callable[[], bool] = lambda: True, + ): + self.span = span + self.wall_ns, self.monotonic_ns = time.time_ns(), time.monotonic_ns() + self.allowed = allowed + self.options = options or AudioRecordingOptions() + self._participants: dict[object, str] = {} + self._speakers: dict[object, str] = {} + self._terminal: set[int] = set() + self._inflight: dict[int, SegmentMeta] = {} + self._closed = False + self._sources = {source: {} for source in sources} + self.recorder = ( + AudioRecorder( + sources=sources, + options=self.options, + middleware=encoding_stack(self.options.encoding, channels=len(sources)), + on_segment=self._export, + on_omitted=self.omitted, + ) + if sources + else None + ) + self.span.log(metadata=self.session_metadata()) + + def session_metadata(self): + return { + "audio.schema.version": "3.0.0", + "audio.timeline": {"origin_unix_ms": self.wall_ns // 1_000_000}, + "audio.sources": dict(self._sources), + } + + def source(self, source_id, member=None): + if source_id not in self._sources: + self._sources[source_id] = dict(member or {}) + self.span.log(metadata={"audio.sources": {source_id: self._sources[source_id]}}) + return source_id + + def elapsed_ns(self): + return time.monotonic_ns() - self.monotonic_ns + + def from_wall(self, seconds: float): + return max(0, round(seconds * 1_000_000_000) - self.wall_ns) + + def participant(self, native_key, member): + return self._identity(self._participants, native_key, member, "audio.participants", "participant") + + def speaker(self, source_id, native_key, member): + return self._identity(self._speakers, (source_id, native_key), member, "audio.speakers", "speaker") + + def _identity(self, registry, key, member, field, prefix): + if not self.allowed(): + return None + if key not in registry: + identifier = f"{prefix}-{len(registry)}" + registry[key] = identifier + self.span.log(metadata={field: {identifier: member}}) + return registry[key] + + def record(self, source_id, pcm, **format_and_time): + if self.recorder is None or not self.allowed(): + if self.recorder is not None: + self.recorder.discard("privacy_restricted") + return False + return self.recorder.record(source_id, pcm, **format_and_time) + + def advance(self, through_ns): + if self.recorder is not None: + self.recorder.advance(through_ns) + + def unavailable(self, start_ns, end_ns, reason): + if self.recorder is not None: + self.recorder.unavailable(start_ns, end_ns, reason) + + def omitted(self, meta: SegmentMeta, reason: str): + if meta.index in self._terminal: + return + self._terminal.add(meta.index) + self._inflight.pop(meta.index, None) + self.span.log( + metadata={ + "audio.segments": { + f"part-{meta.index}": { + "state": "omitted", + **meta.placement(), + "reason": reason, + } + } + } + ) + + async def _export(self, result: AudioResult): + meta = result.meta + self._inflight[meta.index] = meta + identifier = f"part-{meta.index}" + if self._closed or not self.allowed(): + self.omitted(meta, "privacy_restricted" if not self._closed else "shutdown_timeout") + return + self.span.log(metadata={"audio.segments": {identifier: {"state": "pending", **meta.placement()}}}) + extension = {"audio/ogg": "ogg", "audio/wav": "wav", "audio/pcm": "pcm"}[meta.mime_type] + attachment = Attachment(data=result.data, filename=f"{identifier}.{extension}", content_type=meta.mime_type) + attachment._upload_timeout_seconds = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + + def upload(): + # Recheck after waiting for a native worker slot. + if self._closed or not self.allowed(): + return None + return attachment.upload() + + try: + status = await run_native(upload) + except Exception: + self.omitted(meta, "upload_failed") + raise + if meta.index in self._terminal: + return + if status is None or not self.allowed(): + self.omitted(meta, "privacy_restricted") + elif status.get("upload_status") != "done": + self.omitted(meta, "upload_failed") + raise RuntimeError("Audio upload failed") + else: + self._terminal.add(meta.index) + self._inflight.pop(meta.index, None) + self.span.log( + input={"audio": {identifier: attachment.reference}}, + metadata={ + "audio.segments": { + identifier: { + "state": "ready", + **meta.placement(), + "attachment": {"ref": f"/input/audio/{identifier}"}, + "mime_type": meta.mime_type, + "encoding": meta.encoding, + "sample_rate_hz": meta.sample_rate_hz, + "channels": { + str(i): {"sources": channel.sources} for i, channel in enumerate(meta.channels) + }, + "attributes": meta.attributes, + } + } + }, + ) + + async def finish(self) -> RecordingResult: + result = RecordingResult(0) + loop = asyncio.get_running_loop() + deadline = loop.time() + self.options.shutdown_timeout_seconds + try: + if self.recorder is not None: + result = await asyncio.wait_for(self.recorder.finish(), self.options.shutdown_timeout_seconds) + except (asyncio.TimeoutError, asyncio.CancelledError) as error: + self._closed = True + for meta in tuple(self._inflight.values()): + self.omitted(meta, "shutdown_timeout") + if self.recorder is not None: + await self.recorder.abort("shutdown_timeout") + result = RecordingResult(len(self._terminal), "shutdown_timeout") + if isinstance(error, asyncio.CancelledError): + raise + finally: + self._closed = True + if result.reason: + self.span.log(metadata={"contrib.braintrust.recording_stop_reason": result.reason}) + # Terminal updates are queued before flushing. No exporter changes. + flushing = loop.run_in_executor(None, self.span.flush) + flushing.add_done_callback(lambda completed: completed.exception() if not completed.cancelled() else None) + try: + await asyncio.wait_for(asyncio.shield(flushing), max(0, deadline - loop.time())) + except asyncio.TimeoutError: + result = RecordingResult(result.segments, result.reason or "logger_flush_timeout") + except Exception: + logging.getLogger("braintrust").warning("Audio terminal metadata flush failed") + return result diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py new file mode 100644 index 00000000..76b3c030 --- /dev/null +++ b/py/src/braintrust/_audio/test_audio.py @@ -0,0 +1,437 @@ +import asyncio +import io +import subprocess +import sys +import threading +import wave +from array import array +from dataclasses import replace +from fractions import Fraction +from itertools import permutations +from types import MappingProxyType + +import pytest +from braintrust._audio.config import AudioEncoding, AudioRecordingOptions +from braintrust._audio.middleware import encoding_stack, run_middleware +from braintrust._audio.recorder import AudioRecorder +from braintrust._audio.types import AudioChannel, AudioResult, SegmentMeta +from braintrust._audio.worker import run_native + + +def pcm(values): + return array("h", values).tobytes() + + +def recorder(*, duration=0.01, threshold=1024 * 1024, middleware=(), on_segment=None, max_duration=None): + segments, omitted = [], [] + + async def collect(result): + segments.append(result) + if on_segment is not None: + await on_segment(result) + + instance = AudioRecorder( + sources=("input", "output"), + options=AudioRecordingOptions( + segment_duration_seconds=duration, + flush_threshold_bytes=threshold, + max_duration_seconds=max_duration, + ), + middleware=middleware, + on_segment=collect, + on_omitted=lambda *args: omitted.append(args), + ) + return instance, segments, omitted + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rate", [16000, 24000, 44100, 48000]) +async def test_packet_splits_keep_exact_phase_and_length(rate): + async def record(split): + instance, segments, omitted = recorder() + count = rate // 5 + values = [i % 1000 for i in range(count)] + for start in range(0, count, split): + instance.record( + "input", + pcm(values[start : start + split]), + sample_rate_hz=rate, + channels=1, + offset_ns=Fraction(start * 1_000_000_000, rate), + ) + instance.advance(200_000_000) + await instance.finish() + assert not omitted + assert sum(s.meta.sample_count for s in segments) == 9600 + return b"".join(s.data for s in segments) + + assert await record(137) == await record(rate) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("advance", [False, True]) +@pytest.mark.parametrize("count", [480, 960]) +async def test_duration_limit_waits_for_all_sources(advance, count): + instance, segments, omitted = recorder(max_duration=0.01) + assert instance.record("input", pcm([100] * count), sample_rate_hz=48000, channels=1, offset_ns=0) + assert instance.record("output", pcm([200] * 480), sample_rate_hz=48000, channels=1, offset_ns=0) + assert not instance.record("output", pcm([300] * 480), sample_rate_hz=48000, channels=1, offset_ns=10_000_000) + if advance: + instance.advance(10_000_000) + assert instance.reason == "duration_limit" + assert not instance.record("input", pcm([400]), sample_rate_hz=48000, channels=1, offset_ns=10_000_000) + result = await instance.finish() + samples = array("h") + samples.frombytes(b"".join(s.data for s in segments)) + assert samples == array("h", [100, 200] * 480) + assert result.reason == "duration_limit" + assert not omitted + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rate", [16000, 24000, 44100, 48000]) +@pytest.mark.parametrize("threshold", [1, 1024 * 1024]) +async def test_progressive_resampling_waits_for_packet_lookahead(rate, threshold): + async def record(progressive): + instance, segments, omitted = recorder(threshold=threshold) + for index, value in enumerate((0, 10000, -10000)): + assert instance.record( + "input", + pcm([value] * (rate // 100)), + sample_rate_hz=rate, + channels=1, + offset_ns=index * 10_000_000, + ) + if progressive: + instance.advance((index + 1) * 10_000_000) + if instance._task is not None: + await instance._task + if index == 1: + assert segments # Lookahead releases audio before finish. + await instance.finish() + assert not omitted + return b"".join(s.data for s in segments) + + assert await record(True) == await record(False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("order", list(permutations(range(3)))) +@pytest.mark.parametrize("progressive", [False, True]) +async def test_resampling_is_independent_of_packet_arrival_order(order, progressive): + async def record(indices): + instance, segments, omitted = recorder() + for index in indices: + assert instance.record( + "input", + pcm([index * 10000] * 240), + sample_rate_hz=24000, + channels=1, + offset_ns=index * 10_000_000, + ) + if progressive: + instance.advance(20_000_000) + await instance._task + assert len(segments) == 2 # Both boundaries have contiguous lookahead. + await instance.finish() + assert not omitted + return b"".join(s.data for s in segments) + + expected = await record(range(3)) + samples = array("h") + samples.frombytes(expected) + assert samples[959 * 2] == 15000 + assert await record(order) == expected + + +@pytest.mark.asyncio +async def test_resampling_tail_released_when_watermark_passes_packet_end(): + instance, segments, omitted = recorder() + instance.record("input", pcm([100] * 240), sample_rate_hz=24000, channels=1, offset_ns=0) + instance.advance(10_000_000) + if instance._task is not None: + await instance._task + assert not segments + # No contiguous packet can arrive once the watermark passes its start. + instance.advance(10_100_000) + await instance._task + assert len(segments) == 1 + samples = array("h") + samples.frombytes(segments[0].data) + assert samples == array("h", [100, 0] * 480) + await instance.finish() + assert not omitted + + +@pytest.mark.asyncio +async def test_unavailable_audio_is_clamped_to_duration_limit(): + instance, segments, omitted = recorder(max_duration=0.01) + instance.unavailable(0, 20_000_000, "unavailable") + instance.unavailable(20_000_000, 30_000_000, "unavailable") + result = await instance.finish() + assert not segments + assert len(omitted) == 1 + assert omitted[0][0].sample_count == 480 + assert result.reason == "duration_limit" + + +@pytest.mark.asyncio +async def test_sparse_timeline_downmix_and_channel_alignment(): + instance, segments, omitted = recorder(duration=0.1) + instance.record("input", pcm([200, 400] * 480), sample_rate_hz=48000, channels=2, offset_ns=5_000_000) + instance.record("output", pcm([800] * 240), sample_rate_hz=24000, channels=1, offset_ns=20_000_000) + await instance.finish() + samples = array("h") + samples.frombytes(segments[0].data) + assert len(samples) == 1440 * 2 + assert samples[240 * 2 : 720 * 2 : 2] == array("h", [300] * 480) + assert samples[960 * 2 + 1 :: 2] == array("h", [800] * 480) + assert samples[: 240 * 2] == array("h", [0] * 480) + assert not omitted + + +@pytest.mark.asyncio +async def test_size_threshold_waits_for_watermark_and_congestion_preserves_samples(): + gate = asyncio.Event() + + async def slow_upload(_): + await gate.wait() + + instance, segments, omitted = recorder(duration=60, threshold=32, on_segment=slow_upload) + for index in range(10): + instance.record("input", pcm([index] * 480), sample_rate_hz=48000, channels=1, offset_ns=index * 10_000_000) + await asyncio.sleep(0) + assert not segments + instance.advance(50_000_000) + for _ in range(50): + if segments: + break + await asyncio.sleep(0.005) + assert len(segments) == 1 + for index in range(10, 20): + instance.record("input", pcm([index] * 480), sample_rate_hz=48000, channels=1, offset_ns=index * 10_000_000) + instance.advance(200_000_000) + gate.set() + await instance.finish() + samples = array("h") + samples.frombytes(b"".join(s.data for s in segments)) + assert samples[::2] == array("h", (value for index in range(20) for value in [index] * 480)) + assert not omitted + + +@pytest.mark.parametrize("mapping", [dict, MappingProxyType]) +def test_segment_attributes_are_reused_and_isolated_from_caller_mutations(mapping): + attributes = {"nested": {"list": [1]}} + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),), attributes=mapping(attributes)) + changed = replace(meta, mime_type="audio/wav") + assert changed.attributes is meta.attributes + + attributes["nested"]["list"].append(2) + attributes["extra"] = True + assert meta.attributes == {"nested": {"list": (1,)}} + with pytest.raises(TypeError): + meta.attributes["extra"] = True + with pytest.raises(TypeError): + meta.attributes["nested"]["extra"] = True + with pytest.raises(TypeError): + meta.attributes["nested"]["list"][0] = 2 + + replacement = {"nested": {"list": [3]}} + changed = replace(meta, attributes=mapping(replacement)) + replacement["nested"]["list"].append(4) + assert changed.attributes is not meta.attributes + assert changed.attributes == {"nested": {"list": (3,)}} + + +@pytest.mark.asyncio +async def test_middleware_rejects_calling_next_twice(): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) + + async def twice(data, current, next): + await next(data, current) + return await next(data, current) + + with pytest.raises(ValueError, match="more than once"): + await run_middleware((twice,), bytes(96), meta) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("data", "sample_count", "message"), + [ + (None, 48, "final bytes"), + (bytearray(96), 48, "final bytes"), + (bytes(94), 48, "PCM byte length"), + (bytes(96), 49, "identity, placement"), + ], +) +async def test_middleware_validates_final_result_after_nested_layers(data, sample_count, message): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) + + async def passthrough(data, current, next): + return await next(data, current) + + async def invalid(payload, current, next): + return None if data is None else AudioResult(data, replace(current, sample_count=sample_count)) + + with pytest.raises(ValueError, match=message): + await run_middleware((passthrough, passthrough, invalid), bytes(96), meta) + + +@pytest.mark.asyncio +async def test_middleware_rejects_invalid_metadata_before_calling_downstream(): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) + + async def invalid(data, current, next): + return await next(data, replace(current, sample_count=49)) + + async def downstream(data, current, next): + pytest.fail("Invalid segment metadata reached downstream middleware") + + with pytest.raises(ValueError, match="identity, placement"): + await run_middleware((invalid, downstream), bytes(96), meta) + + +@pytest.mark.asyncio +async def test_wav_requires_no_optional_backend_and_preserves_tiny_tail(): + meta = SegmentMeta(0, 0, 137, 48000, (AudioChannel(("input",)),)) + result = await run_middleware(encoding_stack(AudioEncoding(codec="pcm", container="wav")), pcm(range(137)), meta) + with wave.open(io.BytesIO(result.data)) as decoded: + assert decoded.getnframes() == 137 + assert decoded.readframes(137) == pcm(range(137)) + + +@pytest.mark.asyncio +async def test_native_work_completes_after_awaiter_cancellation(): + started, release, completed = threading.Event(), threading.Event(), threading.Event() + payload = bytearray(b"owned") + + def native(data): + started.set() + release.wait(5) + assert data == b"owned" + completed.set() + + task = asyncio.create_task(run_native(native, payload)) + await asyncio.to_thread(started.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + assert await asyncio.to_thread(completed.wait, 5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("separate_loop", [False, True]) +async def test_native_worker_admits_waiting_sessions_in_order_and_removes_cancelled_waiters(separate_loop): + started = [threading.Event() for _ in range(4)] + release = [threading.Event() for _ in range(3)] + loop = asyncio.get_running_loop() + queued = asyncio.Event() + + async def first_waiter(): + task = asyncio.create_task(run_native(blocking, 2)) + # Notify after the task has had a turn to request admission on its loop. + asyncio.get_running_loop().call_soon(loop.call_soon_threadsafe, queued.set) + return await task + + def blocking(index): + started[index].set() + assert release[index].wait(5) + + tasks = [asyncio.create_task(run_native(blocking, index)) for index in range(2)] + try: + for event in started[:2]: + assert await asyncio.to_thread(event.wait, 5) + first = asyncio.create_task( + asyncio.to_thread(asyncio.run, first_waiter()) if separate_loop else first_waiter() + ) + tasks.append(first) + await asyncio.wait_for(queued.wait(), 5) + cancelled = asyncio.create_task(run_native(lambda: pytest.fail("Cancelled job was submitted"))) + last = asyncio.create_task(run_native(started[3].set)) + tasks.extend([cancelled, last]) + + queue_ready = asyncio.Event() + loop.call_soon(queue_ready.set) + await asyncio.wait_for(queue_ready.wait(), 5) + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + release[0].set() + assert await asyncio.to_thread(started[2].wait, 5) + assert not started[3].is_set() + release[2].set() + assert await asyncio.to_thread(started[3].wait, 5) + finally: + for event in release: + event.set() + await asyncio.gather(*tasks, return_exceptions=True) + + +def test_options_reject_unsupported_configuration(): + with pytest.raises(ValueError): + AudioEncoding(codec="opus/ogg") + with pytest.raises(ValueError): + AudioRecordingOptions(segment_duration_seconds=float("nan")) + + +@pytest.mark.asyncio +async def test_unavailable_audio_is_omitted_instead_of_encoded_as_silence(): + instance, segments, omitted = recorder(duration=0.01) + instance.record("input", pcm([100] * 1440), sample_rate_hz=48000, channels=1, offset_ns=0) + instance.unavailable(10_000_000, 20_000_000, "playback_timing_unavailable") + instance.advance(30_000_000) + await instance.finish() + assert [segment.meta.index for segment in segments] == [0, 2] + assert len(omitted) == 1 + assert omitted[0][0].placement() == {"start_offset_ms": 10, "duration_ms": 10} + assert omitted[0][1] == "playback_timing_unavailable" + + +def test_missing_extension_warns_once_and_explicit_opus_fails(caplog, monkeypatch): + from braintrust._audio import middleware + + monkeypatch.setattr(middleware, "_warned_missing_extension", False) + assert encoding_stack(None) == () + assert encoding_stack(None) == () + assert len([record for record in caplog.records if "recording raw PCM" in record.message]) == 1 + with pytest.raises(ValueError, match="requires a compatible"): + encoding_stack(AudioEncoding(codec="opus", container="ogg")) + + +def test_tracing_only_session_has_no_media_imports_or_audio_workers(): + subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +from braintrust.logger import NOOP_SPAN +from braintrust._audio.session import AudioSessionRecorder +from braintrust._audio import worker +session = AudioSessionRecorder(NOOP_SPAN) +assert session.recorder is None +assert worker._executor is None +assert not {'av', 'braintrust_audio', 'livekit', 'numpy'} & sys.modules.keys() +""", + ], + check=True, + ) + + +@pytest.mark.asyncio +async def test_terminal_transform_failure_never_reencodes_omitted_tail_as_silence(): + async def fail(data, meta, next): + raise RuntimeError("terminal codec failure") + + instance, segments, omitted = recorder(duration=0.01, middleware=(fail,)) + instance.record("input", pcm([100] * 1440), sample_rate_hz=48000, channels=1, offset_ns=0) + instance.advance(30_000_000) + await instance.finish() + await instance.finish() + assert not segments + assert [meta.placement() for meta, _ in omitted] == [ + {"start_offset_ms": 0, "duration_ms": 10}, + {"start_offset_ms": 10, "duration_ms": 20}, + ] diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py new file mode 100644 index 00000000..2a1f1130 --- /dev/null +++ b/py/src/braintrust/_audio/test_session.py @@ -0,0 +1,183 @@ +import asyncio +import threading + +import pytest +from braintrust import SpanCustomizer, logger, set_span_customizers +from braintrust._audio.config import AudioEncoding, AudioRecordingOptions +from braintrust._audio.session import AudioSessionRecorder +from braintrust.test_helpers import init_test_logger + + +@pytest.fixture +def memory(monkeypatch): + # init_test_logger replaces this global; restore it after each test. + monkeypatch.setattr(logger, "_compute_logger_metadata", logger._compute_logger_metadata) + init_test_logger("audio-tests") + with logger._internal_with_memory_background_logger() as bgl: + monkeypatch.setattr(logger._state, "_global_bg_logger", logger.LazyValue(lambda: bgl, use_mutex=False)) + yield bgl + set_span_customizers(None) + + +def recording(**kwargs): + span = logger.start_span(name="session") + return AudioSessionRecorder( + span, + sources=("session-input",), + options=AudioRecordingOptions( + encoding=AudioEncoding(codec="pcm", container=None), + **kwargs, + ), + ) + + +@pytest.mark.asyncio +async def test_session_source_registry_is_framework_independent(memory): + span = logger.start_span(name="standalone") + session = AudioSessionRecorder(span) + session.source("microphone-A", {"label": "Input"}) + session.source("speaker-B") + assert session.session_metadata()["audio.sources"] == { + "microphone-A": {"label": "Input"}, + "speaker-B": {}, + } + await session.finish() + metadata = memory.pop()[0]["metadata"] + assert set(metadata["audio.sources"]) == {"microphone-A", "speaker-B"} + + +@pytest.mark.asyncio +async def test_ready_only_after_success_and_reference_is_not_reuploaded(memory, monkeypatch): + entered, release = threading.Event(), threading.Event() + uploads = [] + + def upload(attachment): + uploads.append(attachment) + entered.set() + release.wait(5) + return {"upload_status": "done"} + + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + finishing = asyncio.create_task(session.finish()) + assert await asyncio.to_thread(entered.wait, 5) + records = memory.pop() + assert records[0]["metadata"]["audio.segments"]["part-0"]["state"] == "pending" + assert "audio" not in records[0].get("input", {}) + release.set() + await finishing + records = memory.pop() + segment = records[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "ready" + assert segment["channels"] == {"0": {"sources": ["session-input"]}} + assert segment["attributes"] == {"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"} + assert records[0]["input"]["audio"]["part-0"] == uploads[0].reference + assert not memory.upload_attempts + assert len(uploads) == 1 + + +@pytest.mark.asyncio +async def test_customizer_removal_does_not_veto_upload(memory, monkeypatch): + uploads = [] + + class RemoveAudio(SpanCustomizer): + def on_span_export(self, data): + data.pop("input", None) + return data + + def upload(attachment): + uploads.append(attachment) + return {"upload_status": "done"} + + set_span_customizers([RemoveAudio()]) + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + await session.finish() + assert len(uploads) == 1 + assert "input" not in memory.pop()[0] + + +@pytest.mark.asyncio +async def test_privacy_tightening_before_upload_discards_audio(memory, monkeypatch): + def forbidden(_): + pytest.fail("Restricted audio must never upload") + + monkeypatch.setattr(logger.Attachment, "upload", forbidden) + allowed = True + session = recording() + session.allowed = lambda: allowed + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + allowed = False + await session.finish() + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment == {"state": "omitted", "reason": "privacy_restricted", "start_offset_ms": 0, "duration_ms": 1} + + +@pytest.mark.asyncio +async def test_terminal_upload_failure_preserves_known_placement(memory, monkeypatch): + monkeypatch.setattr(logger.Attachment, "upload", lambda _: {"upload_status": "error"}) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + result = await session.finish() + assert result.reason == "encoding_or_upload_failed" + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "omitted" + assert segment["reason"] == "upload_failed" + assert segment["duration_ms"] == 1 + assert "attachment" not in segment + + +@pytest.mark.asyncio +async def test_finish_propagates_cancellation_after_cleanup(memory): + entered, cleaned_up = asyncio.Event(), asyncio.Event() + + async def blocking(data, meta, next): + entered.set() + try: + await asyncio.Event().wait() + finally: + cleaned_up.set() + + session = recording() + session.recorder.middleware = (blocking,) + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + finishing = asyncio.create_task(session.finish()) + try: + await asyncio.wait_for(entered.wait(), 5) + finishing.cancel() + with pytest.raises(asyncio.CancelledError): + await finishing + assert cleaned_up.is_set() + assert finishing.cancelled() + assert session.recorder.buffered_bytes == 0 + assert not session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "omitted" + assert "attachment" not in segment + finally: + finishing.cancel() + await asyncio.gather(finishing, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_shutdown_deadline_fences_late_success(memory, monkeypatch): + entered, release = threading.Event(), threading.Event() + + def upload(_): + entered.set() + release.wait(5) + return {"upload_status": "done"} + + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording(shutdown_timeout_seconds=0.02) + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + result = await session.finish() + assert result.reason == "shutdown_timeout" + assert entered.is_set() + records = memory.pop() + assert records[0]["metadata"]["audio.segments"]["part-0"]["state"] == "omitted" + release.set() + await asyncio.sleep(0.02) + assert not memory.pop() diff --git a/py/src/braintrust/_audio/types.py b/py/src/braintrust/_audio/types.py new file mode 100644 index 00000000..f8dcd8b3 --- /dev/null +++ b/py/src/braintrust/_audio/types.py @@ -0,0 +1,82 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any + + +@dataclass(frozen=True, eq=False) +class _FrozenMapping(Mapping[str, Any]): + """A snapshot whose values have been recursively frozen.""" + + _data: Mapping[str, Any] + + def __post_init__(self): + object.__setattr__(self, "_data", MappingProxyType({key: freeze(value) for key, value in self._data.items()})) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._data) + + def __len__(self) -> int: + return len(self._data) + + +def freeze(value): + if isinstance(value, _FrozenMapping): + return value + if isinstance(value, Mapping): + return _FrozenMapping(value) + if isinstance(value, (list, tuple)): + return tuple(freeze(item) for item in value) + if value is None or isinstance(value, (str, bool, int, float)): + return value + raise TypeError("Audio attributes must contain JSON values") + + +@dataclass(frozen=True) +class AudioChannel: + sources: tuple[str, ...] + + def __post_init__(self): + object.__setattr__(self, "sources", tuple(self.sources)) + + +@dataclass(frozen=True) +class SegmentMeta: + index: int + start_sample: int + sample_count: int + sample_rate_hz: int + channels: tuple[AudioChannel, ...] + mime_type: str = "audio/pcm" + encoding: str = "pcm" + attributes: Mapping[str, Any] = field( + default_factory=lambda: {"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"} + ) + + def __post_init__(self): + object.__setattr__(self, "channels", tuple(self.channels)) + object.__setattr__(self, "attributes", freeze(self.attributes)) + if self.index < 0 or self.start_sample < 0 or self.sample_count <= 0 or self.sample_rate_hz <= 0: + raise ValueError("Invalid audio segment placement") + if not self.channels or not self.mime_type or not self.encoding: + raise ValueError("Invalid audio segment format") + + def placement(self): + start = round(self.start_sample * 1000 / self.sample_rate_hz) + end = round((self.start_sample + self.sample_count) * 1000 / self.sample_rate_hz) + return {"start_offset_ms": start, "duration_ms": end - start} + + +@dataclass(frozen=True) +class AudioResult: + data: bytes + meta: SegmentMeta + + +@dataclass(frozen=True) +class RecordingResult: + segments: int + reason: str | None = None diff --git a/py/src/braintrust/_audio/worker.py b/py/src/braintrust/_audio/worker.py new file mode 100644 index 00000000..df9d7ca5 --- /dev/null +++ b/py/src/braintrust/_audio/worker.py @@ -0,0 +1,82 @@ +"""Bound native submission; cancellation keeps worker-owned buffers alive.""" + +import asyncio +import os +import threading +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from functools import partial + + +_lock = threading.Lock() +_executor: ThreadPoolExecutor | None = None +_pid: int | None = None +_slots: threading.Semaphore | None = None +_waiters: deque[asyncio.Future[None]] | None = None + + +def _admit(request: asyncio.Future[None], slots: threading.Semaphore, waiters: deque[asyncio.Future[None]]) -> None: + # This callback runs on the waiter's event loop, including cancellation checks. + if request.cancelled(): + _release(slots, waiters) + else: + request.set_result(None) + + +def _release(slots: threading.Semaphore, waiters: deque[asyncio.Future[None]]) -> None: + # Transfer the occupied slot directly to the oldest waiter. + while True: + with _lock: + if not waiters: + slots.release() + return + request = waiters.popleft() + try: + request.get_loop().call_soon_threadsafe(_admit, request, slots, waiters) + return + except RuntimeError: + # A closed event loop cannot accept the slot. Try the next waiter. + continue + + +async def run_native(function, *args): + global _executor, _pid, _slots, _waiters + request = asyncio.get_running_loop().create_future() + with _lock: + if _executor is None or _pid != os.getpid(): + _executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="braintrust-audio") + _slots = threading.Semaphore(2) + _waiters = deque() + _pid = os.getpid() + executor, slots, waiters = _executor, _slots, _waiters + assert slots is not None + assert waiters is not None + # Waiting sessions retain their buffers; only admitted jobs reach the executor. + if slots.acquire(blocking=False): + request.set_result(None) + else: + waiters.append(request) + try: + await request + except BaseException: + with _lock: + try: + waiters.remove(request) + except ValueError: + pass + # If admission won the race with cancellation, return its reserved slot. + # Otherwise a pending _admit callback owns the handoff, if any. + if not request.cancelled(): + _release(slots, waiters) + raise + try: + future = executor.submit(partial(function, *args)) + except BaseException: + _release(slots, waiters) + raise + future.add_done_callback(lambda _: _release(slots, waiters)) + wrapped = asyncio.wrap_future(future) + # A cancelled waiter cannot consume a late native failure. Retrieve it + # without logging potentially private exception details. + wrapped.add_done_callback(lambda completed: completed.exception() if not completed.cancelled() else None) + return await asyncio.shield(wrapped) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index f89c4852..a0b9ced8 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -3830,6 +3830,9 @@ def __init__( "key": str(uuid.uuid4()), } self._data_debug_string = data if isinstance(data, str) else "" + # Internal audio uploads opt into finite request timeouts. Ordinary + # attachments retain their existing transport behavior. + self._upload_timeout_seconds: float | None = None self._data = self._init_data(data) self._uploader = self._init_uploader() @@ -3861,6 +3864,9 @@ def debug_info(self) -> Mapping[str, Any]: return {"input_data": self._data_debug_string, "reference": self._reference} def _init_uploader(self) -> LazyValue[AttachmentStatus]: + def request_options(): + return {"timeout": self._upload_timeout_seconds} if self._upload_timeout_seconds is not None else {} + def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: assert self._reference["type"] == "braintrust_attachment" @@ -3872,7 +3878,7 @@ def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: } try: - metadata_response = api_conn.post("/attachment", json=request_params) + metadata_response = api_conn.post("/attachment", json=request_params, **request_options()) metadata_response.raise_for_status() metadata = metadata_response.json() except Exception as e: @@ -3893,7 +3899,7 @@ def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: # TODO multipart upload. try: obj_conn = HTTPConnection(base_url="", adapter=_http_adapter) - obj_response = obj_conn.put(signed_url, headers=headers, data=data) + obj_response = obj_conn.put(signed_url, headers=headers, data=data, **request_options()) obj_response.raise_for_status() except Exception as e: raise RuntimeError(f"Failed to upload attachment to object store: {e}") from e @@ -3925,7 +3931,7 @@ def error_wrapper() -> AttachmentStatus: "status": status, } try: - status_response = api_conn.post("/attachment/status", json=request_params) + status_response = api_conn.post("/attachment/status", json=request_params, **request_options()) status_response.raise_for_status() except Exception as e: raise RuntimeError(f"Couldn't log attachment status: {e}") from e diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 514e4918..0142b977 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2543,6 +2543,34 @@ def broken_masking_function(data): assert "private" not in json.dumps(record) +@pytest.mark.parametrize("timeout", [None, 12.5]) +def test_attachment_internal_timeout_preserves_ordinary_behavior(monkeypatch, with_simulate_login, timeout): + requests_seen = [] + + def post(connection, path, *args, **kwargs): + requests_seen.append(kwargs.get("timeout")) + response = requests.Response() + response.status_code = 200 + response._content = json.dumps( + {"signedUrl": "https://objectstore.example/upload", "headers": {}} if path == "/attachment" else {} + ).encode() + return response + + def put(connection, path, *args, **kwargs): + requests_seen.append(kwargs.get("timeout")) + response = requests.Response() + response.status_code = 200 + return response + + monkeypatch.setattr(logger.HTTPConnection, "post", post) + monkeypatch.setattr(logger.HTTPConnection, "put", put) + attachment = Attachment(data=b"audio", filename="part.pcm", content_type="audio/pcm") + attachment._upload_timeout_seconds = timeout + assert attachment.upload()["upload_status"] == "done" + assert attachment.upload()["upload_status"] == "done" + assert requests_seen == [timeout, timeout, timeout] + + def test_attachment_unreadable_path_logs_warning(caplog): with caplog.at_level(logging.WARNING, logger="braintrust"): Attachment(