Add federated music similarity search

This commit is contained in:
Aleksandr Bogomiakov
2026-08-10 00:22:20 +01:00
parent c971e81eed
commit 03f74cc91e
22 changed files with 2998 additions and 22 deletions
+61
View File
@@ -8,6 +8,7 @@ use crate::app::Runtime;
use crate::app::command::{self, Command, Parsed};
use crate::app::event::AppEvent;
use crate::app::state::{AppState, GlobalView, SearchState, Tab};
use crate::library::models::SearchResults;
const SEARCH_DEBOUNCE: Duration = Duration::from_millis(180);
const SEARCH_LIMIT: i64 = 12;
@@ -97,6 +98,7 @@ fn set_view_cursor_zero(state: &mut AppState) {
/// spawned task only queries if it is still the latest after the debounce,
/// and the receiver drops responses that arrive out of date.
pub(super) fn schedule_search(state: &mut AppState, runtime: &Runtime) {
state.search.similarity_source = None;
let seq = runtime.search_seq.fetch_add(1, Ordering::SeqCst) + 1;
let query = state.search.query.clone();
if query.is_empty() {
@@ -145,6 +147,62 @@ pub(super) fn schedule_search(state: &mut AppState, runtime: &Runtime) {
}
}
pub(super) fn schedule_similarity_search(
state: &mut AppState,
runtime: &Runtime,
track: &crate::library::models::TrackItem,
) {
let seq = runtime.search_seq.fetch_add(1, Ordering::SeqCst) + 1;
let artist = track.artist_line();
state.search.query = if artist.is_empty() {
track.title.clone()
} else {
format!("{}{artist}", track.title)
};
state.search.similarity_source = Some(track.id);
state.search.loading = true;
state.search.results = None;
state.search.fed_tracks.clear();
state.search.fed_artists.clear();
state.search.fed_loading = false;
state.active_tab = crate::app::state::Tab::Global;
if let Some(crate::app::state::GlobalView::Search { cursor }) = state.global.stack.last_mut() {
*cursor = 0;
} else {
state
.global
.stack
.push(crate::app::state::GlobalView::Search { cursor: 0 });
}
let similarity = Arc::clone(&runtime.similarity);
let tx = runtime.event_tx.clone();
let track_id = track.id;
let source_track = track.clone();
tokio::task::spawn_blocking(move || {
let result = similarity
.search_track(track_id, 49)
.map(|(matches, query)| {
let mut tracks = Vec::with_capacity(1 + matches.len());
tracks.push(source_track);
tracks.extend(matches.into_iter().map(|found| found.track));
(
SearchResults {
artists: Vec::new(),
releases: Vec::new(),
tracks,
},
query,
)
});
let (result, query) = match result {
Ok((results, query)) => (Ok(results), Some(query)),
Err(err) => (Err(format!("{err:#}")), None),
};
let _ = tx.send(AppEvent::SimilaritySearchLoaded { seq, result, query });
});
}
/// Refresh only the local-library half of an already open search.
///
/// Library/device sync notifications can arrive while federated search
@@ -152,6 +210,9 @@ pub(super) fn schedule_search(state: &mut AppState, runtime: &Runtime) {
/// federation rows and bump the shared sequence, causing valid network
/// responses to be dropped or flicker away.
pub(super) fn refresh_local_search(state: &mut AppState, runtime: &Runtime) {
if state.search.similarity_source.is_some() {
return;
}
let query = state.search.query.clone();
if query.is_empty() {
return;
+10
View File
@@ -32,6 +32,16 @@ pub enum AppEvent {
seq: u64,
result: Result<SearchResults, String>,
},
/// Local similar-track search completed. It uses the same sequence as
/// text search so stale pages cannot overwrite a newer request.
SimilaritySearchLoaded {
seq: u64,
result: Result<SearchResults, String>,
query: Option<crate::similarity::QueryVector>,
},
SimilarityStatus(crate::similarity::SimilarityStatus),
/// `None` is emitted after clearing every stored embedding.
SimilarityProfileActivated(Option<String>),
/// Artwork loaded and decoded for the shared art cache.
ArtLoaded {
key: String,
+60
View File
@@ -41,6 +41,7 @@ pub struct Runtime {
pub devices: Arc<crate::devices::DeviceSync>,
pub jam: Arc<crate::jam::JamManager>,
pub federation: Arc<crate::federation::Federation>,
pub similarity: Arc<crate::similarity::Manager>,
/// When the last Federation-tab status snapshot was requested.
pub fed_status_at: Option<std::time::Instant>,
pub library_network_refresh_at: Option<std::time::Instant>,
@@ -262,6 +263,7 @@ pub async fn run(
state.player.volume = settings.volume;
state.global.filters = settings.library;
state.music_dir = settings.music_dir.clone();
state.similarity.settings = settings.similarity.clone();
if let Err(err) = state.visualizer.load_library() {
state.status_message = Some(format!("visualizations disabled: {err:#}"));
}
@@ -269,10 +271,17 @@ pub async fn run(
let devices = crate::devices::DeviceSync::new(Arc::clone(&library))?;
devices.set_event_tx(event_tx.clone());
let jam = crate::jam::JamManager::new(event_tx.clone());
let similarity = crate::similarity::Manager::new(
Arc::clone(&library),
event_tx.clone(),
settings.similarity.clone(),
);
state.similarity.status = similarity.status();
let federation = crate::federation::Federation::new(
Arc::clone(&library),
Arc::clone(&devices),
Arc::clone(&jam),
Arc::clone(&similarity),
settings.music_dir.clone(),
);
state.music_dir = federation.media_dir();
@@ -292,6 +301,7 @@ pub async fn run(
devices,
jam,
federation,
similarity,
fed_status_at: None,
library_network_refresh_at: None,
library_network_refreshing: Arc::new(std::sync::atomic::AtomicBool::new(false)),
@@ -319,6 +329,9 @@ pub async fn run(
status_publisher: crate::status::Publisher::spawn(),
};
spawn_content_id_backfill(&runtime);
if state.similarity.settings.enabled {
runtime.similarity.start();
}
{
let fed = Arc::clone(&runtime.federation);
@@ -1581,6 +1594,15 @@ fn perform_effect(state: &mut AppState, runtime: &mut Runtime, effect: Effect) {
let _ = tx.send(event);
});
}
Effect::SimilarityApplySettings => {
save_app_settings(state);
runtime.similarity.apply(state.similarity.settings.clone());
state.similarity.status = runtime.similarity.status();
}
Effect::SimilarityClear => {
state.status_message = Some("clearing stored embeddings…".to_string());
runtime.similarity.clear();
}
Effect::FedApplySettings => fed_apply_settings(state, runtime),
Effect::FedSyncNow => {
state.federation.publishing = true;
@@ -2785,6 +2807,7 @@ fn save_app_settings(state: &AppState) {
volume: state.player.volume,
library: state.global.filters,
music_dir: state.music_dir.clone(),
similarity: state.similarity.settings.clone(),
};
if let Err(err) = crate::config::settings::save(&settings) {
tracing::warn!(%err, "saving app settings failed");
@@ -3632,6 +3655,40 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent
Err(message) => state.status_message = Some(message),
}
}
AppEvent::SimilaritySearchLoaded { seq, result, query } => {
if seq != runtime.search_seq.load(std::sync::atomic::Ordering::SeqCst) {
return;
}
state.search.loading = false;
match result {
Ok(results) => state.search.results = Some(results),
Err(message) => {
state.status_message = Some(format!("similarity search failed: {message}"));
return;
}
}
if let Some(query) = query
&& state.federation.settings.enabled
&& runtime.similarity.network_allowed()
{
state.search.fed_loading = true;
let federation = Arc::clone(&runtime.federation);
let tx = runtime.event_tx.clone();
tokio::spawn(async move {
let result = federation
.search_similar(query, 50)
.await
.map_err(|err| format!("{err:#}"));
let _ = tx.send(AppEvent::FedSearchLoaded { seq, result });
});
}
}
AppEvent::SimilarityStatus(status) => state.similarity.status = status,
AppEvent::SimilarityProfileActivated(profile_id) => {
state.similarity.settings.active_profile = profile_id;
state.similarity.status = runtime.similarity.status();
save_app_settings(state);
}
AppEvent::ArtLoaded { key, art } => {
let entry = match art {
Some(image) => state::ArtState::Ready(image),
@@ -3857,6 +3914,9 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent
}
AppEvent::LibraryChanged { message } => {
on_library_changed(state, runtime);
if state.similarity.settings.enabled {
runtime.similarity.start();
}
if let Some(message) = message {
state.status_message = Some(message);
}
+70
View File
@@ -115,6 +115,39 @@ pub fn handle_key(state: &mut AppState, runtime: &mut Runtime, key: KeyEvent) {
Popup::ConfirmDelete { target, label } => {
handle_confirm_delete(state, runtime, target, label, key);
}
Popup::SimilarityPrivacyConsent { enable_federation } => match key.code {
KeyCode::Enter | KeyCode::Char('y') => {
state.similarity.settings.federation_consent = true;
state.similarity.settings.enabled = true;
super::perform_effect(
state,
runtime,
crate::app::update::Effect::SimilarityApplySettings,
);
if enable_federation {
super::perform_effect(
state,
runtime,
crate::app::update::Effect::FedApplySettings,
);
}
}
KeyCode::Esc | KeyCode::Char('n') | KeyCode::Char('q') => {
if enable_federation {
state.federation.settings.enabled = false;
}
}
_ => {
state.popup = Some(Popup::SimilarityPrivacyConsent { enable_federation });
}
},
Popup::ConfirmClearEmbeddings => match key.code {
KeyCode::Enter | KeyCode::Char('y') => {
super::perform_effect(state, runtime, crate::app::update::Effect::SimilarityClear)
}
KeyCode::Esc | KeyCode::Char('n') | KeyCode::Char('q') => {}
_ => state.popup = Some(Popup::ConfirmClearEmbeddings),
},
Popup::LibraryFilters { cursor } => handle_library_filters(state, runtime, cursor, key),
Popup::TrackInfo {
tracks,
@@ -544,6 +577,21 @@ fn handle_fed_input(
super::validate_music_directory(state, runtime, value.into());
}
}
FedInputField::SimilarityWorkers => match value.parse::<usize>() {
Ok(workers @ 1..=16) => {
state.similarity.settings.workers = workers;
super::perform_effect(
state,
runtime,
crate::app::update::Effect::SimilarityApplySettings,
);
}
_ => {
state.status_message =
Some("similarity workers must be a number from 1 to 16".into());
state.popup = Some(Popup::FedInput { field, input });
}
},
FedInputField::ConnectTicket => {
if value.is_empty() {
state.status_message = Some("ticket is empty".into());
@@ -1019,6 +1067,28 @@ fn handle_track_info(
scroll,
});
}
KeyCode::Char('s') => {
let Some(track) = tracks.get(cursor.min(len.saturating_sub(1))) else {
return;
};
if !state.similarity.settings.enabled {
state.status_message = Some("enable Similarity search in Settings first".into());
state.popup = Some(Popup::TrackInfo {
tracks,
cursor,
scroll,
});
} else if track.id < 0 || track.file_path.is_empty() {
state.status_message = Some("similarity search starts from a local track".into());
state.popup = Some(Popup::TrackInfo {
tracks,
cursor,
scroll,
});
} else {
super::cmdline::schedule_similarity_search(state, runtime, track);
}
}
_ => {
state.popup = Some(Popup::TrackInfo {
tracks,
+40
View File
@@ -687,6 +687,10 @@ pub enum Popup {
},
/// Delete confirmation; Enter/y deletes, Esc/n cancels.
ConfirmDelete { target: DeleteTarget, label: String },
/// Enabling network similarity reveals an embedding query to peers.
SimilarityPrivacyConsent { enable_federation: bool },
/// Derived data is safe to recreate but potentially expensive.
ConfirmClearEmbeddings,
/// Library-home filters. Cursor is kept for the next filters added here.
LibraryFilters { cursor: usize },
/// Track metadata viewer; left/right switch between selected tracks.
@@ -812,6 +816,7 @@ impl StatusDetailFocus {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FedInputField {
MusicDirectory,
SimilarityWorkers,
NetworkId,
ConnectTicket,
DeviceName,
@@ -823,6 +828,7 @@ impl FedInputField {
pub fn title(self) -> &'static str {
match self {
FedInputField::MusicDirectory => "Music save directory",
FedInputField::SimilarityWorkers => "Similarity background workers",
FedInputField::NetworkId => "Network ID",
FedInputField::ConnectTicket => "Connect to peer (paste ticket)",
FedInputField::DeviceName => "Device name",
@@ -836,6 +842,9 @@ impl FedInputField {
FedInputField::MusicDirectory => {
"Federated tracks saved to your library use this directory. The directory is checked for write access before anything changes."
}
FedInputField::SimilarityWorkers => {
"Enter the maximum number of tracks processed in parallel, from 1 to 16. The change takes effect immediately."
}
FedInputField::NetworkId => {
"A unique network id. It must match exactly on every client that should see and connect to the same peers."
}
@@ -866,6 +875,25 @@ pub enum FedRow {
Connect,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimilarityRow {
Toggle,
Model,
Profile,
Workers,
Clear,
}
impl SimilarityRow {
pub const ALL: [SimilarityRow; 5] = [
SimilarityRow::Toggle,
SimilarityRow::Model,
SimilarityRow::Profile,
SimilarityRow::Workers,
SimilarityRow::Clear,
];
}
impl FedRow {
pub const ALL: [FedRow; 6] = [
FedRow::Toggle,
@@ -883,6 +911,7 @@ impl FedRow {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SettingsRow {
MusicDirectory,
Similarity(SimilarityRow),
Federation(FedRow),
StatusDetails,
DeviceName,
@@ -1021,6 +1050,7 @@ pub fn device_status_order(state: &AppState) -> Vec<usize> {
pub fn settings_rows(state: &AppState) -> Vec<SettingsRow> {
let mut rows = vec![SettingsRow::MusicDirectory];
rows.extend(SimilarityRow::ALL.into_iter().map(SettingsRow::Similarity));
rows.extend(FedRow::ALL.into_iter().map(SettingsRow::Federation));
rows.push(SettingsRow::DeviceName);
rows.push(SettingsRow::DeviceInvite);
@@ -1059,6 +1089,12 @@ pub struct FederationTab {
pub device_syncing: bool,
}
#[derive(Debug, Default)]
pub struct SimilarityTab {
pub settings: crate::config::settings::SimilaritySettings,
pub status: crate::similarity::SimilarityStatus,
}
/// Playlists eligible as add-targets (the virtual Likes playlist is managed
/// through likes, not direct adds).
pub fn addable_playlists(state: &AppState) -> Vec<(i64, String)> {
@@ -1187,6 +1223,9 @@ pub struct SearchState {
/// from the artist names of matching tracks.
pub fed_artists: Vec<crate::federation::FedArtistHit>,
pub fed_loading: bool,
/// Present only for a track-seeded search; text-search refreshes must not
/// replace this page with a title query.
pub similarity_source: Option<i64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -1423,6 +1462,7 @@ pub struct AppState {
pub logs: LogsTab,
pub queue_tab: QueueTab,
pub federation: FederationTab,
pub similarity: SimilarityTab,
/// The one federated artist card being viewed (name + loading state);
/// opening another card replaces it.
pub fed_artist_view: Option<(String, Loadable<crate::federation::FedArtistCard>)>,
+61 -2
View File
@@ -63,6 +63,9 @@ pub enum Effect {
},
/// Persist the federation settings and start/stop the node.
FedApplySettings,
/// Persist/apply embedding model, profile, worker or enable changes.
SimilarityApplySettings,
SimilarityClear,
/// Force an immediate library publish into the DHT.
FedSyncNow,
/// Fetch this peer's ticket and show it in a popup.
@@ -2720,7 +2723,7 @@ fn fed_card_featured_artist_names(track: &crate::federation::FedCardTrack) -> Ve
/// Enter on Settings: toggle switches, open text inputs, run
/// one-shot operations. The heavy lifting happens in perform_effect().
fn federation_select(state: &mut AppState) -> Option<Effect> {
use super::state::{FedInputField, FedRow, Popup, SettingsRow};
use super::state::{FedInputField, FedRow, Popup, SettingsRow, SimilarityRow};
match settings_rows(state).get(state.settings_cursor).copied()? {
SettingsRow::MusicDirectory => {
if state.music_dir_changing {
@@ -2735,6 +2738,52 @@ fn federation_select(state: &mut AppState) -> Option<Effect> {
});
None
}
SettingsRow::Similarity(SimilarityRow::Toggle) => {
if !state.similarity.settings.enabled
&& state.federation.settings.enabled
&& !state.similarity.settings.federation_consent
{
state.popup = Some(Popup::SimilarityPrivacyConsent {
enable_federation: false,
});
return None;
}
state.similarity.settings.enabled = !state.similarity.settings.enabled;
Some(Effect::SimilarityApplySettings)
}
SettingsRow::Similarity(SimilarityRow::Model) => {
let models = crate::similarity::MODELS;
let current = models
.iter()
.position(|model| model.id == state.similarity.settings.model)
.unwrap_or(0);
state.similarity.settings.model = models[(current + 1) % models.len()].id.to_string();
Some(Effect::SimilarityApplySettings)
}
SettingsRow::Similarity(SimilarityRow::Profile) => {
let profile = state.similarity.settings.profile.clone();
let text =
crate::similarity::profile_details(&profile, &state.similarity.settings.model)
.unwrap_or_else(|| format!("Unknown preprocessing profile: {profile}"));
state.popup = Some(Popup::FedText {
title: format!("Preprocessing profile: {profile}"),
text,
});
None
}
SettingsRow::Similarity(SimilarityRow::Workers) => {
state.popup = Some(Popup::FedInput {
field: FedInputField::SimilarityWorkers,
input: crate::app::input::LineEdit::new(
state.similarity.settings.workers.to_string(),
),
});
None
}
SettingsRow::Similarity(SimilarityRow::Clear) => {
state.popup = Some(Popup::ConfirmClearEmbeddings);
None
}
SettingsRow::Federation(FedRow::Toggle) => {
let settings = &mut state.federation.settings;
if !settings.enabled && settings.network_id.trim().is_empty() {
@@ -2744,7 +2793,17 @@ fn federation_select(state: &mut AppState) -> Option<Effect> {
});
return None;
}
settings.enabled = !settings.enabled;
let enabling = !settings.enabled;
settings.enabled = enabling;
if enabling
&& state.similarity.settings.enabled
&& !state.similarity.settings.federation_consent
{
state.popup = Some(Popup::SimilarityPrivacyConsent {
enable_federation: true,
});
return None;
}
Some(Effect::FedApplySettings)
}
SettingsRow::Federation(FedRow::NetworkId) => {