/
/
/
1"""Pure-function helpers for the Sonic Similarity plugin."""
2
3from __future__ import annotations
4
5import re
6from typing import Any
7
8import numpy as np
9
10from music_assistant.providers.sonic_similarity.clap_index import CLAP_EMBEDDING_DIM
11from music_assistant.providers.sonic_similarity.constants import SIMILARITY_PRESETS
12from music_assistant.providers.sonic_similarity.models import SimilarParams
13from music_assistant.providers.sonic_similarity.similarity import ScoredCandidate
14
15_MUSIC_WORD = re.compile(r"\bmusic\b", re.IGNORECASE)
16
17
18def format_text_query(query: str) -> str:
19 """
20 Frame a bare query as ``<query> music`` for the CLAP text encoder.
21
22 Skipped when the query already contains the word "music" (case-insensitive);
23 an empty or whitespace-only query returns "".
24
25 :param query: Raw user query.
26 """
27 # CLAP was trained on audio captions, so the " music" suffix matches that form.
28 cleaned = query.strip()
29 if not cleaned or _MUSIC_WORD.search(cleaned):
30 return cleaned
31 return f"{cleaned} music"
32
33
34def _parse_clap_embedding(raw: Any) -> np.ndarray | None:
35 """Coerce a stored embedding (list/tuple) into a 1024-dim float32 array, or None."""
36 if raw is None:
37 return None
38 try:
39 arr = np.asarray(raw, dtype=np.float32).reshape(-1)
40 except TypeError, ValueError:
41 return None
42 if arr.shape != (CLAP_EMBEDDING_DIM,):
43 return None
44 return arr
45
46
47def _parse_weights(params: dict[str, Any]) -> dict[str, float]:
48 """Parse similarity weights from API parameters."""
49 preset_name = str(params.get("preset", "balanced"))
50 preset = SIMILARITY_PRESETS.get(preset_name, SIMILARITY_PRESETS["balanced"])
51 result = dict(preset)
52
53 def _clamp(val: str, fallback: float) -> float:
54 try:
55 return max(0.0, min(1.0, float(val)))
56 except ValueError, TypeError:
57 return fallback
58
59 for group, default in result.items():
60 key = f"{group}_weight"
61 if key in params:
62 result[group] = _clamp(params[key], default)
63
64 return result
65
66
67def _parse_similar_params( # noqa: PLR0913
68 item_id: str | None = None,
69 item_ids: list[str] | None = None,
70 limit: int = 25,
71 depth: int = 1,
72 branch_factor: int = 5,
73 blend_mode: str = "centroid",
74 seed_weights: list[float] | None = None,
75 diversity: float = 0.0,
76 preset: str = "balanced",
77 candidates: int = 200,
78 filter_genres: list[str] | None = None,
79 filter_providers: list[str] | None = None,
80 exclude_track_ids: list[str] | None = None,
81 exclude_artists: list[str] | None = None,
82 resolve: bool = False,
83 include_group_distances: bool = False,
84 seed_provider: str | None = None,
85 **kwargs: Any,
86) -> SimilarParams:
87 """Validate and normalize parameters for the similar endpoint."""
88 if item_ids is None:
89 if item_id is None:
90 msg = "Either item_id or item_ids must be provided"
91 raise ValueError(msg)
92 item_ids = [item_id]
93
94 limit = max(1, min(100, limit))
95 depth = max(1, min(5, depth))
96 diversity = max(0.0, min(1.0, diversity))
97
98 if blend_mode not in ("centroid", "union"):
99 blend_mode = "centroid"
100
101 if seed_weights is not None and len(seed_weights) != len(item_ids):
102 msg = f"seed_weights length ({len(seed_weights)}) must match item_ids ({len(item_ids)})"
103 raise ValueError(msg)
104
105 has_filters = any(
106 x is not None for x in (filter_genres, filter_providers, exclude_track_ids, exclude_artists)
107 )
108 if has_filters:
109 candidates = candidates * 2
110
111 return SimilarParams(
112 item_ids=item_ids,
113 limit=limit,
114 depth=depth,
115 branch_factor=branch_factor,
116 blend_mode=blend_mode,
117 seed_weights=seed_weights,
118 diversity=diversity,
119 preset=preset,
120 candidates=candidates,
121 filter_genres=filter_genres,
122 filter_providers=filter_providers,
123 exclude_track_ids=exclude_track_ids,
124 exclude_artists=exclude_artists,
125 resolve=resolve,
126 include_group_distances=include_group_distances,
127 seed_provider=seed_provider,
128 weight_overrides=kwargs,
129 )
130
131
132def apply_filters(
133 candidates: list[ScoredCandidate],
134 seed_ids: set[str],
135 exclude_track_ids: set[str] | None,
136 filter_providers: set[str] | None,
137) -> list[ScoredCandidate]:
138 """
139 Apply cheap post-ANN filters to candidate list.
140
141 :param candidates: ScoredCandidate results from the ANN search.
142 :param seed_ids: Seed track IDs to exclude.
143 :param exclude_track_ids: Additional track IDs to exclude.
144 :param filter_providers: If set, only keep candidates from these providers.
145 """
146 result: list[ScoredCandidate] = []
147 exclude = seed_ids | (exclude_track_ids or set())
148
149 for cand in candidates:
150 if cand.item_id in exclude:
151 continue
152 if filter_providers is not None and cand.provider not in filter_providers:
153 continue
154 result.append(cand)
155
156 return result
157