app/backend/tests/flow/test_speech.py

"""The speech nodes: one session per utterance, and each sentence spoken once."""

import asyncio
import hashlib
from typing import Any

from fluksio.ai import client
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.speech import SttNode, TtsNode, sentences, speakable
from fluksio.flow.pipeline import Pipeline

URL = "http://speech:8100"


class Store:
    """The artifact store's surface the nodes use."""

    def __init__(self) -> None:
        self.data: dict[str, bytes] = {}

    def put(self, chunks: list[bytes], name: str, media_type: str) -> dict[str, Any]:
        data = b"".join(chunks)
        digest = "sha256:" + hashlib.sha256(data).hexdigest()
        self.data[digest] = data
        return {"digest": digest, "size": len(data), "media_type": media_type}

    def has(self, digest: str) -> bool:
        return digest in self.data

    def read(self, digest: str) -> list[bytes]:
        return [self.data[digest]]


def test_sentences_are_cut_after_their_ending():
    assert sentences(" Kitchen light on.  Heating off! Is it?\nYes") == [
        "Kitchen light on.",
        "Heating off!",
        "Is it?",
        "Yes",
    ]


def test_units_and_clock_times_are_said_as_words():
    said = speakable("22 °C and 21.5° at 07:30, 3.4 kWh, 505 W, 2 kW, 7:05 or 18:00.")
    assert said == (
        "22 degrees and 21.5 degrees at 7 30, 3.4 kilowatt hours, 505 watts, "
        "2 kilowatts, 7 oh 5 or 18 o'clock."
    )
    assert speakable("Kitchen light on, 3 Watts.") == "Kitchen light on, 3 Watts."


def test_an_utterance_is_one_session_committed_when_listening_ends(monkeypatch):
    heard: list[bytes] = []

    async def session(api_base, audio, *, rate, idle_s, prompt):
        assert (api_base, rate, prompt) == (f"{URL}/v1", 16000, "Driveway.")
        while (chunk := await audio.get()) is not None:
            heard.append(chunk)
        yield False, "turn on"
        yield False, " the"
        yield True, " Turn on the light."

    monkeypatch.setattr(client, "transcribe_stream", session)
    warmed: list[str] = []

    async def warm(api_base):
        warmed.append(api_base)

    monkeypatch.setattr(client, "warm", warm)

    async def scenario() -> None:
        node = SttNode(
            requires=[
                MessageSpec(name="mic", port="audio", dtype=DType.AUDIO),
                MessageSpec(name="listening", dtype=DType.BOOL),
                MessageSpec(name="url", dtype=DType.STR),
            ],
            provides=[
                MessageSpec(name="partial", dtype=DType.STR),
                MessageSpec(name="text", dtype=DType.STR),
            ],
            params={"prompt": "Driveway."},
        )
        node.assign_flow("voice", "stt")
        pipeline = Pipeline([node])
        store = Store()
        chunks = [
            store.put([bytes([n]) * 3200], "voice.mic", "audio/L16;rate=16000")
            for n in range(3)
        ]
        # Left over from before the engine started: not part of what is said.
        pipeline.state["voice.mic"] = chunks[0]
        await node.start()
        node._artifacts = store

        async def hear(audio: dict, listening: bool, url: str = URL) -> None:
            inputs = {"voice.mic": audio, "voice.listening": listening}
            await asyncio.to_thread(node.execute, inputs | {"voice.url": url})

        await hear(chunks[0], True)
        await hear(chunks[1], True)
        await hear(chunks[2], True)
        await hear(chunks[2], False)
        for _ in range(200):
            if "voice.text" in pipeline.state:
                break
            await asyncio.sleep(0.01)

        assert heard == [bytes([1]) * 3200, bytes([2]) * 3200]
        assert pipeline.state["voice.text"] == "Turn on the light."
        assert pipeline.state["voice.partial"] == "Turn on the light."
        assert node._session is None
        # Opened with the session, for the reply's first sentence.
        assert warmed == [f"{URL}/v1"]

        # No server, no session: the node says so instead.
        await hear(chunks[1], True, url="")
        assert node._session is None
        assert node._health == "degraded"

    asyncio.run(scenario())


def _mp3(frames: int) -> bytes:
    """The first frame LAME writes: MPEG-2 layer III at 24 kHz, mono, and the
    Xing header counting the frames, 576 samples each."""
    xing = b"Xing" + (15).to_bytes(4, "big") + frames.to_bytes(4, "big")
    return bytes.fromhex("fff384c4") + bytes(9) + xing + bytes(100)


def test_a_sentence_is_spoken_once_and_the_rest_of_a_reply_queues(monkeypatch):
    spoken: list[str] = []
    published: list[dict] = []

    async def speak(api_base, text, voice):
        spoken.append(text)
        return _mp3(25 * len(spoken))

    monkeypatch.setattr(client, "speak", speak)
    monkeypatch.setattr(TtsNode, "_say", lambda self, speech: published.append(speech))
    node = TtsNode(
        requires=[
            MessageSpec(name="reply", port="text", dtype=DType.STR),
            MessageSpec(name="url", dtype=DType.STR),
        ],
        provides=[MessageSpec(name="speech", dtype=DType.AUDIO)],
    )
    node.assign_flow("voice", "tts")
    Pipeline([node])
    node._artifacts = Store()

    async def scenario() -> None:
        node._loop = asyncio.get_running_loop()
        reply = "Kitchen light on. Heating off."
        await asyncio.to_thread(node.execute, {"voice.reply": reply, "voice.url": URL})
        # Long after both stopped playing.
        node._until = -10.0
        reply = "Kitchen light on."
        await asyncio.to_thread(node.execute, {"voice.reply": reply, "voice.url": URL})

    asyncio.run(scenario())

    assert spoken == ["Kitchen light on.", "Heating off."]
    assert ["append" in speech for speech in published] == [False, True, False]
    assert published[2] == published[0]
    assert published[1]["meta"] == {"seconds": 1.2}
    assert published[1]["media_type"] == "audio/mpeg"