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