mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-29 18:46:59 +00:00
refactor: extract shared audio container sniffer to tools/audio_container.py
One sniffer owns magic-byte container detection (Teknium's one-concept- one-owner rule): the new tools/audio_container.py is used by - gateway/platforms/base.py _sniff_audio_ext (inbound cache — PR #36166's central sniffer, now covering AAC/ADTS, MP4-brand disambiguation, webm) - gateway/platforms/signal.py _guess_extension (audio/AV branches delegated; RIFF/WAVE fix from PR #50690 and M4A-brand fix from PR #72490 now live centrally) - tools/tts_tool.py _sniff_audio_container (outbound repair, PR #73072) cache_audio_from_url inherits the sniff via cache_audio_from_bytes. Adds tests/tools/test_audio_container.py covering every magic-byte type, wrong-extension repair on the inbound cache, unknown passthrough, the URL path, and Signal's delegation.
This commit is contained in:
Teknium
parent
26f50e2f1a
commit
156edc5ded
5 changed files with 320 additions and 51 deletions
|
|
@ -854,21 +854,15 @@ def get_audio_cache_dir() -> Path:
|
|||
|
||||
|
||||
def _sniff_audio_ext(data: bytes, fallback_ext: str) -> str:
|
||||
"""Prefer a container-matching extension when audio magic bytes are obvious."""
|
||||
fallback = fallback_ext if fallback_ext.startswith(".") else f".{fallback_ext}"
|
||||
if len(data) >= 8 and data[4:8] == b"ftyp":
|
||||
return ".m4a"
|
||||
if data.startswith(b"OggS"):
|
||||
return ".ogg"
|
||||
if data.startswith(b"fLaC"):
|
||||
return ".flac"
|
||||
if len(data) >= 12 and data.startswith(b"RIFF") and data[8:12] == b"WAVE":
|
||||
return ".wav"
|
||||
if data.startswith(b"ID3") or data[:2] in {b"\xff\xfb", b"\xff\xf3", b"\xff\xf2"}:
|
||||
return ".mp3"
|
||||
if data.startswith(b"\x1a\x45\xdf\xa3"):
|
||||
return ".webm"
|
||||
return fallback
|
||||
"""Prefer a container-matching extension when audio magic bytes are obvious.
|
||||
|
||||
Thin wrapper around the shared sniffer in ``tools.audio_container`` —
|
||||
ONE module owns container detection for both the outbound TTS repair
|
||||
(``tools/tts_tool.py``) and this inbound cache path.
|
||||
"""
|
||||
from tools.audio_container import sniff_audio_ext
|
||||
|
||||
return sniff_audio_ext(data, fallback_ext)
|
||||
|
||||
|
||||
def cache_audio_from_bytes(data: bytes, ext: str = ".ogg") -> str:
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from gateway.platforms.base import (
|
|||
cache_image_from_url,
|
||||
)
|
||||
from gateway.platforms.helpers import redact_phone
|
||||
from tools.audio_container import CONTAINER_TO_EXT, sniff_container
|
||||
from gateway.platforms.signal_format import markdown_to_signal
|
||||
from gateway.platforms.signal_rate_limit import (
|
||||
SIGNAL_BATCH_PACING_NOTICE_THRESHOLD,
|
||||
|
|
@ -97,34 +98,17 @@ def _guess_extension(data: bytes) -> str:
|
|||
return ".gif"
|
||||
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||||
return ".webp"
|
||||
# RIFF/WAVE shares the ``RIFF`` chunk header with WebP; the form-type at
|
||||
# bytes 8-11 disambiguates. Without this, an inbound WAV voice note fell
|
||||
# through to ``.bin`` (octet-stream) and was cached as a document, so STT
|
||||
# never saw it — even though ``.wav`` is in the audio allowlist + MIME map.
|
||||
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WAVE":
|
||||
return ".wav"
|
||||
if data[:4] == b"%PDF":
|
||||
return ".pdf"
|
||||
if len(data) >= 8 and data[4:8] == b"ftyp":
|
||||
# iOS Signal delivers voice notes as MP4-container AAC carrying an
|
||||
# audio ftyp brand ("M4A ", "M4B "). Returning ".mp4" for those made
|
||||
# them cache as documents and STT reject them ("Invalid file
|
||||
# format") even though the bytes are valid audio. Read the brand so
|
||||
# audio-branded files land as ".m4a" and route to the audio cache.
|
||||
brand = data[8:12].lower() if len(data) >= 12 else b""
|
||||
if brand in (b"m4a ", b"m4b "):
|
||||
return ".m4a"
|
||||
return ".mp4"
|
||||
if data[:4] == b"OggS":
|
||||
return ".ogg"
|
||||
if len(data) >= 2 and data[0] == 0xFF and (data[1] & 0xE0) == 0xE0:
|
||||
# ``0xFF 0xFx`` is shared by MP3 and ADTS AAC. The discriminator
|
||||
# is bits 3-1 of byte 1: ADTS has ``ID=0`` and ``layer=00`` (mask
|
||||
# 0xF6, target 0xF0); MP3 has ``ID=1`` and ``layer`` in {01,10,11}
|
||||
# (mask 0xF6, target in {0xF2, 0xF4, 0xF6}).
|
||||
if (data[1] & 0xF6) == 0xF0:
|
||||
return ".aac"
|
||||
return ".mp3"
|
||||
# Audio/AV containers: delegate to the shared central sniffer
|
||||
# (tools/audio_container.py) — ONE module owns magic-byte container
|
||||
# detection. It handles the brand/form-type disambiguations this
|
||||
# function used to carry locally: RIFF/WAVE vs WEBP (WEBP is claimed
|
||||
# above, before delegation), ftyp audio brands ("M4A ", "M4B ") vs
|
||||
# video brands (isom/mp42/avc1/qt), and MP3 vs ADTS AAC sync words.
|
||||
container = sniff_container(data)
|
||||
if container is not None:
|
||||
return CONTAINER_TO_EXT[container]
|
||||
if data[:2] == b"PK":
|
||||
return ".zip"
|
||||
return ".bin"
|
||||
|
|
|
|||
195
tests/tools/test_audio_container.py
Normal file
195
tests/tools/test_audio_container.py
Normal file
|
|
@ -0,0 +1,195 @@
|
|||
"""Tests for the shared magic-byte audio container sniffer.
|
||||
|
||||
``tools/audio_container.py`` is the single owner of container detection:
|
||||
the outbound TTS repair (``tools/tts_tool.py``), the inbound gateway audio
|
||||
cache (``gateway/platforms/base.py``), and Signal's ``_guess_extension``
|
||||
all delegate to it. These tests cover every magic-byte branch, the
|
||||
wrong-extension repair behaviour on the inbound cache path, and unknown
|
||||
passthrough.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tools.audio_container import CONTAINER_TO_EXT, sniff_audio_ext, sniff_container
|
||||
|
||||
# --- canonical headers ------------------------------------------------------
|
||||
OGG = b"OggS\x00\x02" + b"\x00" * 64
|
||||
FLAC = b"fLaC" + b"\x00" * 64
|
||||
WAV = b"RIFF\x24\x08\x00\x00WAVEfmt " + b"\x00" * 64
|
||||
WEBP = b"RIFF\x24\x08\x00\x00WEBPVP8 " + b"\x00" * 64
|
||||
MP3_ID3 = b"ID3\x04\x00\x00\x00\x00\x00\x00" + b"\x00" * 64
|
||||
MP3_FRAME = b"\xff\xfb\x90\x00" + b"\x00" * 64
|
||||
AAC_ADTS = b"\xff\xf1\x50\x80" + b"\x00" * 64
|
||||
M4A = b"\x00\x00\x00\x1cftypM4A " + b"\x00" * 64
|
||||
M4B = b"\x00\x00\x00\x1cftypM4B " + b"\x00" * 64
|
||||
MP4_ISOM = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 64
|
||||
WEBM = b"\x1a\x45\xdf\xa3" + b"\x00" * 64
|
||||
UNKNOWN = b"not-audio-at-all" + b"\x00" * 64
|
||||
|
||||
|
||||
class TestSniffContainer:
|
||||
@pytest.mark.parametrize(
|
||||
"data,expected",
|
||||
[
|
||||
(OGG, "ogg"),
|
||||
(FLAC, "flac"),
|
||||
(WAV, "wav"),
|
||||
(MP3_ID3, "mp3"),
|
||||
(MP3_FRAME, "mp3"),
|
||||
(AAC_ADTS, "aac"),
|
||||
(M4A, "m4a"),
|
||||
(M4B, "m4a"),
|
||||
(MP4_ISOM, "mp4"),
|
||||
(WEBM, "webm"),
|
||||
],
|
||||
)
|
||||
def test_magic_bytes(self, data, expected):
|
||||
assert sniff_container(data) == expected
|
||||
|
||||
def test_unknown_returns_none(self):
|
||||
assert sniff_container(UNKNOWN) is None
|
||||
assert sniff_container(b"") is None
|
||||
assert sniff_container(b"\x00") is None
|
||||
|
||||
def test_webp_is_not_claimed(self):
|
||||
# Images are the caller's business — the sniffer must not claim
|
||||
# RIFF/WEBP just because it shares the RIFF header with WAV.
|
||||
assert sniff_container(WEBP) is None
|
||||
|
||||
def test_video_ftyp_brands_stay_mp4(self):
|
||||
for brand in (b"isom", b"mp42", b"avc1", b"qt "):
|
||||
data = b"\x00\x00\x00\x1cftyp" + brand + b"\x00" * 64
|
||||
assert sniff_container(data) == "mp4", brand
|
||||
|
||||
def test_every_container_has_an_extension(self):
|
||||
for data in (OGG, FLAC, WAV, MP3_ID3, AAC_ADTS, M4A, MP4_ISOM, WEBM):
|
||||
container = sniff_container(data)
|
||||
assert container in CONTAINER_TO_EXT
|
||||
|
||||
|
||||
class TestSniffAudioExt:
|
||||
@pytest.mark.parametrize(
|
||||
"data,expected",
|
||||
[
|
||||
(OGG, ".ogg"),
|
||||
(FLAC, ".flac"),
|
||||
(WAV, ".wav"),
|
||||
(MP3_ID3, ".mp3"),
|
||||
(MP3_FRAME, ".mp3"),
|
||||
(AAC_ADTS, ".aac"),
|
||||
(M4A, ".m4a"),
|
||||
(WEBM, ".webm"),
|
||||
],
|
||||
)
|
||||
def test_container_wins_over_claimed_ext(self, data, expected):
|
||||
assert sniff_audio_ext(data, ".ogg" if expected != ".ogg" else ".mp3") == expected
|
||||
|
||||
def test_generic_mp4_maps_to_m4a_in_audio_context(self):
|
||||
# In an audio context an MP4 container is AAC audio regardless of
|
||||
# brand — .m4a keeps voice-bubble/audio routing working.
|
||||
assert sniff_audio_ext(MP4_ISOM, ".ogg") == ".m4a"
|
||||
|
||||
def test_unknown_passthrough_keeps_fallback(self):
|
||||
assert sniff_audio_ext(UNKNOWN, ".aac") == ".aac"
|
||||
assert sniff_audio_ext(b"", ".ogg") == ".ogg"
|
||||
|
||||
def test_fallback_without_dot_is_normalized(self):
|
||||
assert sniff_audio_ext(UNKNOWN, "mp3") == ".mp3"
|
||||
|
||||
|
||||
class TestInboundCacheUsesSniffer:
|
||||
"""cache_audio_from_bytes must repair wrong caller-supplied extensions."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data,claimed,expected_suffix",
|
||||
[
|
||||
(MP3_ID3, ".ogg", ".mp3"), # MP3 bytes delivered as "voice.ogg"
|
||||
(WAV, ".ogg", ".wav"), # RIFF/WAVE in an .ogg wrapper claim
|
||||
(M4A, ".ogg", ".m4a"), # iOS voice note claimed as ogg
|
||||
(OGG, ".mp3", ".ogg"), # real opus claimed as mp3
|
||||
(AAC_ADTS, ".ogg", ".aac"), # Android ADTS AAC voice note
|
||||
(WEBM, ".ogg", ".webm"),
|
||||
(FLAC, ".mp3", ".flac"),
|
||||
],
|
||||
)
|
||||
def test_wrong_ext_repaired(self, tmp_path, data, claimed, expected_suffix):
|
||||
from gateway.platforms.base import cache_audio_from_bytes
|
||||
|
||||
with patch("gateway.platforms.base.AUDIO_CACHE_DIR", tmp_path):
|
||||
result = cache_audio_from_bytes(data, ext=claimed)
|
||||
|
||||
saved = tmp_path / os.path.basename(result)
|
||||
assert saved.suffix == expected_suffix
|
||||
assert saved.read_bytes() == data
|
||||
|
||||
def test_unknown_bytes_keep_claimed_ext(self, tmp_path):
|
||||
from gateway.platforms.base import cache_audio_from_bytes
|
||||
|
||||
with patch("gateway.platforms.base.AUDIO_CACHE_DIR", tmp_path):
|
||||
result = cache_audio_from_bytes(UNKNOWN, ext=".amr")
|
||||
|
||||
assert result.endswith(".amr")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_audio_from_url_sniffs_too(self, tmp_path, monkeypatch):
|
||||
"""The URL download path routes through the same sniffer."""
|
||||
import gateway.platforms.base as base
|
||||
|
||||
monkeypatch.setattr(base, "AUDIO_CACHE_DIR", tmp_path)
|
||||
|
||||
class _FakeResponse:
|
||||
status_code = 200
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeStreamCM:
|
||||
async def __aenter__(self):
|
||||
return _FakeResponse()
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
class _FakeClient:
|
||||
def stream(self, *a, **k):
|
||||
return _FakeStreamCM()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
async def _fake_read(response, media_type):
|
||||
return MP3_ID3 # server sent MP3 bytes despite the .ogg claim
|
||||
|
||||
monkeypatch.setattr(base, "_read_httpx_body_with_limit", _fake_read)
|
||||
monkeypatch.setattr(
|
||||
"tools.url_safety.create_ssrf_safe_async_client",
|
||||
lambda **k: _FakeClient(),
|
||||
)
|
||||
monkeypatch.setattr("tools.url_safety.is_safe_url", lambda u: True)
|
||||
|
||||
result = await base.cache_audio_from_url("https://example.com/voice.ogg", ext=".ogg")
|
||||
assert result.endswith(".mp3")
|
||||
|
||||
|
||||
class TestSignalDelegatesToCentralSniffer:
|
||||
"""signal._guess_extension audio branches delegate to the shared module."""
|
||||
|
||||
def test_signal_uses_shared_sniffer(self, monkeypatch):
|
||||
from gateway.platforms import signal as signal_mod
|
||||
|
||||
calls = []
|
||||
real = signal_mod.sniff_container
|
||||
|
||||
def _spy(data):
|
||||
calls.append(data[:4])
|
||||
return real(data)
|
||||
|
||||
monkeypatch.setattr(signal_mod, "sniff_container", _spy)
|
||||
assert signal_mod._guess_extension(M4A) == ".m4a"
|
||||
assert calls, "signal._guess_extension did not delegate to the central sniffer"
|
||||
97
tools/audio_container.py
Normal file
97
tools/audio_container.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
"""Shared magic-byte audio/AV container detection.
|
||||
|
||||
ONE sniffer owns container detection for the whole codebase:
|
||||
|
||||
- **Outbound** (``tools/tts_tool.py``): TTS backends silently ignore the
|
||||
requested opus format (Edge emits MP3, Piper writes WAV, ...), so the
|
||||
synthesized file is sniffed and repaired when the bytes don't match the
|
||||
``.ogg`` extension (PR #73072).
|
||||
- **Inbound** (``gateway/platforms/base.py`` ``cache_audio_from_bytes`` /
|
||||
``cache_audio_from_url``): platform adapters frequently pass a wrong or
|
||||
guessed extension for voice notes (Telegram ``.oga``, iOS Signal M4A-branded
|
||||
MP4, RIFF/WAVE attachments). The cache sniffs the real container so STT and
|
||||
downstream players get an honest extension — the inbound mirror of the
|
||||
outbound repair.
|
||||
- ``gateway/platforms/signal.py`` ``_guess_extension`` delegates its audio/AV
|
||||
branches here instead of duplicating the byte patterns.
|
||||
|
||||
Detection notes:
|
||||
|
||||
- RIFF needs the form-type at bytes 8-11 to split ``WAVE`` (wav) from ``WEBP``
|
||||
(image — deliberately NOT handled here; this module only claims audio/AV
|
||||
containers, callers check images first).
|
||||
- ``ftyp`` needs the brand at bytes 8-11 to split audio brands (``M4A ``,
|
||||
``M4B ``) from video brands (isom/mp42/avc1/qt).
|
||||
- The ``0xFF 0xFx`` sync word is shared by MP3 and ADTS AAC; bits 3-1 of
|
||||
byte 1 disambiguate (ADTS: ``ID=0``, ``layer=00``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
# Container id -> canonical file extension.
|
||||
CONTAINER_TO_EXT = {
|
||||
"m4a": ".m4a",
|
||||
"mp4": ".mp4",
|
||||
"ogg": ".ogg",
|
||||
"flac": ".flac",
|
||||
"wav": ".wav",
|
||||
"mp3": ".mp3",
|
||||
"aac": ".aac",
|
||||
"webm": ".webm",
|
||||
}
|
||||
|
||||
# MP4 ftyp brands that mean "this is audio" (iOS voice notes use M4A ).
|
||||
_MP4_AUDIO_BRANDS = (b"m4a ", b"m4b ")
|
||||
|
||||
|
||||
def sniff_container(data: bytes) -> Optional[str]:
|
||||
"""Return a container id from magic bytes, or ``None`` when unknown.
|
||||
|
||||
Possible ids: ``m4a``, ``mp4``, ``ogg``, ``flac``, ``wav``, ``mp3``,
|
||||
``aac``, ``webm``. Only audio/AV containers are claimed — images
|
||||
(including RIFF/WEBP) return ``None`` so callers can layer their own
|
||||
image detection first.
|
||||
"""
|
||||
if len(data) >= 8 and data[4:8] == b"ftyp":
|
||||
# Brand at bytes 8-11: audio brands ("M4A ", "M4B ") are voice
|
||||
# notes / audiobooks; everything else (isom/mp42/avc1/qt) is video.
|
||||
if len(data) >= 12 and data[8:12].lower() in _MP4_AUDIO_BRANDS:
|
||||
return "m4a"
|
||||
return "mp4"
|
||||
if data.startswith(b"OggS"):
|
||||
return "ogg"
|
||||
if data.startswith(b"fLaC"):
|
||||
return "flac"
|
||||
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WAVE":
|
||||
return "wav"
|
||||
if data.startswith(b"ID3"):
|
||||
return "mp3"
|
||||
if len(data) >= 2 and data[0] == 0xFF and (data[1] & 0xE0) == 0xE0:
|
||||
# ``0xFF 0xFx`` is shared by MP3 and ADTS AAC. Bits 3-1 of byte 1
|
||||
# disambiguate: ADTS has ``ID=0`` and ``layer=00`` (mask 0xF6,
|
||||
# target 0xF0); MP3 has ``ID=1`` and ``layer`` in {01,10,11}.
|
||||
if (data[1] & 0xF6) == 0xF0:
|
||||
return "aac"
|
||||
return "mp3"
|
||||
if data.startswith(b"\x1a\x45\xdf\xa3"):
|
||||
return "webm"
|
||||
return None
|
||||
|
||||
|
||||
def sniff_audio_ext(data: bytes, fallback_ext: str = ".ogg") -> str:
|
||||
"""Return a container-matching extension, or ``fallback_ext`` when unknown.
|
||||
|
||||
Used on inbound audio paths where the caller *claims* the bytes are audio:
|
||||
generic MP4 containers are mapped to ``.m4a`` (audio-in-MP4) because in an
|
||||
audio context the payload is AAC audio regardless of brand — STT accepts
|
||||
``.m4a``/``.mp4`` but voice-bubble routing keys off audio extensions.
|
||||
"""
|
||||
fallback = fallback_ext if fallback_ext.startswith(".") else f".{fallback_ext}"
|
||||
container = sniff_container(data)
|
||||
if container is None:
|
||||
return fallback
|
||||
if container == "mp4":
|
||||
return ".m4a"
|
||||
return CONTAINER_TO_EXT[container]
|
||||
|
|
@ -995,21 +995,20 @@ def _ffmpeg_transcode_to_opus(input_path: str, ogg_path: str) -> Optional[str]:
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _sniff_audio_container(path: str) -> str:
|
||||
"""Return 'ogg' | 'wav' | 'mp3' | 'flac' | 'unknown' from magic bytes."""
|
||||
"""Return a container id ('ogg', 'wav', 'mp3', 'flac', ...) or 'unknown'.
|
||||
|
||||
Delegates to the shared magic-byte sniffer in ``tools.audio_container``
|
||||
(one module owns container detection for both this outbound repair and
|
||||
the inbound gateway audio cache).
|
||||
"""
|
||||
from tools.audio_container import sniff_container
|
||||
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
head = fh.read(12)
|
||||
except OSError:
|
||||
return "unknown"
|
||||
if head[:4] == b"OggS":
|
||||
return "ogg"
|
||||
if head[:4] == b"RIFF" and head[8:12] == b"WAVE":
|
||||
return "wav"
|
||||
if head[:4] == b"fLaC":
|
||||
return "flac"
|
||||
if head[:3] == b"ID3" or (len(head) >= 2 and head[0] == 0xFF and (head[1] & 0xE0) == 0xE0):
|
||||
return "mp3"
|
||||
return "unknown"
|
||||
return sniff_container(head) or "unknown"
|
||||
|
||||
|
||||
def _repair_ogg_container(file_str: str) -> str:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue