Merge local and federated similarity results

This commit is contained in:
Aleksandr Bogomiakov
2026-08-10 20:51:21 +01:00
parent add764e51d
commit 079d87a831
12 changed files with 794 additions and 75 deletions
+9 -15
View File
@@ -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),
+6 -1
View File
@@ -36,9 +36,14 @@ pub enum AppEvent {
/// text search so stale pages cannot overwrite a newer request.
SimilaritySearchLoaded {
seq: u64,
result: Result<SearchResults, String>,
result: Result<Vec<crate::similarity::SimilarTrack>, String>,
query: Option<crate::similarity::QueryVector>,
},
/// Federated candidates and diagnostics for a track-seeded search.
FedSimilaritySearchLoaded {
seq: u64,
result: Result<crate::federation::FedSimilaritySearchResults, String>,
},
SimilarityStatus(crate::similarity::SimilarityStatus),
/// `None` is emitted after clearing every stored embedding.
SimilarityProfileActivated(Option<String>),
+177 -2
View File
@@ -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<state::SimilaritySearchHit>,
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<String, usize> = 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) {
+30
View File
@@ -592,6 +592,36 @@ fn handle_fed_input(
state.popup = Some(Popup::FedInput { field, input });
}
},
FedInputField::SimilarityMinimumScore => match value.parse::<f32>() {
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::<usize>() {
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());
+149 -1
View File
@@ -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<i64>,
/// The source is pinned at row zero; candidates below it are globally
/// ranked across the local library and federation.
pub similarity_source_track: Option<TrackItem>,
pub similarity_tracks: Vec<SimilaritySearchHit>,
pub similarity_stats: Option<crate::federation::SimilaritySearchStats>,
pub similarity_error: Option<String>,
}
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<TrackItem> {
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<String> {
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<usize> {
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)]
+113 -29
View File
@@ -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<TrackItem> {
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<TrackItem> {
_ => 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<Effect> {
},
_ => 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::<Vec<_>>();
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<crate::federation::Fe
if state.active_tab != Tab::Global {
return Vec::new();
}
if state.search.similarity_source.is_some()
&& let Some(GlobalView::Search { cursor }) = state.global.stack.last()
{
let scope = TrackSelectionScope::SimilaritySearch;
let len = state.search.similarity_len();
let indices = state
.track_selection
.indices(&scope, len)
.unwrap_or_else(|| vec![(*cursor).min(len.saturating_sub(1))]);
return indices
.into_iter()
.filter_map(|index| state.search.similarity_fed_track(index).cloned())
.collect();
}
// An active Shift-V range in a federated scope.
if let Some(scope) = state.track_selection.scope.clone() {
match scope {
@@ -2796,6 +2861,25 @@ fn federation_select(state: &mut AppState) -> Option<Effect> {
});
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,
+41 -2
View File
@@ -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);
}
}
+24 -1
View File
@@ -338,6 +338,27 @@ pub struct FedSearchResults {
pub tracks: Vec<FedTrack>,
}
#[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<ScoredFedTrack>,
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<FedSearchResults> {
) -> Result<FedSimilaritySearchResults> {
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
+35 -21
View File
@@ -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<SimilarityDht>,
query: QueryVector,
limit: usize,
minimum_score: f32,
transport: Arc<TransportStats>,
) -> Result<FedSearchResults> {
) -> Result<FedSimilaritySearchResults> {
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<String, usize> = 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::<HashSet<_>>()
.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,
})
}
+49 -3
View File
@@ -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<i64>,
exclude_content_id: Option<&str>,
limit: usize,
) -> Result<Vec<SimilarTrack>> {
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<Vec<SimilarTrack>> {
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<i64>,
exclude_content_id: Option<&str>,
limit: usize,
minimum_score: f32,
max_tracks_per_artist: usize,
) -> Result<Vec<SimilarTrack>> {
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;
+8
View File
@@ -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()),
};
+153
View File
@@ -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<String>, Option<usize>)> = 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)
// ---------------------------------------------------------------------------