diff options
| -rw-r--r-- | llamachat/backend.py | 58 | ||||
| -rwxr-xr-x | test_llamachat.py | 48 |
2 files changed, 98 insertions, 8 deletions
diff --git a/llamachat/backend.py b/llamachat/backend.py index 70a98b1..b79040f 100644 --- a/llamachat/backend.py +++ b/llamachat/backend.py @@ -44,6 +44,14 @@ CODE_SUFFIXES = { IMAGE_MIMES = {"image/png", "image/jpeg", "image/gif", "image/webp", "image/bmp"} +AUDIO_MIMES = {"audio/wav", "audio/x-wav"} +AUDIO_SUFFIXES = {".wav"} + +# Used when a recording is sent with no typed text. llama.cpp pairs audio with +# a text part, and this is the least presumptuous instruction that still tells +# the model what the clip is. +AUDIO_PROMPT = "Listen to this audio and respond." + # Sentinel yielded by the SSE parser at the end of a stream. A distinct # object rather than '' so an empty delta cannot be mistaken for the end. DONE = ("done", "") @@ -102,15 +110,18 @@ class Attachment: sha256: str text: str = "" # kind == 'text': file content, possibly truncated data_url: str = "" # kind == 'image': base64 data: URL for the API + b64: str = "" # kind == 'audio': raw base64 WAV (no data: prefix) thumb: bytes | None = None truncated: bool = False def classify(path: Path) -> str: - """Decide whether a path is a text or image attachment.""" + """Decide whether a path is a text, image or audio attachment.""" mime, _ = mimetypes.guess_type(path.name) if mime in IMAGE_MIMES: return "image" + if mime in AUDIO_MIMES or path.suffix.lower() in AUDIO_SUFFIXES: + return "audio" if path.suffix in CODE_SUFFIXES: return "text" if mime and mime.startswith("text/"): @@ -148,6 +159,11 @@ def load_attachment(path: Path, char_budget: int) -> Attachment: att.text = raw return att + if kind == "audio": + raw = path.read_bytes() + att.b64 = base64.b64encode(raw).decode("ascii") + return att + raw = path.read_bytes() att.data_url = f"data:{mime};base64,{base64.b64encode(raw).decode('ascii')}" att.thumb = _thumbnail(path) @@ -187,7 +203,8 @@ def build_user_content(text: str, attachments: list[Attachment]): joined = "\n\n".join(parts) images = [a for a in attachments if a.kind == "image"] - if not images: + audios = [a for a in attachments if a.kind == "audio"] + if not images and not audios: return joined content: list[dict] = [{"type": "text", "text": joined}] @@ -195,9 +212,33 @@ def build_user_content(text: str, attachments: list[Attachment]): content.append( {"type": "image_url", "image_url": {"url": att.data_url}} ) + for att in audios: + content.append( + { + "type": "input_audio", + "input_audio": {"data": att.b64, "format": "wav"}, + } + ) return content +def audio_attachment(wav: bytes, seconds: int) -> Attachment: + """An Attachment for a recorded clip. + + The bytes are never written to disk: the model gets them in the request, + and the history row records only their size and hash. The name is + synthetic and carries the duration for display. + """ + return Attachment( + path=Path(f"voice-note-{seconds}s.wav"), + kind="audio", + mime="audio/wav", + size=len(wav), + sha256=hashlib.sha256(wav).hexdigest(), + b64=base64.b64encode(wav).decode("ascii"), + ) + + class Client: """Minimal OpenAI-compatible client for the local router.""" @@ -689,12 +730,13 @@ def estimate_tokens(messages: list[dict], chars_per_token: float) -> int: """Rough token count for a request that has not been sent yet. Used for the live meter while typing. Once a reply completes, the exact - count from the usage block replaces this. Images are counted as a flat - allowance rather than by dimension, which is enough to show that a - picture is expensive without pretending to know the tiling. + count from the usage block replaces this. A non-text part (image or + audio) is counted as a flat allowance rather than by dimension, which is + enough to show that a picture is expensive without pretending to know + the tiling. """ chars = 0 - images = 0 + media = 0 for message in messages: chars += len(str(message.get("role", ""))) content = message.get("content") @@ -705,10 +747,10 @@ def estimate_tokens(messages: list[dict], chars_per_token: float) -> int: if part.get("type") == "text": chars += len(part.get("text", "")) else: - images += 1 + media += 1 # A few tokens per message go to the chat template's own markup. overhead = 4 * len(messages) - return int(chars / max(chars_per_token, 1.0)) + overhead + images * 600 + return int(chars / max(chars_per_token, 1.0)) + overhead + media * 600 def _parse_sse_line(line: str): diff --git a/test_llamachat.py b/test_llamachat.py index 586be47..44f569e 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -229,12 +229,28 @@ def test_attachment_truncation(): print("ok attachment truncation") +def test_audio_load(): + """A .wav file loads as an audio attachment, not as text or unknown.""" + with tempfile.TemporaryDirectory() as tmp: + src = Path(tmp) / "note.wav" + src.write_bytes(b"RIFF\x00\x00\x00\x00WAVEfmt ") + att = backend.load_attachment(src, 1000) + assert att.kind == "audio" + assert att.b64 + assert not att.b64.startswith("data:") + assert att.thumb is None + assert att.text == "" + print("ok audio attachment load") + + def test_classify(): assert backend.classify(Path("a.py")) == "text" assert backend.classify(Path("a.SlackBuild")) == "text" assert backend.classify(Path("a.png")) == "image" assert backend.classify(Path("a.jpg")) == "image" assert backend.classify(Path("a.so")) == "unknown" + assert backend.classify(Path("a.wav")) == "audio" + assert backend.classify(Path("a.WAV")) == "audio" print("ok file classification") @@ -376,6 +392,25 @@ def test_user_content(): assert content[0]["type"] == "text" assert content[1]["type"] == "image_url" assert content[1]["image_url"]["url"].startswith("data:image/png") + + # Audio -> the input_audio part llama.cpp expects, raw base64. The + # data must not carry a data: prefix, unlike the image URL above. + rec = backend.audio_attachment(b"\x00\x01" * 8, seconds=2) + assert rec.kind == "audio" + assert rec.path.name == "voice-note-2s.wav" + assert rec.sha256 and rec.size == 16 + + content = backend.build_user_content("listen", [rec]) + assert isinstance(content, list) + assert content[0] == {"type": "text", "text": "listen"} + assert content[1]["type"] == "input_audio" + assert content[1]["input_audio"]["format"] == "wav" + assert content[1]["input_audio"]["data"] == rec.b64 + assert not content[1]["input_audio"]["data"].startswith("data:") + + # Image and audio together keep both parts, text first. + both = backend.build_user_content("look and listen", [img, rec]) + assert [p["type"] for p in both] == ["text", "image_url", "input_audio"] print("ok user content assembly") @@ -758,6 +793,18 @@ def test_token_estimate(): ] assert backend.estimate_tokens(with_image, 3.5) > 500 + # An audio part costs like an image, not free. + with_audio = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "listen"}, + {"type": "input_audio", "input_audio": {"data": "AAA", "format": "wav"}}, + ], + } + ] + assert backend.estimate_tokens(with_audio, 3.5) > 500 + # A silly ratio must not divide by zero. assert backend.estimate_tokens(small, 0) > 0 print("ok token estimate") @@ -3867,6 +3914,7 @@ if __name__ == "__main__": test_fts_query_escaping() test_history_roundtrip() test_attachment_truncation() + test_audio_load() test_classify() test_sse_parsing() test_reasoning_storage() |