Melvin Strobl
16383ea6e2
- `stt`: one realtime transcription session per utterance on the speech
server's url, opened at `listening` true or the first chunk, committed at
false; a whole clip goes to /v1/audio/transcriptions. Returns at once and
publishes `partial` and `text` when the server answers.
- `tts`: one /v1/audio/speech call per sentence, cached by voice and text in
node state; a reference is marked `append` while earlier ones still play.
- wren: `reply` says what was applied, `unhandled` carries the prompt when
nothing was. The install hint names the git pin (PyPI's `wren` is another
package).
- ai/client: transcribe, speak and transcribe_stream beside stream_completion.
- A device socket publishes scalar `in` ports as {"type": "publish"} text,
in order with its binary frames.
- Pipeline.publish walks only what its values reach: a node downstream of a
consumer that published nothing no longer runs again (wren re-acted on
every microphone chunk).
354 lines
14 KiB
Python
354 lines
14 KiB
Python
"""Speech: a microphone's audio to text, and text to audio a speaker plays.
|
|
|
|
Both talk to an OpenAI-compatible speech server, such as `parrot serve` kept
|
|
running by a Service node, whose address arrives on a ``url`` input. An empty
|
|
url is a server that is not answering: the node reads degraded and drops what
|
|
it gets.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io
|
|
import logging
|
|
import re
|
|
import threading
|
|
import time
|
|
import wave
|
|
from collections.abc import Coroutine, Iterable
|
|
from concurrent.futures import Future
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from fluksio.ai import client
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.nodes.base import Node
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: Seconds a session waits for its next chunk, and the most one utterance may
|
|
#: take from its first chunk to its text. Past either it is given up.
|
|
IDLE_S = 10.0
|
|
UTTERANCE_S = 60.0
|
|
#: The rate of a microphone's PCM chunks, the panels' (openpanel V3).
|
|
RATE = 16000
|
|
#: How many spoken sentences a tts node keeps.
|
|
SPOKEN = 256
|
|
#: How long after its estimated end a sentence still counts as playing: the
|
|
#: device fetches it before it starts.
|
|
PLAYING_SLACK_S = 1.0
|
|
|
|
|
|
def _check(specs: list[MessageSpec], ports: dict[str, DType], node: str) -> None:
|
|
"""Refuse a port the node does not have, or one of another type."""
|
|
for spec in specs:
|
|
dtype = ports.get(spec.port)
|
|
if dtype is None:
|
|
raise ValueError(
|
|
f"'{spec.port}' is not a port of {node}. Its ports are "
|
|
f"{', '.join(ports)}."
|
|
)
|
|
if spec.dtype is not dtype:
|
|
raise ValueError(f"{node}'s '{spec.port}' is {dtype.value}.")
|
|
|
|
|
|
def _is_pcm(audio: dict[str, Any]) -> bool:
|
|
"""A microphone chunk rather than a whole clip."""
|
|
media_type = str(audio.get("media_type") or "").lower()
|
|
return media_type.startswith(("audio/l16", "audio/pcm"))
|
|
|
|
|
|
def sentences(text: str) -> list[str]:
|
|
"""``text`` cut after each ``.``, ``!`` and ``?``, one call to speak each."""
|
|
return [s for s in (p.strip() for p in re.split(r"(?<=[.!?])\s+", text)) if s]
|
|
|
|
|
|
class _Speech(Node):
|
|
"""What both nodes share: the url, the artifact store and their health."""
|
|
|
|
__slots__ = ("_artifacts", "_loop")
|
|
|
|
def __init__(
|
|
self,
|
|
f: Any,
|
|
requires: Iterable[MessageSpec],
|
|
provides: Iterable[MessageSpec],
|
|
params: dict[str, Any] | None,
|
|
name: str,
|
|
inputs: dict[str, DType],
|
|
outputs: dict[str, DType],
|
|
):
|
|
# Set on start; a test sets them directly.
|
|
self._artifacts: Any = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
super().__init__(
|
|
f, requires=requires, provides=provides, params=params, name=name
|
|
)
|
|
_check(self.input_ports, inputs, name)
|
|
_check(self.output_ports, outputs, name)
|
|
if not any(spec.port == "url" for spec in self.input_ports):
|
|
raise ValueError(
|
|
f"{name} needs a 'url' input: the speech server's address."
|
|
)
|
|
|
|
async def start(self, app: FastAPI | None = None) -> None:
|
|
self._loop = asyncio.get_running_loop()
|
|
self._artifacts = getattr(app.state, "artifact_store", None) if app else None
|
|
|
|
def _reachable(self, url: str) -> bool:
|
|
"""Whether there is a server to ask, said as the node's health."""
|
|
if not url:
|
|
if self._health != "degraded":
|
|
self.report_health("degraded", "no speech server: its url is empty")
|
|
return False
|
|
if self._health != "ok":
|
|
self.report_health("ok")
|
|
return True
|
|
|
|
def _load(self, reference: dict[str, Any]) -> bytes:
|
|
return b"".join(self._artifacts.read(str(reference["digest"])))
|
|
|
|
def _say(self, **values: Any) -> None:
|
|
"""Publish on the ports that are wired."""
|
|
wired = {spec.port for spec in self.output_ports if spec.name}
|
|
self.emit({port: value for port, value in values.items() if port in wired})
|
|
|
|
|
|
class SttNode(_Speech):
|
|
"""Speech to text: a microphone's audio in, what was said out.
|
|
|
|
Inputs: ``audio``, a stream of 16 kHz PCM16 chunks or a whole clip,
|
|
``listening``, true while the device listens, and ``url``.
|
|
Outputs: ``partial``, the words agreed on so far, and ``text``.
|
|
|
|
One realtime session per utterance, opened when ``listening`` turns true
|
|
or at the first chunk, and committed when it turns false. A clip goes to
|
|
the server whole. Either way the node returns at once, since a device's
|
|
next frame waits for it, and publishes once the server answers. An
|
|
utterance that fails publishes nothing and is logged.
|
|
"""
|
|
|
|
__slots__ = ("_heard", "_listening", "_lock", "_session", "_tasks")
|
|
|
|
def __init__(
|
|
self,
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
self._lock = threading.Lock()
|
|
#: The chunks of the utterance being spoken, None between utterances.
|
|
self._session: asyncio.Queue[bytes | None] | None = None
|
|
#: The digest last heard, and whether the device was listening.
|
|
self._heard = ""
|
|
self._listening: bool | None = None
|
|
self._tasks: set[Future[None]] = set()
|
|
inputs = [
|
|
# The url is read, not waited for: its changing is nothing to hear.
|
|
spec.model_copy(update={"trigger": spec.port != "url"})
|
|
for spec in self._normalize_ports(requires)
|
|
]
|
|
if not any(spec.port == "audio" for spec in inputs):
|
|
raise ValueError("stt needs an 'audio' input: the microphone.")
|
|
super().__init__(
|
|
self._hear,
|
|
inputs,
|
|
self._normalize_ports(provides),
|
|
params,
|
|
name or "stt",
|
|
{"audio": DType.AUDIO, "listening": DType.BOOL, "url": DType.STR},
|
|
{"partial": DType.STR, "text": DType.STR},
|
|
)
|
|
|
|
async def start(self, app: FastAPI | None = None) -> None:
|
|
await super().start(app)
|
|
if self._pipeline is None:
|
|
return
|
|
# What the inputs hold already was heard before this start.
|
|
held = await asyncio.to_thread(
|
|
self._pipeline.state.get_present, list(self.requires)
|
|
)
|
|
for spec in self.input_ports:
|
|
value = held.get(spec.name)
|
|
if spec.port == "audio" and isinstance(value, dict):
|
|
self._heard = str(value.get("digest") or "")
|
|
elif spec.port == "listening" and isinstance(value, bool):
|
|
self._listening = value
|
|
|
|
async def stop(self, app: FastAPI | None = None) -> None:
|
|
for task in list(self._tasks):
|
|
task.cancel()
|
|
self._session = None
|
|
|
|
def _hear(
|
|
self,
|
|
params: dict[str, Any],
|
|
audio: dict[str, Any] | None = None,
|
|
listening: bool | None = None,
|
|
url: str = "",
|
|
) -> None:
|
|
with self._lock:
|
|
digest = str((audio or {}).get("digest") or "")
|
|
fresh = bool(digest) and digest != self._heard
|
|
self._heard = digest or self._heard
|
|
was, self._listening = self._listening, listening
|
|
if not self._reachable(url):
|
|
self._session = None
|
|
return None
|
|
if fresh and audio is not None and not _is_pcm(audio):
|
|
self._spawn(self._clip(url, audio, self._load(audio)))
|
|
elif fresh and audio is not None:
|
|
if self._session is None:
|
|
self._session = self._open(url)
|
|
self._feed(self._session, self._load(audio))
|
|
elif listening and not was and self._session is None:
|
|
self._session = self._open(url)
|
|
elif was and not listening and self._session is not None:
|
|
self._feed(self._session, None)
|
|
self._session = None
|
|
return None
|
|
|
|
def _spawn(self, work: Coroutine[Any, Any, None]) -> None:
|
|
"""Run on the engine's loop, so this run can return at once."""
|
|
if self._loop is None:
|
|
work.close()
|
|
raise RuntimeError("stt only hears in a live flow")
|
|
task = asyncio.run_coroutine_threadsafe(work, self._loop)
|
|
self._tasks.add(task)
|
|
task.add_done_callback(self._tasks.discard)
|
|
|
|
def _open(self, url: str) -> asyncio.Queue[bytes | None]:
|
|
chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
|
|
self._spawn(self._transcribe(url, chunks))
|
|
return chunks
|
|
|
|
def _feed(self, chunks: asyncio.Queue[bytes | None], chunk: bytes | None) -> None:
|
|
assert self._loop is not None
|
|
self._loop.call_soon_threadsafe(chunks.put_nowait, chunk)
|
|
|
|
async def _transcribe(self, url: str, chunks: asyncio.Queue[bytes | None]) -> None:
|
|
words = ""
|
|
try:
|
|
async with asyncio.timeout(UTTERANCE_S):
|
|
async for final, text in client.transcribe_stream(
|
|
url + "/v1", chunks, rate=RATE, idle_s=IDLE_S
|
|
):
|
|
words = text.strip() if final else words + text
|
|
if final and words:
|
|
await asyncio.to_thread(self._say, partial=words, text=words)
|
|
elif words:
|
|
await asyncio.to_thread(self._say, partial=words)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"'%s' lost an utterance: %s: %s", self.id, type(exc).__name__, exc
|
|
)
|
|
finally:
|
|
# Ended without a commit, by a timeout or a failure: the next chunk
|
|
# opens a session of its own rather than feeding this dead one.
|
|
with self._lock:
|
|
if self._session is chunks:
|
|
self._session = None
|
|
|
|
async def _clip(self, url: str, audio: dict[str, Any], data: bytes) -> None:
|
|
try:
|
|
text = await client.transcribe(
|
|
url + "/v1", data, "speech.wav", str(audio.get("media_type"))
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("'%s' lost a clip: %s: %s", self.id, type(exc).__name__, exc)
|
|
return
|
|
if text.strip():
|
|
await asyncio.to_thread(self._say, partial=text.strip(), text=text.strip())
|
|
|
|
|
|
class TtsNode(_Speech):
|
|
"""Text to speech: each sentence of ``text`` spoken, out on ``speech``.
|
|
|
|
One call to the server per sentence, each sentence published as an
|
|
``audio`` reference as soon as it is spoken, so a reply plays before its
|
|
last sentence is. A sentence spoken before in the same voice is not spoken
|
|
again.
|
|
|
|
A reference is marked ``"append": true`` while the ones before it are
|
|
still playing, which queues it on the device. The first sentence of a
|
|
reply carries no mark and cuts off whatever was playing, however its
|
|
sentences arrive: as one text, or one value each while they are written.
|
|
"""
|
|
|
|
class Params(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
voice: str = Field(
|
|
default="af_heart",
|
|
description="The voice, one of Kokoro's such as af_heart or bm_george.",
|
|
)
|
|
|
|
__slots__ = ("cfg", "_lock", "_until")
|
|
|
|
def __init__(
|
|
self,
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
self.cfg = self.Params.model_validate(params or {})
|
|
# One reply at a time, so its sentences go out in order.
|
|
self._lock = threading.Lock()
|
|
#: When what was published so far stops playing, on the monotonic clock.
|
|
self._until = 0.0
|
|
inputs = [
|
|
spec.model_copy(update={"trigger": spec.port != "url"})
|
|
for spec in self._normalize_ports(requires)
|
|
]
|
|
if not any(spec.port == "text" for spec in inputs):
|
|
raise ValueError("tts needs a 'text' input: what to say.")
|
|
super().__init__(
|
|
self._speak,
|
|
inputs,
|
|
self._normalize_ports(provides),
|
|
params,
|
|
name or "tts",
|
|
{"text": DType.STR, "url": DType.STR},
|
|
{"speech": DType.AUDIO},
|
|
)
|
|
|
|
def _speak(self, params: dict[str, Any], text: str = "", url: str = "") -> None:
|
|
if not text.strip() or not self._reachable(url):
|
|
return None
|
|
with self._lock:
|
|
for sentence in sentences(text):
|
|
reference = self._spoken(url, sentence)
|
|
now = time.monotonic()
|
|
if now < self._until + PLAYING_SLACK_S:
|
|
reference = {**reference, "append": True}
|
|
self._until = max(now, self._until) + reference["meta"]["seconds"]
|
|
self._say(speech=reference)
|
|
return None
|
|
|
|
def _spoken(self, url: str, sentence: str) -> dict[str, Any]:
|
|
"""The reference to ``sentence`` spoken, from the cache or the server."""
|
|
key = hashlib.sha256(f"{self.cfg.voice}\n{sentence}".encode()).hexdigest()
|
|
# Kept in state, so the sweep keeps the bytes it refers to.
|
|
spoken = dict(self.recall("spoken") or {})
|
|
reference = spoken.get(key)
|
|
if reference is not None and self._artifacts.has(reference["digest"]):
|
|
return dict(reference)
|
|
data = asyncio.run(client.speak(url + "/v1", sentence, self.cfg.voice))
|
|
with wave.open(io.BytesIO(data)) as clip:
|
|
seconds = round(clip.getnframes() / clip.getframerate(), 3)
|
|
stored: dict[str, Any] = self._artifacts.put(
|
|
[data], name="speech.wav", media_type="audio/wav"
|
|
)
|
|
reference = stored | {"meta": {"seconds": seconds}}
|
|
spoken.pop(key, None)
|
|
spoken[key] = reference
|
|
self.remember("spoken", dict(list(spoken.items())[-SPOKEN:]))
|
|
return reference
|