From 079d87a831435e38bb2439bdd613da702c4caefd Mon Sep 17 00:00:00 2001 From: Aleksandr Bogomiakov Date: Mon, 10 Aug 2026 20:51:21 +0100 Subject: [PATCH] Merge local and federated similarity results --- src/app/cmdline.rs | 24 ++--- src/app/event.rs | 7 +- src/app/mod.rs | 179 ++++++++++++++++++++++++++++++++++- src/app/popup.rs | 30 ++++++ src/app/state.rs | 150 ++++++++++++++++++++++++++++- src/app/update.rs | 142 +++++++++++++++++++++------ src/config/settings.rs | 43 ++++++++- src/federation/mod.rs | 25 ++++- src/federation/similarity.rs | 56 +++++++---- src/similarity.rs | 52 +++++++++- src/ui/federation.rs | 8 ++ src/ui/global.rs | 153 ++++++++++++++++++++++++++++++ 12 files changed, 794 insertions(+), 75 deletions(-) diff --git a/src/app/cmdline.rs b/src/app/cmdline.rs index 15286a4..c14c1c7 100644 --- a/src/app/cmdline.rs +++ b/src/app/cmdline.rs @@ -8,7 +8,6 @@ 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; @@ -99,6 +98,10 @@ fn set_view_cursor_zero(state: &mut AppState) { /// 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; + state.search.similarity_source_track = None; + state.search.similarity_tracks.clear(); + state.search.similarity_stats = None; + state.search.similarity_error = None; let seq = runtime.search_seq.fetch_add(1, Ordering::SeqCst) + 1; let query = state.search.query.clone(); if query.is_empty() { @@ -160,6 +163,10 @@ pub(super) fn schedule_similarity_search( format!("{} — {artist}", track.title) }; state.search.similarity_source = Some(track.id); + state.search.similarity_source_track = Some(track.clone()); + state.search.similarity_tracks.clear(); + state.search.similarity_stats = None; + state.search.similarity_error = None; state.search.loading = true; state.search.results = None; state.search.fed_tracks.clear(); @@ -178,23 +185,10 @@ pub(super) fn schedule_similarity_search( 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, - ) - }); + .map(|(matches, query)| (matches, query)); let (result, query) = match result { Ok((results, query)) => (Ok(results), Some(query)), Err(err) => (Err(format!("{err:#}")), None), diff --git a/src/app/event.rs b/src/app/event.rs index 3e1cbf1..5025389 100644 --- a/src/app/event.rs +++ b/src/app/event.rs @@ -36,9 +36,14 @@ pub enum AppEvent { /// text search so stale pages cannot overwrite a newer request. SimilaritySearchLoaded { seq: u64, - result: Result, + result: Result, String>, query: Option, }, + /// Federated candidates and diagnostics for a track-seeded search. + FedSimilaritySearchLoaded { + seq: u64, + result: Result, + }, SimilarityStatus(crate::similarity::SimilarityStatus), /// `None` is emitted after clearing every stored embedding. SimilarityProfileActivated(Option), diff --git a/src/app/mod.rs b/src/app/mod.rs index 8e57e3f..68f3a65 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -3348,6 +3348,39 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent Err(message) => tracing::warn!(%message, "federated search failed"), } } + AppEvent::FedSimilaritySearchLoaded { seq, result } => { + if runtime.search_seq.load(std::sync::atomic::Ordering::SeqCst) != seq { + return; + } + state.search.fed_loading = false; + let selected_key = state.global.stack.last().and_then(|view| match view { + state::GlobalView::Search { cursor } => state.search.similarity_key(*cursor), + _ => None, + }); + match result { + Ok(results) => { + let remote = results.tracks.into_iter().map(|hit| { + state::SimilaritySearchHit::Federated { + track: hit.track, + score: hit.score, + embedding_signature: hit.embedding_signature, + } + }); + state.search.similarity_tracks.extend(remote); + rank_similarity_search_tracks( + &mut state.search.similarity_tracks, + state.similarity.settings.max_tracks_per_artist, + ); + state.search.similarity_stats = Some(results.stats); + state.search.similarity_error = None; + } + Err(message) => { + tracing::warn!(%message, "federated similarity search failed"); + state.search.similarity_error = Some(message); + } + } + restore_similarity_cursor(state, selected_key.as_deref()); + } AppEvent::FedTrackResolved { placeholder_id, resolve_key, @@ -3696,7 +3729,20 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent } state.search.loading = false; match result { - Ok(results) => state.search.results = Some(results), + Ok(results) => { + state.search.similarity_tracks = results + .into_iter() + .map(|hit| state::SimilaritySearchHit::Local { + track: hit.track, + score: hit.score, + embedding_signature: hit.embedding_signature, + }) + .collect(); + rank_similarity_search_tracks( + &mut state.search.similarity_tracks, + state.similarity.settings.max_tracks_per_artist, + ); + } Err(message) => { state.status_message = Some(format!("similarity search failed: {message}")); return; @@ -3714,7 +3760,7 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent .search_similar(query, 50) .await .map_err(|err| format!("{err:#}")); - let _ = tx.send(AppEvent::FedSearchLoaded { seq, result }); + let _ = tx.send(AppEvent::FedSimilaritySearchLoaded { seq, result }); }); } } @@ -4007,6 +4053,135 @@ fn handle_app_event(state: &mut AppState, runtime: &mut Runtime, event: AppEvent } } +fn rank_similarity_search_tracks( + tracks: &mut Vec, + max_tracks_per_artist: usize, +) { + const RESULT_LIMIT: usize = 49; + const MAX_NEAR_DUPLICATE_SIGNATURE_DISTANCE: u32 = 8; + + tracks.sort_by(|left, right| right.score().total_cmp(&left.score())); + let candidates = std::mem::take(tracks); + let mut content = std::collections::HashSet::new(); + let mut signatures = Vec::new(); + let mut artist_counts: std::collections::HashMap = Default::default(); + for hit in candidates { + if !content.insert(hit.content_key()) { + continue; + } + if hit.embedding_signature().is_some_and(|candidate| { + signatures.iter().any(|existing| { + music_dht::similarity::signature_distance(&candidate, existing) + <= MAX_NEAR_DUPLICATE_SIGNATURE_DISTANCE + }) + }) { + continue; + } + let artist = hit.primary_artist_key(); + let count = artist_counts.entry(artist.clone()).or_default(); + if !artist.is_empty() && *count >= max_tracks_per_artist.clamp(1, RESULT_LIMIT) { + continue; + } + *count += 1; + if let Some(signature) = hit.embedding_signature() { + signatures.push(signature); + } + tracks.push(hit); + if tracks.len() >= RESULT_LIMIT { + break; + } + } +} + +fn restore_similarity_cursor(state: &mut AppState, selected_key: Option<&str>) { + let selected_index = selected_key.and_then(|key| state.search.similarity_index_for_key(key)); + let len = state.search.similarity_len(); + if let Some(state::GlobalView::Search { cursor }) = state.global.stack.last_mut() { + *cursor = selected_index.unwrap_or(*cursor).min(len.saturating_sub(1)); + } +} + +#[cfg(test)] +mod similarity_search_tests { + use super::*; + use crate::library::models::{ArtistRef, TrackItem}; + + fn local_hit(id: i64, artist: &str, score: f32, signature: u8) -> state::SimilaritySearchHit { + state::SimilaritySearchHit::Local { + track: TrackItem { + id, + title: format!("local {id}"), + track_number: None, + disc_number: None, + duration_seconds: 1.0, + artists: vec![ArtistRef { + id, + name: artist.to_string(), + }], + featured_artists: Vec::new(), + release_id: id, + release_title: "release".to_string(), + release_year: None, + file_path: format!("/music/{id}"), + content_id: Some(format!("local-{id}")), + cover_path: None, + audio_format: None, + audio_bitrate: None, + audio_sample_rate: None, + audio_bit_depth: None, + file_size_bytes: None, + play_count: 0, + fed: None, + }, + score, + embedding_signature: [signature; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES], + } + } + + fn remote_hit(artist: &str, score: f32, signature: u8) -> state::SimilaritySearchHit { + state::SimilaritySearchHit::Federated { + track: crate::federation::FedTrack { + item_id: format!("remote-{signature}"), + owner: "peer".to_string(), + own: false, + title: format!("remote {signature}"), + artist_names: vec![artist.to_string()], + featured_artist_names: Vec::new(), + year: None, + duration_seconds: Some(1), + content_id: Some(format!("remote-{signature}")), + release_title: None, + track_number: None, + disc_number: None, + }, + score, + embedding_signature: Some( + [signature; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES], + ), + } + } + + #[test] + fn similarity_results_rank_local_and_remote_together_with_one_artist_cap() { + let mut tracks = vec![ + local_hit(1, "same artist", 0.70, 1), + remote_hit("other artist", 0.90, 2), + remote_hit("same artist", 0.80, 3), + local_hit(2, "same artist", 0.60, 4), + ]; + + rank_similarity_search_tracks(&mut tracks, 1); + + assert_eq!(tracks.len(), 2); + assert_eq!(tracks[0].score(), 0.90); + assert_eq!(tracks[1].score(), 0.80); + assert!(matches!( + tracks[0], + state::SimilaritySearchHit::Federated { .. } + )); + } +} + /// Mirror the playback state to the OS now-playing surface. `force` skips /// the position throttle (track switches, pauses). fn push_media_update(state: &AppState, runtime: &mut Runtime, force: bool) { diff --git a/src/app/popup.rs b/src/app/popup.rs index 79509f9..1d82205 100644 --- a/src/app/popup.rs +++ b/src/app/popup.rs @@ -592,6 +592,36 @@ fn handle_fed_input( state.popup = Some(Popup::FedInput { field, input }); } }, + FedInputField::SimilarityMinimumScore => match value.parse::() { + Ok(score) if score.is_finite() && (0.0..=1.0).contains(&score) => { + state.similarity.settings.minimum_score = score; + super::perform_effect( + state, + runtime, + crate::app::update::Effect::SimilarityApplySettings, + ); + } + _ => { + state.status_message = + Some("minimum similarity must be a number from 0.00 to 1.00".into()); + state.popup = Some(Popup::FedInput { field, input }); + } + }, + FedInputField::SimilarityMaxTracksPerArtist => match value.parse::() { + Ok(limit @ 1..=50) => { + state.similarity.settings.max_tracks_per_artist = limit; + super::perform_effect( + state, + runtime, + crate::app::update::Effect::SimilarityApplySettings, + ); + } + _ => { + state.status_message = + Some("tracks per artist must be a number from 1 to 50".into()); + state.popup = Some(Popup::FedInput { field, input }); + } + }, FedInputField::ConnectTicket => { if value.is_empty() { state.status_message = Some("ticket is empty".into()); diff --git a/src/app/state.rs b/src/app/state.rs index dd6c9eb..25d0221 100644 --- a/src/app/state.rs +++ b/src/app/state.rs @@ -509,6 +509,8 @@ pub enum TrackSelectionScope { Release(i64), Playlist(i64), Queue, + /// The unified local + federated similar-track result list. + SimilaritySearch, /// The federated section of the search results (its tracks). FedSearch, /// The tracklist of the open federated release view. @@ -817,6 +819,8 @@ impl StatusDetailFocus { pub enum FedInputField { MusicDirectory, SimilarityWorkers, + SimilarityMinimumScore, + SimilarityMaxTracksPerArtist, NetworkId, ConnectTicket, DeviceName, @@ -829,6 +833,8 @@ impl FedInputField { match self { FedInputField::MusicDirectory => "Music save directory", FedInputField::SimilarityWorkers => "Similarity background workers", + FedInputField::SimilarityMinimumScore => "Minimum similarity score", + FedInputField::SimilarityMaxTracksPerArtist => "Tracks per artist", FedInputField::NetworkId => "Network ID", FedInputField::ConnectTicket => "Connect to peer (paste ticket)", FedInputField::DeviceName => "Device name", @@ -845,6 +851,12 @@ impl FedInputField { FedInputField::SimilarityWorkers => { "Enter the maximum number of tracks processed in parallel, from 1 to 16. The change takes effect immediately." } + FedInputField::SimilarityMinimumScore => { + "Enter the minimum cosine similarity from 0.00 to 1.00. Lower values show broader matches; higher values hide weak matches. Embeddings are not recalculated." + } + FedInputField::SimilarityMaxTracksPerArtist => { + "Enter how many tracks by one primary artist may appear in similarity results, from 1 to 50. Embeddings are not recalculated." + } FedInputField::NetworkId => { "A unique network id. It must match exactly on every client that should see and connect to the same peers." } @@ -880,15 +892,19 @@ pub enum SimilarityRow { Toggle, Model, Profile, + MinimumScore, + MaxTracksPerArtist, Workers, Clear, } impl SimilarityRow { - pub const ALL: [SimilarityRow; 5] = [ + pub const ALL: [SimilarityRow; 7] = [ SimilarityRow::Toggle, SimilarityRow::Model, SimilarityRow::Profile, + SimilarityRow::MinimumScore, + SimilarityRow::MaxTracksPerArtist, SimilarityRow::Workers, SimilarityRow::Clear, ]; @@ -1211,6 +1227,85 @@ mod cmdline_history_tests { } /// Live search state driven by the `:/query` command. +#[derive(Debug, Clone)] +pub enum SimilaritySearchHit { + Local { + track: TrackItem, + score: f32, + embedding_signature: [u8; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES], + }, + Federated { + track: crate::federation::FedTrack, + score: f32, + embedding_signature: Option<[u8; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES]>, + }, +} + +impl SimilaritySearchHit { + pub fn score(&self) -> f32 { + match self { + Self::Local { score, .. } | Self::Federated { score, .. } => *score, + } + } + + pub fn embedding_signature( + &self, + ) -> Option<[u8; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES]> { + match self { + Self::Local { + embedding_signature, + .. + } => Some(*embedding_signature), + Self::Federated { + embedding_signature, + .. + } => *embedding_signature, + } + } + + pub fn track_item(&self) -> TrackItem { + match self { + Self::Local { track, .. } => track.clone(), + Self::Federated { track, .. } => crate::federation::pending_track(track), + } + } + + pub fn federated_track(&self) -> Option<&crate::federation::FedTrack> { + match self { + Self::Federated { track, .. } => Some(track), + Self::Local { .. } => None, + } + } + + pub fn primary_artist_key(&self) -> String { + match self { + Self::Local { track, .. } => track + .artists + .first() + .map(|artist| music_dht::normalize_name(&artist.name)) + .unwrap_or_default(), + Self::Federated { track, .. } => track + .artist_names + .first() + .map(|artist| music_dht::normalize_name(artist)) + .unwrap_or_default(), + } + } + + pub fn content_key(&self) -> String { + match self { + Self::Local { track, .. } => track + .content_id + .clone() + .unwrap_or_else(|| format!("local:{}", track.id)), + Self::Federated { track, .. } => track + .content_id + .clone() + .unwrap_or_else(|| format!("remote:{}:{}", track.owner, track.item_id)), + } + } +} + #[derive(Debug, Default)] pub struct SearchState { pub query: String, @@ -1226,6 +1321,59 @@ pub struct SearchState { /// Present only for a track-seeded search; text-search refreshes must not /// replace this page with a title query. pub similarity_source: Option, + /// The source is pinned at row zero; candidates below it are globally + /// ranked across the local library and federation. + pub similarity_source_track: Option, + pub similarity_tracks: Vec, + pub similarity_stats: Option, + pub similarity_error: Option, +} + +impl SearchState { + pub fn similarity_len(&self) -> usize { + usize::from(self.similarity_source_track.is_some()) + self.similarity_tracks.len() + } + + pub fn similarity_track(&self, index: usize) -> Option { + if index == 0 { + return self.similarity_source_track.clone(); + } + self.similarity_tracks + .get(index.checked_sub(1)?) + .map(SimilaritySearchHit::track_item) + } + + pub fn similarity_fed_track(&self, index: usize) -> Option<&crate::federation::FedTrack> { + self.similarity_tracks + .get(index.checked_sub(1)?)? + .federated_track() + } + + pub fn similarity_key(&self, index: usize) -> Option { + if index == 0 { + return self + .similarity_source_track + .as_ref() + .map(|track| format!("source:{}", track.id)); + } + self.similarity_tracks + .get(index.checked_sub(1)?) + .map(SimilaritySearchHit::content_key) + } + + pub fn similarity_index_for_key(&self, key: &str) -> Option { + if self + .similarity_source_track + .as_ref() + .is_some_and(|track| key == format!("source:{}", track.id)) + { + return Some(0); + } + self.similarity_tracks + .iter() + .position(|hit| hit.content_key() == key) + .map(|index| index + 1) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] diff --git a/src/app/update.rs b/src/app/update.rs index 0b1310b..af5ba62 100644 --- a/src/app/update.rs +++ b/src/app/update.rs @@ -905,6 +905,9 @@ fn set_track_scope_cursor(state: &mut AppState, scope: &TrackSelectionScope, val TrackSelectionScope::Queue => { state.queue_tab.cursor = value; } + TrackSelectionScope::SimilaritySearch => { + set_view_cursor(state, value); + } TrackSelectionScope::FedSearch => { let base = state.search.results.as_ref().map_or(0, |r| r.len()) + state.search.fed_artists.len(); @@ -952,6 +955,14 @@ fn current_track_list_context(state: &AppState) -> Option<(TrackSelectionScope, _ => None, }, GlobalView::Search { cursor } => { + if state.search.similarity_source.is_some() { + let len = state.search.similarity_len(); + return (*cursor < len).then_some(( + TrackSelectionScope::SimilaritySearch, + *cursor, + len, + )); + } // Only the federated tracks section is selectable here. let base = state.search.results.as_ref().map_or(0, |r| r.len()) + state.search.fed_artists.len(); @@ -1044,6 +1055,20 @@ fn current_track_list(state: &AppState) -> Option<(TrackSelectionScope, usize, V } pub fn selected_tracks(state: &AppState) -> Vec { + if state.active_tab == Tab::Global + && state.search.similarity_source.is_some() + && let Some(GlobalView::Search { cursor }) = state.global.stack.last() + { + let len = state.search.similarity_len(); + let indices = state + .track_selection + .indices(&TrackSelectionScope::SimilaritySearch, len) + .unwrap_or_else(|| vec![(*cursor).min(len.saturating_sub(1))]); + return indices + .into_iter() + .filter_map(|index| state.search.similarity_track(index)) + .collect(); + } // Federated contexts produce queueable placeholders that behave like // regular tracks (queue, info, playback-on-demand). { @@ -1276,6 +1301,9 @@ pub fn selected_track(state: &AppState) -> Option { _ => None, }, GlobalView::Search { cursor } => { + if state.search.similarity_source.is_some() { + return state.search.similarity_track(*cursor); + } let results = state.search.results.as_ref()?; let offset = cursor.checked_sub(results.artists.len() + results.releases.len())?; match results.tracks.get(offset) { @@ -1877,10 +1905,13 @@ fn move_selection(state: &mut AppState, dx: isize, dy: isize) { refresh_track_selection_cursor(state); } Some(GlobalView::Search { cursor }) => { - // Local results plus the federated section below them. - let total = (state.search.results.as_ref().map_or(0, |r| r.len()) - + state.search.fed_artists.len() - + state.search.fed_tracks.len()) as isize; + let total = if state.search.similarity_source.is_some() { + state.search.similarity_len() + } else { + state.search.results.as_ref().map_or(0, |r| r.len()) + + state.search.fed_artists.len() + + state.search.fed_tracks.len() + } as isize; if total == 0 { return; } @@ -2000,9 +2031,13 @@ fn current_view_len(state: &AppState) -> usize { _ => 0, }, Some(GlobalView::Search { .. }) => { - state.search.results.as_ref().map_or(0, |r| r.len()) - + state.search.fed_artists.len() - + state.search.fed_tracks.len() + if state.search.similarity_source.is_some() { + state.search.similarity_len() + } else { + state.search.results.as_ref().map_or(0, |r| r.len()) + + state.search.fed_artists.len() + + state.search.fed_tracks.len() + } } Some(GlobalView::FedArtist { .. }) => fed_card_len(state), Some(GlobalView::FedRelease { index, .. }) => fed_card_release(state, *index) @@ -2292,31 +2327,47 @@ fn select_current(state: &mut AppState) -> Option { }, _ => Outcome::Nothing, }, - Some(GlobalView::Search { cursor }) => match &state.search.results { - Some(results) => { - let artists = results.artists.len(); - let releases = results.releases.len(); - if cursor < artists { - Outcome::Push(GlobalView::Artist { - id: results.artists[cursor].id, - cursor: 0, - }) - } else if cursor < artists + releases { - Outcome::Push(GlobalView::Release { - id: results.releases[cursor - artists].id, - cursor: 0, - }) - } else if results.tracks.get(cursor - artists - releases).is_some() { - Outcome::Play { - tracks: results.tracks.clone(), - start: cursor - artists - releases, - } + Some(GlobalView::Search { cursor }) => { + if state.search.similarity_source.is_some() { + let tracks = (0..state.search.similarity_len()) + .filter_map(|index| state.search.similarity_track(index)) + .collect::>(); + if tracks.is_empty() { + Outcome::Nothing } else { - fed_outcome(state, cursor - artists - releases - results.tracks.len()) + Outcome::Play { + start: cursor.min(tracks.len() - 1), + tracks, + } + } + } else { + match &state.search.results { + Some(results) => { + let artists = results.artists.len(); + let releases = results.releases.len(); + if cursor < artists { + Outcome::Push(GlobalView::Artist { + id: results.artists[cursor].id, + cursor: 0, + }) + } else if cursor < artists + releases { + Outcome::Push(GlobalView::Release { + id: results.releases[cursor - artists].id, + cursor: 0, + }) + } else if results.tracks.get(cursor - artists - releases).is_some() { + Outcome::Play { + tracks: results.tracks.clone(), + start: cursor - artists - releases, + } + } else { + fed_outcome(state, cursor - artists - releases - results.tracks.len()) + } + } + None => fed_outcome(state, cursor), } } - None => fed_outcome(state, cursor), - }, + } Some(GlobalView::FedArtist { cursor }) => match &state.fed_artist_view { Some((_, Loadable::Ready(card))) => { let release_indices = fed_artist_visible_release_indices(state, card); @@ -2492,6 +2543,20 @@ pub(crate) fn selected_fed_tracks(state: &AppState) -> Vec Option { }); None } + SettingsRow::Similarity(SimilarityRow::MinimumScore) => { + state.popup = Some(Popup::FedInput { + field: FedInputField::SimilarityMinimumScore, + input: crate::app::input::LineEdit::new(format!( + "{:.2}", + state.similarity.settings.minimum_score + )), + }); + None + } + SettingsRow::Similarity(SimilarityRow::MaxTracksPerArtist) => { + state.popup = Some(Popup::FedInput { + field: FedInputField::SimilarityMaxTracksPerArtist, + input: crate::app::input::LineEdit::new( + state.similarity.settings.max_tracks_per_artist.to_string(), + ), + }); + None + } SettingsRow::Similarity(SimilarityRow::Workers) => { state.popup = Some(Popup::FedInput { field: FedInputField::SimilarityWorkers, diff --git a/src/config/settings.rs b/src/config/settings.rs index 21d485d..2183641 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -45,7 +45,7 @@ pub struct LibraryFilters { pub source_mode: LibrarySourceMode, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct SimilaritySettings { /// Local embedding/search master switch. Network participation follows /// federation and additionally requires the explicit privacy consent. @@ -57,6 +57,14 @@ pub struct SimilaritySettings { pub profile: String, #[serde(default = "default_similarity_workers")] pub workers: usize, + /// Requester-side cosine score floor. This is search policy, not part of + /// the embedding profile, so changing it never invalidates vectors. + #[serde(default = "default_similarity_minimum_score")] + pub minimum_score: f32, + /// Requester-side diversity cap applied independently to local and + /// federated candidates. + #[serde(default = "default_similarity_max_tracks_per_artist")] + pub max_tracks_per_artist: usize, #[serde(default)] pub federation_consent: bool, /// Exact fingerprint of the last fully usable profile. Keeping this @@ -73,13 +81,15 @@ impl Default for SimilaritySettings { model: default_similarity_model(), profile: default_similarity_profile(), workers: default_similarity_workers(), + minimum_score: default_similarity_minimum_score(), + max_tracks_per_artist: default_similarity_max_tracks_per_artist(), federation_consent: false, active_profile: None, } } } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AppSettings { #[serde(default = "default_volume")] pub volume: u8, @@ -116,6 +126,11 @@ impl AppSettings { self.similarity.profile = default_similarity_profile(); } self.similarity.workers = self.similarity.workers.clamp(1, 16); + if !self.similarity.minimum_score.is_finite() { + self.similarity.minimum_score = default_similarity_minimum_score(); + } + self.similarity.minimum_score = self.similarity.minimum_score.clamp(0.0, 1.0); + self.similarity.max_tracks_per_artist = self.similarity.max_tracks_per_artist.clamp(1, 50); self } } @@ -138,6 +153,14 @@ fn default_similarity_workers() -> usize { .unwrap_or(1) } +fn default_similarity_minimum_score() -> f32 { + 0.70 +} + +fn default_similarity_max_tracks_per_artist() -> usize { + 5 +} + /// The historical permanent-download location, kept as the default for /// backward compatibility with existing installations. pub fn default_music_dir() -> PathBuf { @@ -204,5 +227,21 @@ hide_featured_only = true assert!(settings.library.hide_featured_only); assert_eq!(settings.library.source_mode, LibrarySourceMode::Global); assert_eq!(settings.music_dir, default_music_dir()); + assert_eq!(settings.similarity.minimum_score, 0.70); + assert_eq!(settings.similarity.max_tracks_per_artist, 5); + } + + #[test] + fn similarity_search_policy_is_normalized_without_changing_the_profile() { + let mut settings = AppSettings::default(); + let profile = settings.similarity.profile.clone(); + settings.similarity.minimum_score = f32::NAN; + settings.similarity.max_tracks_per_artist = 0; + + let settings = settings.normalized(); + + assert_eq!(settings.similarity.minimum_score, 0.70); + assert_eq!(settings.similarity.max_tracks_per_artist, 1); + assert_eq!(settings.similarity.profile, profile); } } diff --git a/src/federation/mod.rs b/src/federation/mod.rs index 3aef329..2784033 100644 --- a/src/federation/mod.rs +++ b/src/federation/mod.rs @@ -338,6 +338,27 @@ pub struct FedSearchResults { pub tracks: Vec, } +#[derive(Debug, Clone)] +pub struct ScoredFedTrack { + pub track: FedTrack, + pub score: f32, + pub embedding_signature: Option<[u8; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES]>, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct SimilaritySearchStats { + pub tracks: usize, + pub artists: usize, + pub peers_queried: usize, + pub elapsed_ms: u64, +} + +#[derive(Debug, Clone, Default)] +pub struct FedSimilaritySearchResults { + pub tracks: Vec, + pub stats: SimilaritySearchStats, +} + /// A track found through federated search. #[derive(Debug, Clone, PartialEq, Eq)] pub struct FedTrack { @@ -1213,17 +1234,19 @@ impl Federation { &self, query: crate::similarity::QueryVector, limit: usize, - ) -> Result { + ) -> Result { anyhow::ensure!( self.similarity.network_allowed(), "similarity federation has no consent" ); let (service, similarity_dht) = self.similarity_services().await?; + let settings = self.similarity.settings(); similarity::search( service, similarity_dht, query, limit, + settings.minimum_score, Arc::clone(&self.transport_stats), ) .await diff --git a/src/federation/similarity.rs b/src/federation/similarity.rs index 8bce47d..d2991dd 100644 --- a/src/federation/similarity.rs +++ b/src/federation/similarity.rs @@ -4,9 +4,9 @@ //! module owns application policy: consent, peer fan-out, local index access, //! result conversion, deduplication, and ranking limits. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use anyhow::{Context as _, Result}; use futures_util::stream::{self, StreamExt as _}; @@ -16,7 +16,9 @@ use music_dht::{ ByteStream, EndpointId, ItemId, ItemKind, MusicDhtService, PeerTicket, StreamAcceptor, }; -use crate::federation::{FedSearchResults, FedTrack, TransportStats}; +use crate::federation::{ + FedSimilaritySearchResults, FedTrack, ScoredFedTrack, SimilaritySearchStats, TransportStats, +}; use crate::similarity::{Manager, QueryVector}; pub use music_dht::similarity::SIMILARITY_ALPN; @@ -26,7 +28,6 @@ const MAX_QUERY_PEERS: usize = 48; const QUERY_CONCURRENCY: usize = 8; const QUERY_TIMEOUT: Duration = Duration::from_secs(5); const ROUTING_TIMEOUT: Duration = Duration::from_secs(5); -const MAX_PER_ARTIST: usize = 3; const MAX_NEAR_DUPLICATE_SIGNATURE_DISTANCE: u32 = 8; pub async fn serve_peers( @@ -62,7 +63,7 @@ async fn serve_one( let vector = request.vector; let limit = request.limit; let matches = tokio::task::spawn_blocking(move || { - similarity.search_vector(&profile, &vector, None, None, limit) + similarity.search_vector_for_peer(&profile, &vector, limit) }) .await .context("local similarity task failed") @@ -130,8 +131,10 @@ pub async fn search( routing: Arc, query: QueryVector, limit: usize, + minimum_score: f32, transport: Arc, -) -> Result { +) -> Result { + let started = Instant::now(); let own = service.endpoint_id(); let routed = match tokio::time::timeout( ROUTING_TIMEOUT, @@ -192,6 +195,7 @@ pub async fn search( ) .await; let mut successful = 0usize; + let mut peers_queried = initial; for response in responses { match response { Ok(peer_hits) => { @@ -202,6 +206,7 @@ pub async fn search( } } if initial < peers.len() && (hits.len() < limit || successful < initial.min(4)) { + peers_queried = peers.len(); for response in query_peers( Arc::clone(&service), &peers[initial..], @@ -219,9 +224,11 @@ pub async fn search( hits.sort_by(|left, right| right.1.total_cmp(&left.1)); let mut dedup = HashSet::new(); let mut embedding_signatures = vec![query_signature]; - let mut artist_counts: HashMap = HashMap::new(); let mut tracks = Vec::new(); - for (track, _, embedding_signature) in hits { + for (track, score, embedding_signature) in hits { + if score < minimum_score { + break; + } if query .source_content_id .as_deref() @@ -244,26 +251,33 @@ pub async fn search( }) { continue; } - let artist = track - .artist_names - .first() - .map(|name| music_dht::normalize_name(name)) - .unwrap_or_default(); - let count = artist_counts.entry(artist.clone()).or_default(); - if !artist.is_empty() && *count >= MAX_PER_ARTIST { - continue; - } - *count += 1; if let Some(signature) = embedding_signature { embedding_signatures.push(signature); } - tracks.push(track); + tracks.push(ScoredFedTrack { + track, + score, + embedding_signature, + }); if tracks.len() >= limit.min(wire::MAX_SIMILARITY_RESULTS) { break; } } - Ok(FedSearchResults { - artists: Vec::new(), + let artists = tracks + .iter() + .filter_map(|hit| hit.track.artist_names.first()) + .map(|name| music_dht::normalize_name(name)) + .filter(|name| !name.is_empty()) + .collect::>() + .len(); + let elapsed_ms = started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64; + Ok(FedSimilaritySearchResults { + stats: SimilaritySearchStats { + tracks: tracks.len(), + artists, + peers_queried, + elapsed_ms, + }, tracks, }) } diff --git a/src/similarity.rs b/src/similarity.rs index 4c732a8..60f7907 100644 --- a/src/similarity.rs +++ b/src/similarity.rs @@ -36,7 +36,7 @@ const EMBEDDING_DIMENSIONS: usize = 1280; const MODEL_BATCH: usize = 8; const MAX_MODEL_BYTES: usize = 64 * 1024 * 1024; const RESULT_LIMIT: usize = 50; -const MAX_PER_ARTIST: usize = 3; +const PEER_CANDIDATE_MAX_PER_ARTIST: usize = 10; const NEAR_DUPLICATE_COSINE: f32 = 0.995; const FULL_TRACK_MAX_SECONDS: u32 = 5 * 60; const LONG_TRACK_WINDOW_SECONDS: u32 = 60; @@ -368,6 +368,52 @@ impl Manager { exclude_track_id: Option, exclude_content_id: Option<&str>, limit: usize, + ) -> Result> { + let settings = lock(&self.settings); + let minimum_score = settings.minimum_score; + let max_tracks_per_artist = settings.max_tracks_per_artist; + drop(settings); + self.search_vector_with_policy( + profile_id, + vector, + exclude_track_id, + exclude_content_id, + limit, + minimum_score, + max_tracks_per_artist, + ) + } + + /// Returns a wider, policy-neutral candidate set to a remote requester. + /// The requester applies its own score threshold and artist diversity + /// limit; neither value is part of embedding compatibility. + pub(crate) fn search_vector_for_peer( + &self, + profile_id: &str, + vector: &[f32], + limit: usize, + ) -> Result> { + self.search_vector_with_policy( + profile_id, + vector, + None, + None, + limit, + -1.0, + PEER_CANDIDATE_MAX_PER_ARTIST, + ) + } + + #[allow(clippy::too_many_arguments)] + fn search_vector_with_policy( + &self, + profile_id: &str, + vector: &[f32], + exclude_track_id: Option, + exclude_content_id: Option<&str>, + limit: usize, + minimum_score: f32, + max_tracks_per_artist: usize, ) -> Result> { anyhow::ensure!( !vector.is_empty() && vector.len() <= 4096, @@ -396,7 +442,7 @@ impl Manager { entry.vector.as_slice(), ) }) - .filter(|(_, score, _, _)| score.is_finite()) + .filter(|(_, score, _, _)| score.is_finite() && *score >= minimum_score) .collect(); scores.sort_by(|left, right| right.1.total_cmp(&left.1)); @@ -410,7 +456,7 @@ impl Manager { continue; } let count = artist_counts.entry(artist.to_string()).or_default(); - if !artist.is_empty() && *count >= MAX_PER_ARTIST { + if !artist.is_empty() && *count >= max_tracks_per_artist { continue; } *count += 1; diff --git a/src/ui/federation.rs b/src/ui/federation.rs index 29ade78..90edaa8 100644 --- a/src/ui/federation.rs +++ b/src/ui/federation.rs @@ -103,6 +103,14 @@ fn draw_settings_rows(frame: &mut Frame, area: Rect, state: &AppState) { "Preprocessing profile", format!("{} (enter for details)", similarity.profile), ), + SimilarityRow::MinimumScore => ( + "Minimum similarity", + format!("{:.2}", similarity.minimum_score), + ), + SimilarityRow::MaxTracksPerArtist => ( + "Tracks per artist", + similarity.max_tracks_per_artist.to_string(), + ), SimilarityRow::Workers => ("Background workers", similarity.workers.to_string()), SimilarityRow::Clear => ("Clear all stored embeddings", "↵".to_string()), }; diff --git a/src/ui/global.rs b/src/ui/global.rs index ba3a004..cbadfd9 100644 --- a/src/ui/global.rs +++ b/src/ui/global.rs @@ -900,6 +900,9 @@ fn draw_search(frame: &mut Frame, area: Rect, state: &AppState, cursor: usize) { title.push_str("· searching… "); } let inner = bordered(frame, area, state, title); + if search.similarity_source.is_some() { + return draw_similarity_search(frame, inner, state, cursor); + } let empty_results = SearchResults::default(); let results = match &search.results { @@ -1102,6 +1105,156 @@ fn draw_search(frame: &mut Frame, area: Rect, state: &AppState, cursor: usize) { } } +fn draw_similarity_search(frame: &mut Frame, area: Rect, state: &AppState, cursor: usize) { + let search = &state.search; + let status = if search.fed_loading { + super::loading_line(state, "searching federation…") + } else if let Some(stats) = &search.similarity_stats { + let elapsed = if stats.elapsed_ms < 1_000 { + format!("{} ms", stats.elapsed_ms) + } else { + format!("{:.2} s", stats.elapsed_ms as f64 / 1_000.0) + }; + Line::from(vec![ + Span::styled("Federation · ", theme::accent_for(state)), + Span::styled( + format!( + "{} tracks · {} artists · {} peers queried · {elapsed}", + stats.tracks, stats.artists, stats.peers_queried + ), + theme::dim(), + ), + ]) + } else if search.similarity_error.is_some() { + Line::styled( + "Federation search failed · showing local results", + error_style(), + ) + } else if search.loading { + super::loading_line(state, "preparing local similarity search…") + } else { + Line::styled("Federation disabled · showing local results", theme::dim()) + }; + let [status_area, content] = + Layout::vertical([Constraint::Length(1), Constraint::Min(0)]).areas(area); + frame.render_widget(Paragraph::new(status), status_area); + + let mut rows: Vec<(Line, Option, Option)> = Vec::new(); + rows.push((Line::styled("Tracks", theme::header_for(state)), None, None)); + if let Some(track) = &search.similarity_source_track { + let heart = if state.track_liked(track) { + Span::styled("♥ ", theme::accent_for(state)) + } else { + Span::raw(" ") + }; + rows.push(( + Line::from(vec![ + heart, + Span::raw(track.title.clone()), + Span::styled( + format!(" {} · {}", track.artist_line(), track.release_title), + theme::dim(), + ), + ]), + Some(super::track_meta_suffix(track, true)), + Some(0), + )); + } + for (offset, hit) in search.similarity_tracks.iter().enumerate() { + let index = offset + 1; + match hit { + crate::app::state::SimilaritySearchHit::Local { track, .. } => { + let heart = if state.track_liked(track) { + Span::styled("♥ ", theme::accent_for(state)) + } else { + Span::raw(" ") + }; + rows.push(( + Line::from(vec![ + heart, + Span::raw(track.title.clone()), + Span::styled( + format!(" {} · {}", track.artist_line(), track.release_title), + theme::dim(), + ), + ]), + Some(super::track_meta_suffix(track, true)), + Some(index), + )); + } + crate::app::state::SimilaritySearchHit::Federated { track, .. } => { + let heart = if state.fed_track_liked(track) { + Span::styled("♥ ", theme::accent_for(state)) + } else { + Span::raw(" ") + }; + let origin = if track.own { + "your library".to_string() + } else { + format!("peer {}…", track.owner_short()) + }; + let mut meta = track.duration_label(); + if let Some(year) = track.year { + if !meta.is_empty() { + meta.push_str(" · "); + } + meta.push_str(&year.to_string()); + } + rows.push(( + Line::from(vec![ + heart, + fed_track_availability_prefix(state, track), + Span::raw(track.title.clone()), + Span::styled( + format!(" {} · {origin}", track.artist_line()), + theme::dim(), + ), + ]), + Some(meta), + Some(index), + )); + } + } + } + + let scope = crate::app::state::TrackSelectionScope::SimilaritySearch; + let mut selected = std::collections::HashSet::new(); + if state.track_selection.is_active_for(&scope) + && let Some(indices) = state + .track_selection + .indices(&scope, search.similarity_len()) + { + selected.extend(indices); + } + let cursor_row = rows + .iter() + .position(|(_, _, row_cursor)| *row_cursor == Some(cursor)) + .unwrap_or(0); + let visible = usize::from(content.height.max(1)); + let first = cursor_row + .saturating_sub(visible / 2) + .min(rows.len().saturating_sub(visible)); + for (offset, (line, right, row_cursor)) in + rows.into_iter().enumerate().skip(first).take(visible) + { + let rect = Rect { + x: content.x, + y: content.y + (offset - first) as u16, + width: content.width, + height: 1, + }; + if let Some(row_index) = row_cursor + && selected.contains(&row_index) + && row_index != cursor + { + frame + .buffer_mut() + .set_style(rect, theme::selection_for(state)); + } + draw_row(frame, rect, state, line, right, row_cursor == Some(cursor)); + } +} + // --------------------------------------------------------------------------- // Federated artist card (assembled from peer catalogs) // ---------------------------------------------------------------------------