/
/
/
1"""Helpers to talk to an OpenAI-compatible API."""
2
3from __future__ import annotations
4
5import re
6from typing import TYPE_CHECKING, Any
7
8import aiohttp
9from aiohttp import ClientTimeout
10from music_assistant_models.errors import (
11 InvalidDataError,
12 LoginFailed,
13 ProviderUnavailableError,
14 RateLimited,
15)
16
17from music_assistant.helpers.json import JSON_DECODE_EXCEPTIONS, json_loads
18
19if TYPE_CHECKING:
20 from music_assistant.mass import MusicAssistant
21
22# a reasoning model may narrate its thinking inline ahead of the answer, while the
23# consumers want the answer alone - some of them parse it as strict JSON
24THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
25
26# an upstream failure message is echoed into the raised error, so keep it bounded
27MAX_ERROR_DETAIL_CHARS = 200
28
29
30async def list_models(mass: MusicAssistant, base_url: str, api_key: str, timeout: int) -> list[str]:
31 """
32 Return the model ids the endpoint advertises, sorted.
33
34 Raises when the endpoint cannot be asked at all, so that a wrong address is caught
35 rather than mistaken for a service that simply has nothing to offer. An endpoint
36 that answers without a usable list yields an empty list.
37
38 :param mass: The Music Assistant instance, for its shared HTTP session.
39 :param base_url: Base URL of the API, without a trailing slash.
40 :param api_key: API key to authenticate with, empty for an endpoint without auth.
41 :param timeout: Maximum request duration in seconds.
42 """
43 payload = await _request(mass, "get", f"{base_url}/models", api_key, timeout=timeout)
44 data = payload.get("data")
45 if not isinstance(data, list):
46 return []
47 models = {
48 model["id"].strip()
49 for model in data
50 if isinstance(model, dict) and isinstance(model.get("id"), str) and model["id"].strip()
51 }
52 return sorted(models)
53
54
55async def chat_completion(
56 mass: MusicAssistant,
57 base_url: str,
58 api_key: str,
59 model: str,
60 prompt: str,
61 timeout: int,
62) -> str:
63 """
64 Return the reply of a single-turn chat completion.
65
66 :param mass: The Music Assistant instance, for its shared HTTP session.
67 :param base_url: Base URL of the API, without a trailing slash.
68 :param api_key: API key to authenticate with, empty for an endpoint without auth.
69 :param model: Id of the model to answer the prompt.
70 :param prompt: The prompt to send.
71 :param timeout: Maximum request duration in seconds.
72 """
73 payload = await _request(
74 mass,
75 "post",
76 f"{base_url}/chat/completions",
77 api_key,
78 timeout=timeout,
79 # deliberately no token limit: reasoning models reject max_tokens outright and
80 # a truncated answer is worthless to every consumer anyway
81 json={"model": model, "messages": [{"role": "user", "content": prompt}]},
82 )
83 choices = payload.get("choices")
84 if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict):
85 msg = f"Model {model} returned no choices"
86 raise InvalidDataError(msg)
87 message = choices[0].get("message")
88 content = message.get("content") if isinstance(message, dict) else None
89 if not isinstance(content, str):
90 msg = f"Model {model} returned no message content"
91 raise InvalidDataError(msg)
92 if not (answer := THINK_BLOCK.sub("", content).strip()):
93 msg = f"Model {model} returned an empty answer"
94 raise InvalidDataError(msg)
95 return answer
96
97
98async def _request(
99 mass: MusicAssistant,
100 method: str,
101 url: str,
102 api_key: str,
103 *,
104 timeout: int,
105 json: dict[str, Any] | None = None,
106) -> dict[str, Any]:
107 """Perform one API request and return its parsed JSON object."""
108 headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
109 try:
110 async with mass.http_session.request(
111 method, url, headers=headers, json=json, timeout=ClientTimeout(total=timeout)
112 ) as response:
113 status = response.status
114 retry_after = response.headers.get("Retry-After", "")
115 body = await response.read()
116 except TimeoutError as err:
117 msg = f"Request to {url} timed out after {timeout} seconds"
118 raise ProviderUnavailableError(msg) from err
119 except aiohttp.ClientError as err:
120 msg = f"Could not reach {url}: {err}"
121 raise ProviderUnavailableError(msg) from err
122 if status in (401, 403):
123 msg = f"The endpoint rejected the API key: {_error_detail(body)}"
124 raise LoginFailed(msg)
125 if status == 429:
126 msg = f"The endpoint is rate limiting us: {_error_detail(body)}"
127 raise RateLimited(msg, backoff_time=int(retry_after) if retry_after.isdigit() else 0)
128 if status >= 400:
129 msg = f"The endpoint returned an error ({status}): {_error_detail(body)}"
130 raise InvalidDataError(msg)
131 try:
132 payload = json_loads(body)
133 except JSON_DECODE_EXCEPTIONS as err:
134 msg = f"The endpoint returned a malformed response: {_error_detail(body)}"
135 raise InvalidDataError(msg) from err
136 if not isinstance(payload, dict):
137 msg = "The endpoint returned an unexpected response"
138 raise InvalidDataError(msg)
139 return payload
140
141
142def _error_detail(body: bytes) -> str:
143 """Return a short, readable description of a failed request's response body."""
144 text = body.decode(errors="ignore").strip()
145 try:
146 payload = json_loads(body)
147 except JSON_DECODE_EXCEPTIONS:
148 payload = None
149 if isinstance(payload, dict) and isinstance(error := payload.get("error"), dict):
150 if isinstance(message := error.get("message"), str) and message.strip():
151 text = message.strip()
152 return text[:MAX_ERROR_DETAIL_CHARS] or "no details provided"
153