music-assistant-server

1"""Just-in-time clip rendering for AI Radio."""
2# mypy: disable-error-code="attr-defined"
3
4from __future__ import annotations
5
6import asyncio
7import logging
8from dataclasses import dataclass
9from typing import TYPE_CHECKING, Any, cast
10
11from music_assistant_models.enums import ContentType, StreamType
12from music_assistant_models.errors import (
13    InvalidDataError,
14    MediaNotFoundError,
15    MusicAssistantError,
16)
17from music_assistant_models.media_items import AudioFormat
18from music_assistant_models.streamdetails import StreamDetails
19
20from music_assistant.helpers.tags import async_parse_tags
21from music_assistant.helpers.tts import query_tts_engine, resolve_tts_stream_path
22
23from .constants import (
24    ATTR_HOST_ID,
25    ATTR_MAX_CHARS,
26    ATTR_PROMPT,
27    ATTR_RENDERED_TEXT,
28    ATTR_SESSION_ID,
29    ATTR_WEB_SEARCH_MODE,
30    CLIP_STREAMDETAILS_EXPIRATION,
31    DEFERRED_PLACEHOLDERS,
32    MIN_CLIP_MEDIA_LIFETIME,
33    TTS_SERVER_ERROR_MARKERS,
34)
35from .helpers import soft_limit_text
36
37if TYPE_CHECKING:
38    from music_assistant_models.enums import MediaType
39    from music_assistant_models.queue_item import QueueItem
40
41    from music_assistant.mass import MusicAssistant
42
43    from .models import SessionState
44
45
46@dataclass(slots=True)
47class _CachedClipMedia:
48    """Media previously minted for a clip, kept until it expires."""
49
50    path: str
51    stream_type: StreamType
52    audio_format: AudioFormat
53    duration: int | None
54    minted_at: float
55
56
57class AIRadioRenderMixin:
58    """Renders an AI Radio clip at the moment MA needs its audio."""
59
60    if TYPE_CHECKING:
61        mass: MusicAssistant
62        logger: logging.Logger
63        _hosts: dict[str, dict[str, Any]]
64        _sessions: dict[str, SessionState]
65
66    _render_locks: dict[str, asyncio.Lock]
67    _media_cache: dict[str, _CachedClipMedia]
68
69    async def get_stream_details(self, item_id: str, media_type: MediaType) -> StreamDetails:
70        """
71        Render the AI Radio clip with the given id and return its StreamDetails.
72
73        :param item_id: The clip id of the queue item MA wants to play.
74        :param media_type: The media type of the requested item.
75        """
76        queue_item = self._find_clip_item(item_id)
77        if queue_item is None:
78            raise MediaNotFoundError(f"AI Radio clip {item_id} is not in any queue")
79        prompt = str(queue_item.extra_attributes.get(ATTR_PROMPT) or "")
80        if not prompt:
81            self._record_skip(queue_item, "clip has no prompt to render")
82            raise MediaNotFoundError(f"AI Radio clip {item_id} has no prompt to render")
83
84        async with self._lock_for(item_id):
85            text = str(queue_item.extra_attributes.get(ATTR_RENDERED_TEXT) or "")
86            if not text:
87                text = await self._generate_script(queue_item, prompt, item_id)
88                queue_item.extra_attributes[ATTR_RENDERED_TEXT] = text
89                # the signal is what marks the items cache dirty and schedules the persist
90                self.mass.player_queues.signal_update(queue_item.queue_id, items_changed=True)
91            media = await self._cached_clip_media(queue_item, text, item_id)
92
93        return StreamDetails(
94            provider=self.instance_id,
95            item_id=item_id,
96            audio_format=media.audio_format,
97            media_type=media_type,
98            stream_type=media.stream_type,
99            path=media.path,
100            duration=media.duration,
101            # a talk clip has nothing worth seeking to, and a seek is the one path that
102            # would re-fetch a possibly-expired HA url mid-playback
103            can_seek=False,
104            allow_seek=False,
105            # a cache hit serves a url that was minted earlier, so it may only claim the life
106            # that url has left or the stream outlives the token behind it
107            expiration=self._remaining_media_lifetime(media),
108        )
109
110    def _lock_for(self, clip_id: str) -> asyncio.Lock:
111        """Return the per-clip render lock, creating it on first use."""
112        if not hasattr(self, "_render_locks"):
113            self._render_locks = {}
114        if clip_id not in self._render_locks:
115            self._render_locks[clip_id] = asyncio.Lock()
116        return self._render_locks[clip_id]
117
118    async def _cached_clip_media(
119        self, queue_item: QueueItem, text: str, clip_id: str
120    ) -> _CachedClipMedia:
121        """Return the clip's minted media, re-minting only once the cache entry has expired."""
122        if not hasattr(self, "_media_cache"):
123            self._media_cache = {}
124        now = asyncio.get_running_loop().time()
125        cached = self._media_cache.get(clip_id)
126        if cached is not None and self._remaining_media_lifetime(cached) > MIN_CLIP_MEDIA_LIFETIME:
127            return cached
128        # the caller holds the per-clip render lock, so of the several uncoordinated paths
129        # that resolve the same clip only the first one mints; the rest hit the cache above
130        path, stream_type, audio_format, duration = await self._mint_clip_media(
131            queue_item, text, clip_id, self._tts_language()
132        )
133        media = _CachedClipMedia(path, stream_type, audio_format, duration, now)
134        # clips are minted per queue item, so without pruning the cache grows for as long as
135        # the server runs. an entry past its window can never be served again anyway
136        for expired_id in [
137            key
138            for key, entry in self._media_cache.items()
139            if now - entry.minted_at >= CLIP_STREAMDETAILS_EXPIRATION
140        ]:
141            del self._media_cache[expired_id]
142        self._media_cache[clip_id] = media
143        return media
144
145    def _remaining_media_lifetime(self, media: _CachedClipMedia) -> int:
146        """Return the seconds the given minted media is still usable for."""
147        elapsed = asyncio.get_running_loop().time() - media.minted_at
148        return max(MIN_CLIP_MEDIA_LIFETIME, round(CLIP_STREAMDETAILS_EXPIRATION - elapsed))
149
150    def _tts_language(self) -> str | None:
151        """Return the configured locale as a hyphenated language code, or None when unset."""
152        locale = self.mass.metadata.locale
153        return locale.replace("_", "-") if locale else None
154
155    def _find_clip_item(self, clip_id: str) -> QueueItem | None:
156        """Return the queue item holding the given clip, or None when no queue holds it."""
157        for queue_id in self._candidate_queue_ids(clip_id):
158            if (item := self._find_clip_in_queue(clip_id, queue_id)) is not None:
159                return item
160        return None
161
162    def _candidate_queue_ids(self, clip_id: str) -> list[str]:
163        """
164        Return the queue ids to search for a clip, the most likely one first.
165
166        The owning session knows its queue, but the session registry is empty after a
167        restart while the clip lives on in the persisted queue, so every queue stays a
168        candidate. Clip ids carry a uuid4-based session id, so a hit is unambiguous.
169        """
170        queue_ids = [queue.queue_id for queue in self.mass.player_queues.all()]
171        session = self._sessions.get(clip_id.rpartition("_")[0])
172        if session is not None and session.queue_id in queue_ids:
173            queue_ids.remove(session.queue_id)
174            queue_ids.insert(0, session.queue_id)
175        return queue_ids
176
177    def _find_clip_in_queue(self, clip_id: str, queue_id: str) -> QueueItem | None:
178        """Return the queue item holding the given clip, paging through the queue."""
179        page_size = 500
180        offset = 0
181        while True:
182            page = self.mass.player_queues.items(queue_id, limit=page_size, offset=offset)
183            if not page:
184                return None
185            for item in page:
186                if item.media_item is not None and item.media_item.item_id == clip_id:
187                    return item
188            if len(page) < page_size:
189                return None
190            offset += page_size
191
192    async def _generate_script(self, queue_item: QueueItem, prompt: str, clip_id: str) -> str:
193        """Resolve the deferred placeholders and generate the spoken script."""
194        attributes = queue_item.extra_attributes
195        deferred = await self._resolve_deferred_placeholders(prompt)
196        resolved = prompt
197        for key, value in deferred.items():
198            resolved = resolved.replace(key, value)
199        host = self._hosts.get(str(attributes.get(ATTR_HOST_ID) or "")) or {}
200        instructions = str(host.get("instructions") or "")
201        max_chars = int(attributes.get(ATTR_MAX_CHARS) or 0)
202        web_mode = str(attributes.get(ATTR_WEB_SEARCH_MODE) or "disabled")
203        try:
204            text = cast(
205                "str",
206                await self._generate_text(
207                    instructions=instructions, prompt=resolved, web_mode=web_mode
208                ),
209            )
210        except Exception as err:
211            self.logger.warning(
212                "AI Radio clip %s (%s) failed to generate: %s", clip_id, queue_item.name, err
213            )
214            self._record_skip(queue_item, f"generation failed: {err}")
215            raise MediaNotFoundError(f"AI Radio clip {clip_id} failed to generate") from err
216        if max_chars > 0:
217            text = soft_limit_text(text, max_chars=max_chars)
218        self.logger.debug(
219            "AI Radio clip %s (%s) rendered: %d chars", clip_id, queue_item.name, len(text)
220        )
221        return text
222
223    async def _resolve_deferred_placeholders(self, prompt: str) -> dict[str, str]:
224        """Return freshly resolved values for the placeholders deferred until airtime."""
225        values = dict.fromkeys(DEFERRED_PLACEHOLDERS, "")
226        values["<timestamp>"] = self._configured_now().strftime("%Y-%m-%d %H:%M %Z")
227        # weather is the only deferred placeholder that costs a network round-trip, so it is
228        # only fetched when the prompt actually references it
229        weather_tokens = ("<weather_hourly>", "<weather_daily>")
230        if any(token in prompt for token in weather_tokens):
231            values.update(await self._prepare_weather_tokens())
232        return values
233
234    async def _mint_clip_media(
235        self, queue_item: QueueItem, text: str, clip_id: str, language: str | None = None
236    ) -> tuple[str, StreamType, AudioFormat, int | None]:
237        """Convert the script to playable audio via the configured TTS engine."""
238        host = self._hosts.get(str(queue_item.extra_attributes.get(ATTR_HOST_ID) or "")) or {}
239        engine_uid = str(host.get("tts_engine") or "") or None
240        try:
241            path, stream_type, audio_format = await self._render_tts_media(
242                text, engine_uid, language
243            )
244            # the probe is the first fetch, so a failed render surfaces here and not in playback
245            duration = await self._probe_duration(path)
246            return path, stream_type, audio_format, duration
247        except Exception as err:
248            self.logger.warning("AI Radio clip %s failed TTS: %s", clip_id, err)
249            self._record_skip(queue_item, f"TTS failed: {err}")
250            raise MediaNotFoundError(f"AI Radio clip {clip_id} failed TTS") from err
251
252    async def _render_tts_media(
253        self, text: str, engine_uid: str | None = None, language: str | None = None
254    ) -> tuple[str, StreamType, AudioFormat]:
255        """Ask the TTS engine for audio and return the path, stream type and format to play it."""
256        engine = await self._get_tts_engine(engine_uid)
257        try:
258            stream_details = await query_tts_engine(engine, text, language)
259        except TimeoutError, MusicAssistantError:
260            # a timeout or our own structured failure is not a language rejection, so a
261            # language-less retry would not help and would only double the wait
262            raise
263        except Exception as err:
264            if language is None:
265                raise
266            # some engines reject a language they don't support; fall back to the engine's
267            # own default voice rather than losing the clip entirely
268            self.logger.warning(
269                "AI Radio TTS engine '%s' rejected language '%s' (%s), retrying with its "
270                "default voice",
271                engine.uid,
272                language,
273                err,
274            )
275            stream_details = await query_tts_engine(engine, text, None)
276        path, stream_type = await resolve_tts_stream_path(engine, stream_details)
277        audio_format = stream_details.audio_format
278        if audio_format.content_type == ContentType.UNKNOWN:
279            audio_format = AudioFormat(content_type=ContentType.MP3)
280        return path, stream_type, audio_format
281
282    async def _probe_duration(self, path: str) -> int | None:
283        """Return the clip duration in seconds, or None when it cannot be determined."""
284        try:
285            tags = await async_parse_tags(path, require_duration=True)
286        except (InvalidDataError, OSError) as err:
287            if any(marker in str(err) for marker in TTS_SERVER_ERROR_MARKERS):
288                raise MusicAssistantError(
289                    "Error during TTS generation. Does your TTS provider have enough credit? "
290                    "Check the logs of your TTS provider for the reason."
291                ) from err
292            self.logger.warning("Could not determine AI Radio clip duration: %s", err)
293            return None
294        return int(tags.duration) if tags.duration else None
295
296    def _record_skip(self, queue_item: QueueItem, error: str) -> None:
297        """Record a skipped clip on its owning session."""
298        session_id = str(queue_item.extra_attributes.get(ATTR_SESSION_ID) or "")
299        if (session := self._sessions.get(session_id)) is None:
300            return
301        session.skipped_sections += 1
302        session.last_render_error = error
303