skills/tests/test_hub_speech.py

from __future__ import annotations

import dataclasses
import hashlib
import io
import json
import os
import subprocess
import sys
import unittest
import wave
from unittest import mock

from hub_case import HubCase
from hub_engines import OPERATOR_VOICE, VOICES, log_rows, wav_bytes
from hub_support import HOST_STATE, HUB, server
from kvb_hub import council, engines, exchange, handler, phrasing


class TestHubSpeech(HubCase):
    def test_voices_come_from_the_models_and_stick_to_their_session(self):
        speech = self.httpd.hub.speech
        self.assertEqual(list(speech.voices), VOICES)
        self.assertEqual((speech.voices["de_DE-solo-medium:0"], speech.voices["en_US-test-high:3"]), (("de_DE-solo-medium", None), ("en_US-test-high", 3)))
        pool = [voice for voice in VOICES if voice != OPERATOR_VOICE]
        self.assertEqual((speech.pool, speech.operator), (pool, OPERATOR_VOICE))
        keys = [f"session-{n}" for n in range(40)] + ["s1", "name:Docs"]
        picked = [engines.voice_of(key, pool, OPERATOR_VOICE) for key in keys]
        self.assertEqual(picked, [pool[int(hashlib.sha256(key.encode()).hexdigest(), 16) % len(pool)] for key in keys])
        self.assertEqual(picked, [speech.voice_of(key) for key in keys])
        self.assertEqual(picked, [engines.voice_of(key, list(pool), OPERATOR_VOICE) for key in keys])
        self.assertEqual(set(picked), set(pool))
        self.assertEqual(speech.voice_of("Operator"), OPERATOR_VOICE)
        self.assertEqual(engines.voice_of("s1", [], OPERATOR_VOICE), OPERATOR_VOICE)
        self.assertEqual(engines.voice_of("s1", [], ""), "")
        with self.assertRaises(server.ToolError):
            engines.Speech(dataclasses.replace(self.engines, operator_voice="en_US-test-high:4"), self.state / "other")
        silent = engines.Speech(dataclasses.replace(self.engines, voices=self.outside, operator_voice="anything"), self.state / "silent")
        self.assertEqual((silent.speaks(), silent.voice_of("s1"), silent.describe()["operator"], silent.describe()["tts"]), (False, "", "", False))
        deaf = engines.Speech(dataclasses.replace(self.engines, whisper=self.outside / "missing"), self.state / "deaf")
        self.assertFalse(deaf.hears())

    def test_speech_output_is_guarded_whitelisted_capped_and_cached(self):
        good = self.guard(self.host)
        calls = self.engines.scratch / "piper.jsonl"
        cache = self.state / "tts"
        request = {"text": "all stations, this is Docs, board, out", "voice": "en_US-test-high:2"}
        for headers in ({k: v for k, v in good.items() if k != "X-Kvb-Hub"}, {k: v for k, v in good.items() if k != "Origin"}, {**good, "Origin": "http://evil.example"}):
            with self.subTest(headers=headers):
                self.assertEqual(self.post_json("/api/tts", request, headers)[0], 403)
        for broken in (
            {**request, "voice": "en_US-test-high:4"},
            {**request, "voice": "zero:0"},
            {**request, "voice": "bad name!:0"},
            {**request, "voice": "broken:0"},
            {**request, "voice": ["en_US-test-high:2"]},
            {**request, "text": "x" * (phrasing.SPEECH_LIMIT + 1)},
            {**request, "text": " \n\t "},
            {**request, "text": "bell \x07 here"},
            {**request, "text": 7},
            {**request, "extra": 1},
            {"text": request["text"]},
        ):
            with self.subTest(broken=broken):
                self.assertEqual(self.post_json("/api/tts", broken, good)[0], 400)
        self.assertEqual(log_rows(calls), [])
        status, headers, body = self.post_json("/api/tts", request, good)
        self.assertEqual((status, headers["content-type"], headers["content-security-policy"]), (200, "audio/wav", exchange.PLAIN_CSP), body)
        with wave.open(io.BytesIO(body)) as reader:
            self.assertEqual((reader.getnchannels(), reader.getframerate(), reader.getnframes()), (1, 22050, len(request["text"]) + 1))
        key = hashlib.sha256(f"{request['voice']}\n{request['text']}".encode()).hexdigest()
        self.assertEqual(sorted(path.name for path in cache.iterdir()), [f"{key}.wav"])
        self.assertEqual(self.post_json("/api/tts", {**request, "text": "  all stations,\nthis is Docs, board, out "}, good)[2], body)
        self.assertEqual(self.post_json("/api/tts", {**request, "text": "x" * phrasing.SPEECH_LIMIT}, good)[0], 200)
        rows = log_rows(calls)
        model = self.engines.voices / "en_US-test-high.onnx"
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0]["argv"][:5], ["--model", str(model), "--config", f"{model}.json", "--output-file"])
        self.assertEqual(rows[0]["argv"][6:], ["--speaker", "2"])
        self.assertTrue(rows[0]["argv"][5].startswith(f"{cache}/.{key}-"))
        self.assertEqual((rows[0]["text"], rows[0]["tmp"], rows[0]["cwd"]), (f"{request['text']}\n", str(self.engines.scratch), str(self.engines.scratch)))
        self.assertEqual(self.post_json("/api/tts", {"text": "solo", "voice": "de_DE-solo-medium:0"}, good)[0], 200)
        self.assertNotIn("--speaker", log_rows(calls)[-1]["argv"])
        status, _, body = self.post_json("/api/tts", {**request, "text": "FAIL now"}, good)
        self.assertEqual(status, 502, body)
        self.assertIn(b"exit code 3", body)
        self.assertEqual([path.name for path in cache.iterdir() if path.name.startswith(".")], [])
        with mock.patch.object(engines, "TTS_KEEP", 2):
            for number in range(3):
                self.assertEqual(self.post_json("/api/tts", {**request, "text": f"line {number}"}, good)[0], 200)
        kept = {hashlib.sha256(f"{request['voice']}\nline {number}".encode()).hexdigest() + ".wav" for number in (1, 2)}
        self.assertEqual({path.name for path in cache.iterdir()}, kept)
        with mock.patch.object(self.httpd.hub.speech, "voices", {}):
            self.assertEqual(self.post_json("/api/tts", request, good)[0], 503)
        self.assertEqual(self.scripts.calls, [])

    def test_speech_input_checks_the_guard_type_size_and_format(self):
        good = self.guard(self.host)
        calls = self.engines.scratch / "whisper.jsonl"
        clip = wav_bytes(1, 2, engines.STT_RATE, b"\x01\x00" * engines.STT_RATE)

        def post(body: bytes, headers: dict[str, str]) -> tuple[int, dict, bytes]:
            return self.call("POST", "/api/stt", body, {**headers, "Content-Type": "audio/wav"})

        for headers in ({k: v for k, v in good.items() if k != "X-Kvb-Hub"}, {k: v for k, v in good.items() if k != "Origin"}, {**good, "Origin": "http://evil.example"}):
            with self.subTest(headers=headers):
                self.assertEqual(post(clip, headers)[0], 403)
        for content_type in ("application/octet-stream", "audio/x-wav", "text/plain", "application/json"):
            with self.subTest(content_type=content_type):
                self.assertEqual(self.call("POST", "/api/stt", clip, {**good, "Content-Type": content_type})[0], 415)
        self.assertEqual(self.call("POST", "/api/stt", clip, good)[0], 415)
        self.assertEqual(self.call("POST", "/api/stt", b"RIFF", {**good, "Content-Type": "audio/wav", "Content-Length": str(engines.STT_LIMIT + 1)})[0], 413)
        for broken in (
            b"not a wav",
            wav_bytes(1, 2, 44100, b"\x00\x00" * 100),
            wav_bytes(2, 2, engines.STT_RATE, b"\x00\x00" * 200),
            wav_bytes(1, 1, engines.STT_RATE, b"\x00" * 100),
            wav_bytes(1, 2, engines.STT_RATE, b""),
        ):
            with self.subTest(broken=broken[:40]):
                self.assertEqual(post(broken, good)[0], 400)
        with mock.patch.object(engines, "STT_SECONDS", 1):
            self.assertEqual(post(wav_bytes(1, 2, engines.STT_RATE, b"\x00\x00" * (engines.STT_RATE + 1)), good)[0], 400)
        self.assertEqual(log_rows(calls), [])
        status, headers, body = post(clip, good)
        self.assertEqual((status, headers["content-type"]), (200, "application/json"), body)
        self.assertEqual(json.loads(body), {"text": "hello from the mic second line"})
        argv = ["--model", str(self.engines.whisper_model), "--file", "-", "--output-txt", "--language", "auto", "--no-prints", "--suppress-nst"]
        self.assertEqual(log_rows(calls), [{"argv": argv, "sha256": hashlib.sha256(clip).hexdigest()}])
        long = wav_bytes(1, 2, engines.STT_RATE, b"\x00\x00" * engines.STT_RATE * 40)
        self.assertGreater(len(long), handler.BODY_LIMIT)
        self.assertEqual(post(long, good)[0], 200)
        status, _, body = post(wav_bytes(1, 2, engines.STT_RATE, b"FAIL" * 100), good)
        self.assertEqual(status, 502, body)
        with mock.patch.object(self.httpd.hub.speech, "hears", return_value=False):
            self.assertEqual(post(clip, good)[0], 503)
        self.assertEqual(self.scripts.calls, [])
        self.assertEqual(self.get("/api/stt")[0], 405)

    def test_the_server_requires_every_engine_argument(self):
        args = {
            "--bind": "127.0.0.1", "--listen-port": "0", "--public-port": "1", "--council": str(self.council), "--pages": str(self.pages), "--state": str(self.state),
            "--host-state": str(HOST_STATE), "--skills": str(self.outside), "--transcripts": str(self.projects), "--templates": str(HUB / "templates"), "--repos": str(self.repos),
            "--worktrees": str(self.worktrees), "--timers": str(self.timers), "--llm-key-file": os.devnull, "--llm-url": "", "--llm-model": "",
            "--piper": str(self.engines.piper), "--voices": str(self.engines.voices), "--operator-voice": "en_US-test-high:9", "--whisper": str(self.engines.whisper),
            "--whisper-model": str(self.engines.whisper_model), "--scratch": str(self.engines.scratch), "--approvers": str(self.approvers),
        }

        def run(given: dict[str, str]) -> subprocess.CompletedProcess:
            argv = [part for pair in given.items() for part in pair]
            return subprocess.run([sys.executable, str(HUB / "server.py"), *argv], capture_output=True, text=True, timeout=60, check=False)

        for name in ("--piper", "--voices", "--operator-voice", "--whisper", "--whisper-model", "--scratch", "--approvers"):
            with self.subTest(name=name):
                done = run({key: value for key, value in args.items() if key != name})
                self.assertEqual(done.returncode, 2, done.stderr)
                self.assertIn(f"required: {name}", done.stderr)
        done = run(args)
        self.assertEqual(done.returncode, 2, done.stderr)
        self.assertIn("the operator voice 'en_US-test-high:9' is not installed", done.stderr)


class TestPhraseology(unittest.TestCase):
    def test_every_class_speaks_in_maritime_phraseology(self):
        cases = (
            (("mayday", "Docs", [], "MAYDAY disk full"), "MAYDAY MAYDAY MAYDAY, this is Docs Docs Docs, disk full, over"),
            (("mayday", "Docs", ["Skills"], "mayday, mayday: disk full"), "MAYDAY MAYDAY MAYDAY, this is Docs Docs Docs, disk full, over"),
            (("pan-pan", "Docs", ["Skills"], "PAN-PAN: slow"), "PAN-PAN PAN-PAN PAN-PAN, all stations, this is Docs, slow, over"),
            (("security", "Maintainer", ["Skills", "Docs"], "SECURITE rotate"), "SECURITE SECURITE SECURITE, all stations, this is Maintainer, rotate, over"),
            (("direct", "Skills", ["Maintainer"], "ready"), "Maintainer, Maintainer, this is Skills, ready, over"),
            (("direct", "Skills", ["Alpha", "Beta"], "ready"), "Alpha and Beta, Alpha and Beta, this is Skills, ready, over"),
            (("direct", "Docs", ["Ops"], "MAYDAY is only a word here"), "Ops, Ops, this is Docs, MAYDAY is only a word here, over"),
            (("direct", "Skills", [], "lost target"), "all stations, this is Skills, lost target, out"),
            (("broadcast", "Maintainer", [], "board"), "all stations, this is Maintainer, board, out"),
            (("broadcast", "Operator", [], "[kvb-hub msg 0123456789abcdef] deploy now file: /home/op/x.txt"), "all stations, this is Operator, deploy now, out"),
            (("broadcast", "Docs", [], "  two\n lines\t here "), "all stations, this is Docs, two lines here, out"),
            (("broadcast", "", [], "-"), "all stations, this is unknown station, out"),
        )
        for arguments, expected in cases:
            with self.subTest(arguments=arguments):
                self.assertEqual(phrasing.phrase(*arguments), expected)

    def test_a_phrase_fits_the_speech_limit_on_one_line(self):
        for cls in council.CLASSES:
            with self.subTest(cls=cls):
                spoken = phrasing.phrase(cls, "n" * 900, ["a\nb " * 300] * 9, "word\n" * 9000)
                self.assertLessEqual(len(spoken), phrasing.SPEECH_LIMIT)
                self.assertTrue(phrasing.SPOKEN_LINE.fullmatch(spoken))
        self.assertTrue(phrasing.phrase("broadcast", "Docs", [], "word " * 300).endswith("word, out"))


if __name__ == "__main__":
    unittest.main()