music-assistant-server

1"""
2OpenAI Text-to-speech provider for Music Assistant.
3
4Renders speech through the OpenAI speech API and any self-hostable server implementing
5the same endpoint (Kokoro-FastAPI, LocalAI, Speaches, ...), exposing each voice as a TTS
6engine. Rendered clips are cached on disk and served from a route on the streamserver.
7"""
8
9from __future__ import annotations
10
11import asyncio
12import hashlib
13import os
14import re
15import shutil
16import time
17from contextlib import suppress
18from pathlib import Path
19from typing import TYPE_CHECKING
20
21import aiofiles
22from aiofiles.os import makedirs, remove, replace
23from aiofiles.os import path as aiopath
24from aiohttp import ClientTimeout, web
25from music_assistant_models.config_entries import ConfigEntry
26from music_assistant_models.enums import (
27    ConfigEntryType,
28    ContentType,
29    MediaType,
30    ProviderFeature,
31    StreamType,
32)
33from music_assistant_models.errors import AudioError
34from music_assistant_models.media_items import AudioFormat
35from music_assistant_models.streamdetails import StreamDetails
36
37from music_assistant.models.plugin import PluginProvider, TTSEngine
38
39if TYPE_CHECKING:
40    from collections.abc import Callable
41
42    from aiohttp import ClientSession
43    from music_assistant_models.config_entries import ProviderConfig
44    from music_assistant_models.provider import ProviderManifest
45
46    from music_assistant.mass import MusicAssistant
47    from music_assistant.models import ProviderInstanceType
48
49CONF_BASE_URL = "base_url"
50CONF_API_KEY = "api_key"
51CONF_MODEL = "model"
52CONF_VOICES = "voices"
53
54DEFAULT_BASE_URL = "https://api.openai.com/v1"
55DEFAULT_MODEL = "tts-1"
56# the voices the tts-1 model supports, used when the backend does not advertise its own
57DEFAULT_VOICES = ("alloy", "echo", "fable", "nova", "onyx", "shimmer")
58# the one response format all compatible backends implement
59RESPONSE_FORMAT = "mp3"
60# rendered clips older than this are removed from the on-disk cache
61CACHE_MAX_AGE = 24 * 3600
62# the shared http session has no default timeout: discovery runs during provider load so
63# it must give up quickly, while rendering can legitimately be slow on a cpu-bound backend
64VOICES_TIMEOUT = ClientTimeout(total=10)
65SPEECH_TIMEOUT = ClientTimeout(total=120)
66
67# clip file names are sha256 hexdigests; anything else in the cache directory is not ours
68FILE_ID_PATTERN = re.compile(r"^[0-9a-f]{64}$")
69
70SUPPORTED_FEATURES = {ProviderFeature.TTS}
71
72
73async def setup(
74    mass: MusicAssistant, manifest: ProviderManifest, config: ProviderConfig
75) -> ProviderInstanceType:
76    """Initialize provider(instance) with given configuration."""
77    return OpenAITTSProvider(mass, manifest, config, SUPPORTED_FEATURES)
78
79
80async def fetch_backend_voices(
81    http_session: ClientSession, base_url: str, api_key: str = ""
82) -> list[str]:
83    """
84    Return the voices the backend advertises, empty when it does not advertise any.
85
86    Only some of the self-hostable servers implement a voice listing; the OpenAI cloud
87    API does not. Never raises: an unreachable or silent backend yields an empty list.
88
89    :param http_session: The HTTP session to request with.
90    :param base_url: The API endpoint, without trailing slash.
91    :param api_key: The API key to authenticate with, empty for backends without auth.
92    """
93    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
94    try:
95        async with http_session.get(
96            f"{base_url}/audio/voices", headers=headers, timeout=VOICES_TIMEOUT
97        ) as response:
98            response.raise_for_status()
99            data = await response.json()
100    except Exception:
101        return []
102    items = data.get("voices") if isinstance(data, dict) else data
103    if not isinstance(items, list):
104        return []
105    # per item, so one unexpected entry does not discard the voices around it
106    voices: list[str] = []
107    for item in items:
108        name: object
109        if isinstance(item, str):
110            name = item
111        elif isinstance(item, dict):
112            name = item.get("id") or item.get("name")
113        else:
114            continue
115        if isinstance(name, str) and name:
116            voices.append(name)
117    return voices
118
119
120class OpenAITTSProvider(PluginProvider):
121    """Text-to-speech provider backed by the OpenAI (compatible) speech API."""
122
123    _cache_dir: str
124    _render_lock: asyncio.Lock
125    _voices: list[str]
126    # rendered clips by file id; the route only serves paths from this index
127    _clips: dict[str, str]
128    _unregister_route: Callable[[], None] | None = None
129
130    async def handle_async_init(self) -> None:
131        """Handle async initialization of the provider."""
132        # scoped per instance: different endpoints can render the same input differently
133        self._cache_dir = os.path.join(self.mass.cache_path, self.domain, self.instance_id)
134        self._render_lock = asyncio.Lock()
135        self._voices = await self._resolve_voices()
136        # clips rendered before a restart stay playable
137        self._clips = await self._index_cache()
138        self._unregister_route = self.mass.streams.register_dynamic_route(
139            self._route_path, self._handle_speech_request
140        )
141
142    async def unload(self, is_removed: bool = False) -> None:
143        """
144        Handle unload/close of the provider.
145
146        Called when provider is deregistered (e.g. MA exiting or config reloading).
147        is_removed will be set to True when the provider is removed from the configuration.
148        """
149        await super().unload(is_removed)
150        if unregister := self._unregister_route:
151            self._unregister_route = None
152            unregister()
153        if is_removed:
154            await asyncio.to_thread(shutil.rmtree, self._cache_dir, ignore_errors=True)
155
156    async def get_config_entries(self) -> tuple[ConfigEntry, ...]:
157        """
158        Return the (options) config entries for this provider instance.
159
160        The connection details (endpoint, api key and model) are collected by the setup
161        flow (see setup_flow.py); the only option here overrides the list of voices that
162        is exposed as TTS engines.
163        """
164        return (
165            ConfigEntry(
166                key=CONF_VOICES,
167                type=ConfigEntryType.STRING,
168                # no options: the override exists to name voices we could not discover,
169                # and a populated options list would render as a strict select
170                multi_value=True,
171                required=False,
172                default_value=[],
173                category="features",
174                requires_reload=True,
175            ),
176        )
177
178    async def get_tts_engines(self) -> list[TTSEngine]:
179        """Return one TTS engine per available voice."""
180        return [TTSEngine(id=voice, name=voice, provider=self) for voice in self._voices]
181
182    async def get_tts_message(
183        self, message: str, language: str | None = None, engine_id: str | None = None
184    ) -> StreamDetails:
185        """
186        Render the given message as speech.
187
188        :param message: The text to convert to speech.
189        :param language: Ignored: the speech endpoint has no language parameter, the
190            backend infers the language from the message itself.
191        :param engine_id: The voice to render with; defaults to the first available voice.
192        :return: StreamDetails for the (cached) audio clip.
193        """
194        voice = engine_id or self._voices[0]
195        file_id = await self._render_speech(message, voice)
196        return StreamDetails(
197            provider=self.instance_id,
198            item_id=file_id,
199            audio_format=AudioFormat(content_type=ContentType.MP3),
200            media_type=MediaType.SOUND_EFFECT,
201            stream_type=StreamType.HTTP,
202            path=f"{self.mass.streams.base_url}{self._route_path}?id={file_id}",
203        )
204
205    @property
206    def _route_path(self) -> str:
207        """Return the path of this instance's dynamic route on the streamserver."""
208        return f"/{self.instance_id}_speech"
209
210    @property
211    def _base_url(self) -> str:
212        """Return the configured API endpoint, without trailing slash."""
213        return str(self.get_setup_value(CONF_BASE_URL, DEFAULT_BASE_URL)).rstrip("/")
214
215    async def _handle_speech_request(self, request: web.Request) -> web.FileResponse | web.Response:
216        """Serve a rendered speech clip by its file id."""
217        if not (file_id := request.query.get("id")):
218            return web.Response(status=400, text="Missing id")
219        # the id from the url is only a lookup key, never a path component
220        if not (file_path := self._clips.get(file_id)):
221            raise web.HTTPNotFound
222        if not await aiopath.isfile(file_path):
223            self._clips.pop(file_id, None)
224            raise web.HTTPNotFound
225        return web.FileResponse(file_path)
226
227    async def _render_speech(self, message: str, voice: str) -> str:
228        """
229        Render the message to a file in the cache directory and return its file id.
230
231        :param message: The text to convert to speech.
232        :param voice: The voice to render with.
233        """
234        model = str(self.get_setup_value(CONF_MODEL, DEFAULT_MODEL))
235        # the endpoint is part of the identity: repointing an instance must not reuse clips
236        digest = f"{self._base_url}\0{model}\0{voice}\0{message}"
237        file_id = hashlib.sha256(digest.encode()).hexdigest()
238        file_path = os.path.join(self._cache_dir, f"{file_id}.{RESPONSE_FORMAT}")
239        async with self._render_lock:
240            if await aiopath.isfile(file_path):
241                # a reused clip is played right away, so keep it out of reach of the reaper
242                with suppress(OSError):
243                    await asyncio.to_thread(os.utime, file_path, None)
244                self._clips[file_id] = file_path
245                return file_id
246            await makedirs(self._cache_dir, exist_ok=True)
247            audio_data = await self._request_speech(message, voice)
248            # write to a temp file first so an interrupted render leaves no partial clip
249            tmp_path = f"{file_path}.tmp"
250            try:
251                async with aiofiles.open(tmp_path, "wb") as _file:
252                    await _file.write(audio_data)
253                await replace(tmp_path, file_path)
254            except OSError:
255                with suppress(OSError):
256                    await remove(tmp_path)
257                raise
258            self._clips[file_id] = file_path
259        await self._reap_cache()
260        return file_id
261
262    async def _request_speech(self, message: str, voice: str) -> bytes:
263        """
264        Request the rendered audio for the message from the speech API.
265
266        :param message: The text to convert to speech.
267        :param voice: The voice to render with.
268        """
269        api_key = self.get_setup_value(CONF_API_KEY)
270        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
271        payload = {
272            "model": str(self.get_setup_value(CONF_MODEL, DEFAULT_MODEL)),
273            "voice": voice,
274            "input": message,
275            "response_format": RESPONSE_FORMAT,
276        }
277        async with self.mass.http_session.post(
278            f"{self._base_url}/audio/speech",
279            headers=headers,
280            json=payload,
281            timeout=SPEECH_TIMEOUT,
282        ) as response:
283            if response.status != 200:
284                detail = (await response.text())[:200]
285                raise AudioError(f"Speech request failed (HTTP {response.status}): {detail}")
286            return await response.read()
287
288    async def _reap_cache(self) -> None:
289        """Remove rendered clips that are older than the maximum cache age."""
290
291        def _reap() -> list[str]:
292            reaped: list[str] = []
293            cutoff = time.time() - CACHE_MAX_AGE
294            with os.scandir(self._cache_dir) as entries:
295                for entry in entries:
296                    with suppress(OSError):
297                        if (
298                            entry.is_file(follow_symlinks=False)
299                            and entry.stat(follow_symlinks=False).st_mtime < cutoff
300                        ):
301                            Path(entry.path).unlink()
302                            reaped.append(Path(entry.name).stem)
303            return reaped
304
305        # reaping is opportunistic housekeeping and must never break playback
306        with suppress(OSError):
307            for file_id in await asyncio.to_thread(_reap):
308                self._clips.pop(file_id, None)
309
310    async def _index_cache(self) -> dict[str, str]:
311        """Return the clips present in the cache directory, keyed by file id."""
312
313        def _index() -> dict[str, str]:
314            # only adopt what this provider could have written itself: anything else that
315            # ended up in the directory must never become reachable through the route
316            with os.scandir(self._cache_dir) as entries:
317                return {
318                    Path(entry.name).stem: entry.path
319                    for entry in entries
320                    if entry.is_file(follow_symlinks=False)
321                    and entry.name.endswith(f".{RESPONSE_FORMAT}")
322                    and FILE_ID_PATTERN.match(Path(entry.name).stem)
323                }
324
325        try:
326            return await asyncio.to_thread(_index)
327        except OSError:
328            # nothing rendered yet: the cache directory is created on the first render
329            return {}
330
331    async def _resolve_voices(self) -> list[str]:
332        """Return the voices to expose, from the config override, the backend or the defaults."""
333        if voices := self._voice_override():
334            self.logger.debug("Using %s voice(s) from the config override", len(voices))
335            return voices
336        api_key = str(self.get_setup_value(CONF_API_KEY) or "")
337        if voices := await fetch_backend_voices(self.mass.http_session, self._base_url, api_key):
338            self.logger.debug("Discovered %s voice(s) on the backend", len(voices))
339            return voices
340        self.logger.debug("Falling back to the default voices")
341        return list(DEFAULT_VOICES)
342
343    def _voice_override(self) -> list[str]:
344        """Return the configured voices, empty when the override is unset."""
345        configured = self.get_config_value(CONF_VOICES)
346        values = configured if isinstance(configured, list) else [configured]
347        # a single value can hold a comma separated list of its own, e.g. when the whole
348        # list was pasted into one field
349        names = [
350            name
351            for value in values
352            if isinstance(value, str)
353            for part in value.split(",")
354            if (name := part.strip())
355        ]
356        return list(dict.fromkeys(names))
357