music-assistant-server

1"""Helpers to discover, select and configure the AI/TTS engines exposed by plugins."""
2
3from __future__ import annotations
4
5import logging
6from typing import TYPE_CHECKING, Any
7
8from music_assistant_models.config_entries import ConfigEntry, ConfigValueOption
9from music_assistant_models.enums import ConfigEntryType, ProviderFeature, ProviderType
10
11from music_assistant.constants import MASS_LOGGER_NAME
12from music_assistant.models.plugin import (
13    AIEngine,
14    PluginEngine,
15    PluginProvider,
16    TTSEngine,
17)
18
19if TYPE_CHECKING:
20    from collections.abc import Callable, Coroutine
21
22    from music_assistant.mass import MusicAssistant
23    from music_assistant.models.provider import Provider
24
25LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.helpers.plugin_engines")
26
27
28async def get_ai_engines(mass: MusicAssistant) -> list[AIEngine]:
29    """
30    Return all AI engines currently exposed by the loaded plugins, in a stable order.
31
32    :param mass: The Music Assistant instance to query.
33    """
34    return await _collect_engines(
35        mass, ProviderFeature.AI_QUERY, lambda provider: provider.get_ai_engines()
36    )
37
38
39async def get_tts_engines(mass: MusicAssistant) -> list[TTSEngine]:
40    """
41    Return all TTS engines currently exposed by the loaded plugins, in a stable order.
42
43    :param mass: The Music Assistant instance to query.
44    """
45    return await _collect_engines(
46        mass, ProviderFeature.TTS, lambda provider: provider.get_tts_engines()
47    )
48
49
50async def resolve_ai_engine(mass: MusicAssistant, selected: str | None) -> AIEngine | None:
51    """
52    Return the AI engine for a configured selection.
53
54    :param mass: The Music Assistant instance to query.
55    :param selected: The configured engine uid. Returns None when it is empty/unset or names
56        an engine that no longer exists - another engine is never substituted for it.
57    """
58    return _resolve(await get_ai_engines(mass), selected)
59
60
61async def resolve_tts_engine(mass: MusicAssistant, selected: str | None) -> TTSEngine | None:
62    """
63    Return the TTS engine for a configured selection.
64
65    :param mass: The Music Assistant instance to query.
66    :param selected: The configured engine uid. Returns None when it is empty/unset or names
67        an engine that no longer exists - another engine is never substituted for it.
68    """
69    return _resolve(await get_tts_engines(mass), selected)
70
71
72async def select_ai_engine(
73    provider: Provider, key: str, *, in_setup_data: bool = False
74) -> AIEngine | None:
75    """
76    Return the provider's AI engine, adopting the first available one when it has none yet.
77
78    A selection that was made explicitly is honoured or, when that engine no longer exists,
79    reported as missing - another engine is never substituted for it.
80
81    :param provider: The provider whose selection to read and, if unset, persist.
82    :param key: The config key holding the selected engine uid.
83    :param in_setup_data: Store in (and read from) the provider's setup_data instead of its
84        regular config values.
85    """
86    return await _select_engine(
87        provider, key, await get_ai_engines(provider.mass), in_setup_data=in_setup_data
88    )
89
90
91async def select_tts_engine(
92    provider: Provider, key: str, *, in_setup_data: bool = False
93) -> TTSEngine | None:
94    """
95    Return the provider's TTS engine, adopting the first available one when it has none yet.
96
97    A selection that was made explicitly is honoured or, when that engine no longer exists,
98    reported as missing - another engine is never substituted for it.
99
100    :param provider: The provider whose selection to read and, if unset, persist.
101    :param key: The config key holding the selected engine uid.
102    :param in_setup_data: Store in (and read from) the provider's setup_data instead of its
103        regular config values.
104    """
105    return await _select_engine(
106        provider, key, await get_tts_engines(provider.mass), in_setup_data=in_setup_data
107    )
108
109
110async def create_ai_engine_config_entries(
111    mass: MusicAssistant, key: str, depends_on: str | None = None, required: bool = False
112) -> tuple[ConfigEntry, ...]:
113    """
114    Return the config entries letting the user pick an AI engine.
115
116    Adds a second (alert) entry with key ``<key>_unavailable`` when no engine is available.
117
118    :param mass: The Music Assistant instance to query.
119    :param key: The config entry key holding the selected engine uid.
120    :param depends_on: Optional key of the entry this picker should be shown for.
121    :param required: Reject a submission that leaves the picker empty. Only for setup flows;
122        an options entry must stay optional so a provider with no selection still loads.
123    """
124    return _create_engine_config_entries(await get_ai_engines(mass), key, depends_on, required)
125
126
127async def create_tts_engine_config_entries(
128    mass: MusicAssistant, key: str, depends_on: str | None = None, required: bool = False
129) -> tuple[ConfigEntry, ...]:
130    """
131    Return the config entries letting the user pick a TTS engine.
132
133    Adds a second (alert) entry with key ``<key>_unavailable`` when no engine is available.
134
135    :param mass: The Music Assistant instance to query.
136    :param key: The config entry key holding the selected engine uid.
137    :param depends_on: Optional key of the entry this picker should be shown for.
138    :param required: Reject a submission that leaves the picker empty. Only for setup flows;
139        an options entry must stay optional so a provider with no selection still loads.
140    """
141    return _create_engine_config_entries(await get_tts_engines(mass), key, depends_on, required)
142
143
144async def _collect_engines[EngineT: PluginEngine](
145    mass: MusicAssistant,
146    feature: ProviderFeature,
147    fetch: Callable[[PluginProvider], Coroutine[Any, Any, list[EngineT]]],
148) -> list[EngineT]:
149    """Collect the engines of every available plugin declaring the given feature."""
150    result: list[EngineT] = []
151    for provider in mass.get_providers_supporting_feature(feature, priority=(ProviderType.PLUGIN,)):
152        if not isinstance(provider, PluginProvider):
153            continue
154        try:
155            engines = await fetch(provider)
156        except Exception as err:
157            LOGGER.warning(
158                "Could not retrieve %s engines from %s: %s",
159                feature,
160                provider.instance_id,
161                err,
162            )
163            continue
164        result.extend(sorted(engines, key=lambda engine: engine.name))
165    return result
166
167
168def _resolve[EngineT: PluginEngine](engines: list[EngineT], selected: str | None) -> EngineT | None:
169    """Return the available engine matching a configured selection, if any."""
170    if not selected:
171        return None
172    return next((engine for engine in engines if engine.uid == selected), None)
173
174
175async def _select_engine[EngineT: PluginEngine](
176    provider: Provider,
177    key: str,
178    engines: list[EngineT],
179    *,
180    in_setup_data: bool,
181) -> EngineT | None:
182    """Return the stored engine, or persist and return the first available one."""
183    config = provider.mass.config
184    stored = (
185        config.get_provider_setup_value(provider.instance_id, key)
186        if in_setup_data
187        else config.get_raw_provider_config_value(provider.instance_id, key)
188    )
189    if isinstance(stored, str) and stored:
190        return _resolve(engines, stored)
191    if not engines:
192        return None
193    engine = engines[0]
194    if in_setup_data:
195        provider._update_setup_data(key, engine.uid)
196    else:
197        provider._update_config_value(key, engine.uid)
198    LOGGER.debug("Selected engine %s as %s for %s", engine.uid, key, provider.instance_id)
199    return engine
200
201
202def _create_engine_config_entries(
203    engines: list[AIEngine] | list[TTSEngine],
204    key: str,
205    depends_on: str | None,
206    required: bool = False,
207) -> tuple[ConfigEntry, ...]:
208    """Build the picker (and unavailable alert) config entries for the given engines."""
209    entry = ConfigEntry(
210        key=key,
211        type=ConfigEntryType.STRING,
212        required=required,
213        # a concrete default would make the seeded selection equal to it and therefore
214        # not persisted (to_raw only stores values differing from the default)
215        default_value=None,
216        # the picker aggregates engines from every plugin, so each option names the
217        # plugin it came from - engine names alone are ambiguous across providers
218        options=[
219            ConfigValueOption(engine.uid, title=f"{engine.provider.name} | {engine.name}")
220            for engine in engines
221        ],
222        depends_on=depends_on,
223        category="features",
224        read_only=not engines,
225    )
226    if engines:
227        return (entry,)
228    return (
229        entry,
230        ConfigEntry(
231            key=f"{key}_unavailable",
232            type=ConfigEntryType.ALERT,
233            depends_on=depends_on,
234        ),
235    )
236