From bfb8a970dde79efe1ebe8f5f893e558e349791e0 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 15 Jun 2026 14:32:35 -0700 Subject: [PATCH] feat(attachments): pdf + audio attachment kinds (dormant spine) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Provider-neutral plumbing for PDF and audio attachments, with no user-facing change yet: the upload classifier still rejects them and the capability tables stay unpopulated (both land in the native-translator phase). No migration — workstream_attachments.kind is free-text. - attachments.py: PDF/audio byte caps, allowed-audio MIMEs + format map, magic-byte sniffers (sniff_pdf_mime / sniff_audio_mime), Attachment.is_pdf / is_audio - providers/_protocol.py: supports_pdf / supports_audio_input capability fields (default False; orthogonal to the STT/TTS roles) - storage/_utils.py: attachment_to_content_part emits the internal document(application/pdf, base64) and input_audio shapes - session.py: by-reference placeholder branches for pdf / audio - trajectory.py: AttachmentRef docstring (dict-bridge already kind-agnostic) - tests: test_attachments_pdf_audio.py --- tests/test_attachments_pdf_audio.py | 133 ++++++++++++++++++++++++++ turnstone/core/attachments.py | 74 +++++++++++++- turnstone/core/providers/_protocol.py | 9 ++ turnstone/core/session.py | 4 + turnstone/core/storage/_utils.py | 26 ++++- turnstone/core/trajectory.py | 7 +- 6 files changed, 247 insertions(+), 6 deletions(-) create mode 100644 tests/test_attachments_pdf_audio.py diff --git a/tests/test_attachments_pdf_audio.py b/tests/test_attachments_pdf_audio.py new file mode 100644 index 00000000..638d50c3 --- /dev/null +++ b/tests/test_attachments_pdf_audio.py @@ -0,0 +1,133 @@ +"""Phase 1 (spine) tests for PDF + audio attachment kinds. + +Pure-function coverage for the provider-neutral plumbing: magic-byte sniffers, +``Attachment`` kind predicates, and the internal content-part shapes the wire +builder emits. No DB / provider wiring yet (Phase 2) — these pin the shapes the +later phases translate. +""" + +from __future__ import annotations + +import base64 + +from turnstone.core.attachments import ( + AUDIO_MIME_TO_FORMAT, + Attachment, + sniff_audio_mime, + sniff_pdf_mime, +) +from turnstone.core.storage._utils import attachment_to_content_part + +# --- sample bytes (just enough magic for the sniffers) --------------------- # +PDF = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n1 0 obj\n" +WAV = b"RIFF\x24\x00\x00\x00WAVEfmt " +MP3_ID3 = b"ID3\x04\x00\x00\x00\x00\x00\x00\x00\x00" +MP3_SYNC = b"\xff\xfb\x90\x00" + b"\x00" * 8 +OGG = b"OggS\x00\x02" + b"\x00" * 8 +FLAC = b"fLaC\x00\x00\x00\x22" + b"\x00" * 8 +M4A = b"\x00\x00\x00\x20ftypM4A \x00\x00\x00\x00" +WEBM = b"\x1aE\xdf\xa3" + b"\x00" * 8 +PNG = b"\x89PNG\r\n\x1a\n" + b"\x00" * 8 + + +class TestSniffPdf: + def test_pdf_magic(self) -> None: + assert sniff_pdf_mime(PDF) == "application/pdf" + + def test_rejects_non_pdf(self) -> None: + assert sniff_pdf_mime(PNG) is None + assert sniff_pdf_mime(b"not a pdf at all") is None + + def test_too_short(self) -> None: + assert sniff_pdf_mime(b"%PD") is None + assert sniff_pdf_mime(b"") is None + + +class TestSniffAudio: + def test_each_format(self) -> None: + assert sniff_audio_mime(WAV) == "audio/wav" + assert sniff_audio_mime(MP3_ID3) == "audio/mpeg" + assert sniff_audio_mime(MP3_SYNC) == "audio/mpeg" + assert sniff_audio_mime(OGG) == "audio/ogg" + assert sniff_audio_mime(FLAC) == "audio/flac" + assert sniff_audio_mime(M4A) == "audio/mp4" + assert sniff_audio_mime(WEBM) == "audio/webm" + + def test_rejects_non_audio(self) -> None: + assert sniff_audio_mime(PNG) is None + assert sniff_audio_mime(PDF) is None + + def test_too_short(self) -> None: + assert sniff_audio_mime(b"RIFF") is None + assert sniff_audio_mime(b"") is None + + +class TestAttachmentKindPredicates: + def _att(self, kind: str) -> Attachment: + return Attachment( + attachment_id="a", + filename="f", + mime_type="application/octet-stream", + kind=kind, + content=b"x", + ) + + def test_pdf(self) -> None: + a = self._att("pdf") + assert a.is_pdf and not (a.is_image or a.is_text or a.is_audio) + + def test_audio(self) -> None: + a = self._att("audio") + assert a.is_audio and not (a.is_image or a.is_text or a.is_pdf) + + def test_existing_kinds_unaffected(self) -> None: + assert self._att("image").is_image + assert self._att("text").is_text + + +class TestContentPartBuilder: + def test_pdf_part_is_base64_document(self) -> None: + raw = PDF + part = attachment_to_content_part( + {"kind": "pdf", "content": raw, "mime_type": "application/pdf", "filename": "doc.pdf"} + ) + assert part is not None + assert part["type"] == "document" + doc = part["document"] + assert doc["name"] == "doc.pdf" + assert doc["media_type"] == "application/pdf" + # base64 (not utf-8 text) — round-trips to the original bytes. + assert base64.b64decode(doc["data"]) == raw + + def test_audio_part_is_input_audio(self) -> None: + raw = WAV + part = attachment_to_content_part( + {"kind": "audio", "content": raw, "mime_type": "audio/wav", "filename": "a.wav"} + ) + assert part is not None + assert part["type"] == "input_audio" + ia = part["input_audio"] + assert ia["format"] == "wav" + assert base64.b64decode(ia["data"]) == raw + + def test_audio_format_falls_back_to_codec_token(self) -> None: + part = attachment_to_content_part( + { + "kind": "audio", + "content": b"\x00" * 16, + "mime_type": "audio/x-exotic", + "filename": "x", + } + ) + assert part is not None + assert part["input_audio"]["format"] == "x-exotic" + + def test_unknown_kind_returns_none(self) -> None: + assert attachment_to_content_part({"kind": "weird", "content": b"x"}) is None + + +class TestAudioFormatMap: + def test_known_mimes_map_to_codec_tokens(self) -> None: + assert AUDIO_MIME_TO_FORMAT["audio/mpeg"] == "mp3" + assert AUDIO_MIME_TO_FORMAT["audio/wav"] == "wav" + assert AUDIO_MIME_TO_FORMAT["audio/mp4"] == "m4a" diff --git a/turnstone/core/attachments.py b/turnstone/core/attachments.py index 9672b227..f79f9b74 100644 --- a/turnstone/core/attachments.py +++ b/turnstone/core/attachments.py @@ -29,19 +29,50 @@ if TYPE_CHECKING: # constants live here so the session / tests share the same definitions. IMAGE_SIZE_CAP: int = 4 * 1024 * 1024 TEXT_DOC_SIZE_CAP: int = 512 * 1024 +PDF_SIZE_CAP: int = 32 * 1024 * 1024 +AUDIO_SIZE_CAP: int = 25 * 1024 * 1024 ALLOWED_IMAGE_MIMES: frozenset[str] = frozenset( {"image/png", "image/jpeg", "image/gif", "image/webp"} ) +# Audio MIMEs accepted as chat attachments (sniffed by magic bytes; the +# client-claimed Content-Type is never trusted alone). ``AUDIO_MIME_TO_FORMAT`` +# maps each to the OpenAI ``input_audio.format`` token the wire builder emits. +ALLOWED_AUDIO_MIMES: frozenset[str] = frozenset( + { + "audio/wav", + "audio/x-wav", + "audio/mpeg", + "audio/mp3", + "audio/ogg", + "audio/flac", + "audio/mp4", + "audio/aac", + "audio/webm", + } +) + +AUDIO_MIME_TO_FORMAT: dict[str, str] = { + "audio/wav": "wav", + "audio/x-wav": "wav", + "audio/mpeg": "mp3", + "audio/mp3": "mp3", + "audio/ogg": "ogg", + "audio/flac": "flac", + "audio/mp4": "m4a", + "audio/aac": "aac", + "audio/webm": "webm", +} + @dataclass(frozen=True) class Attachment: """An attachment resolved from storage, ready for injection into a turn. - ``kind`` is ``"image"`` or ``"text"``. ``content`` is raw bytes — for - text attachments, UTF-8 decoded at the point of content-part - construction. + ``kind`` is ``"image"``, ``"text"``, ``"pdf"``, or ``"audio"``. ``content`` + is raw bytes — text attachments are UTF-8 decoded at content-part + construction; image/pdf/audio are base64-encoded at the wire boundary. """ attachment_id: str @@ -58,6 +89,14 @@ class Attachment: def is_text(self) -> bool: return self.kind == "text" + @property + def is_pdf(self) -> bool: + return self.kind == "pdf" + + @property + def is_audio(self) -> bool: + return self.kind == "audio" + # --------------------------------------------------------------------------- # Upload classification @@ -112,6 +151,35 @@ def sniff_image_mime(data: bytes) -> str | None: return None +def sniff_pdf_mime(data: bytes) -> str | None: + """Return ``"application/pdf"`` if ``data`` starts with the PDF magic, else None.""" + return "application/pdf" if data[:5] == b"%PDF-" else None + + +def sniff_audio_mime(data: bytes) -> str | None: + """Return a canonical audio MIME type by inspecting magic bytes. + + Covers WAV, MP3 (ID3 tag or MPEG frame sync), OGG, FLAC, ISO-BMFF + (m4a / mp4 audio), and WebM/Matroska. Returns ``None`` on no match — the + client-provided ``Content-Type`` is never trusted alone. + """ + if len(data) < 12: + return None + if data[:4] == b"RIFF" and data[8:12] == b"WAVE": + return "audio/wav" + if data[:3] == b"ID3" or data[:2] in (b"\xff\xfb", b"\xff\xf3", b"\xff\xf2"): + return "audio/mpeg" + if data[:4] == b"OggS": + return "audio/ogg" + if data[:4] == b"fLaC": + return "audio/flac" + if data[4:8] == b"ftyp": + return "audio/mp4" + if data[:4] == b"\x1aE\xdf\xa3": + return "audio/webm" + return None + + def classify_text_attachment( filename: str, claimed_mime: str, data: bytes ) -> tuple[str | None, str | None]: diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index f13f02ed..ef1a5612 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -86,6 +86,15 @@ class ModelCapabilities: supports_web_search: bool = False supports_tool_search: bool = False supports_vision: bool = False + # Chat-input modalities carried as user-turn attachments — distinct from the + # STT/TTS *roles* below. ``supports_pdf``: native PDF document ingest; + # ``supports_audio_input``: native audio ingest (OpenAI ``input_audio`` / + # vLLM omni). When False the wire-build path falls back client-side (PDF → + # rasterize/extract, audio → STT transcription) — see core/attachments.py. + # ``supports_audio_input`` is orthogonal to ``supports_transcription``: an + # omni model has the former and lacks the latter (it has no /audio endpoint). + supports_pdf: bool = False + supports_audio_input: bool = False # Audio I/O roles (STT / TTS) — not chat behavior; consumed by the audio # endpoints and the Models -> Roles capability gate (turnstone/core/audio.py). supports_transcription: bool = False diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 81b69c79..f7000b1a 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -3583,6 +3583,10 @@ class ChatSession: parts.append({"type": "image", "attachment_id": att.attachment_id}) elif att.is_text: parts.append({"type": "document", "attachment_id": att.attachment_id}) + elif att.is_pdf: + parts.append({"type": "pdf", "attachment_id": att.attachment_id}) + elif att.is_audio: + parts.append({"type": "audio", "attachment_id": att.attachment_id}) else: log.warning( "attachment id=%s has unknown kind=%r; injecting placeholder", diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index b882f713..d9aeafb3 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -10,7 +10,7 @@ from typing import Any import sqlalchemy as sa -from turnstone.core.attachments import unreadable_placeholder +from turnstone.core.attachments import AUDIO_MIME_TO_FORMAT, unreadable_placeholder from turnstone.core.log import get_logger from turnstone.core.storage._schema import ( conversations, @@ -324,6 +324,30 @@ def attachment_to_content_part(att: dict[str, Any]) -> dict[str, Any] | None: "data": text, }, } + if kind == "pdf" and isinstance(raw, bytes): + # PDF rides as a ``document`` part discriminated by media_type: + # base64 bytes (vs. a text doc's utf-8 ``data``). Per-provider + # translators branch on ``application/pdf`` (Phase 2); the client-side + # fallback for non-PDF models lands in Phase 3. + b64 = base64.b64encode(raw).decode("ascii") + return { + "type": "document", + "document": { + "name": att.get("filename") or "", + "media_type": "application/pdf", + "data": b64, + }, + } + if kind == "audio" and isinstance(raw, bytes): + # OpenAI-style ``input_audio`` part — passes through the openai-compat + # lane untouched (omni models); other lanes translate / fall back in + # Phase 2/3. ``format`` is the bare codec token derived from the MIME. + b64 = base64.b64encode(raw).decode("ascii") + fmt = AUDIO_MIME_TO_FORMAT.get(mime) or (mime.split("/", 1)[-1] if "/" in mime else "wav") + return { + "type": "input_audio", + "input_audio": {"data": b64, "format": fmt}, + } return None diff --git a/turnstone/core/trajectory.py b/turnstone/core/trajectory.py index c501a6c4..ab3b0f73 100644 --- a/turnstone/core/trajectory.py +++ b/turnstone/core/trajectory.py @@ -49,8 +49,11 @@ class AttachmentRef: """A reference to attachment bytes held in the content-addressed blob store. Non-text content is carried *by reference* (never inline bytes): the translator - resolves ``attachment_id`` to bytes and expands it to the provider's image / - document format at wire time. ``kind`` is ``"image"`` or ``"document"``. + resolves ``attachment_id`` to bytes and expands it to the provider's native + format at wire time. ``kind`` is the by-reference placeholder type — + ``"image"``, ``"document"`` (text docs), ``"pdf"``, or ``"audio"``. The + dict-bridge keys off ``attachment_id`` and is kind-agnostic, so new kinds + need no change here. """ attachment_id: str