Furumi/furumi-tui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

Add federated music similarity search

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
src/federation/capabilities.rs
+1
View File
@@ -170,6 +170,7 @@ mod tests {
"music_dht",
"catalog",
"audio",
"similarity",
"device_sync",
"jam",
] {
src/federation/mod.rs
+32
View File
@@ -15,6 +15,7 @@
mod audio;
mod capabilities;
pub mod catalog;
mod similarity;
use std::collections::{HashMap, VecDeque};
use std::path::{Path, PathBuf};
@@ -39,6 +40,7 @@ use crate::library::models::{ArtistRef, TrackItem};
pub use audio::{AUDIO_ALPN, DownloadProgress, StreamingStart, TrackMetadata};
pub use capabilities::ProtocolVersions;
pub use catalog::{CATALOG_ALPN, FedAppearsOn, FedArtistCard, FedCardTrack, FedRelease};
pub use similarity::SIMILARITY_ALPN;
/// How often the published library is re-synchronized with the local index.
const SYNC_INTERVAL: Duration = Duration::from_secs(60);
@@ -420,6 +422,7 @@ pub struct Federation {
library: Arc<Library>,
devices: Arc<crate::devices::DeviceSync>,
jam: Arc<crate::jam::JamManager>,
similarity: Arc<crate::similarity::Manager>,
data_dir: PathBuf,
cache_dir: PathBuf,
media_dir: std::sync::Mutex<PathBuf>,
@@ -526,6 +529,7 @@ impl Federation {
library: Arc<Library>,
devices: Arc<crate::devices::DeviceSync>,
jam: Arc<crate::jam::JamManager>,
similarity: Arc<crate::similarity::Manager>,
media_dir: PathBuf,
) -> Arc<Self> {
let dirs = crate::config::project_dirs();
@@ -549,6 +553,7 @@ impl Federation {
library,
devices,
jam,
similarity,
data_dir,
cache_dir,
media_dir: std::sync::Mutex::new(media_dir),
@@ -692,6 +697,8 @@ impl Federation {
.stream_protocol(AUDIO_ALPN)
// ...and browse each other's per-artist catalogs over this one.
.stream_protocol(CATALOG_ALPN)
// Anonymous, bounded direct embedding queries.
.stream_protocol(SIMILARITY_ALPN)
// Personal-device sync (likes, playlists, trusted devices).
.stream_protocol(crate::devices::SYNC_ALPN)
// Capability-scoped shared playback control.
@@ -746,6 +753,15 @@ impl Federation {
service.endpoint_id(),
Arc::clone(&self.transport_stats),
));
let similarity_acceptor = service
.stream_acceptor(SIMILARITY_ALPN)
.map_err(|err| anyhow::anyhow!("failed to take the similarity acceptor: {err}"))?;
let similarity_task = tokio::spawn(similarity::serve_peers(
similarity_acceptor,
Arc::clone(&self.similarity),
service.endpoint_id(),
Arc::clone(&self.transport_stats),
));
let sync_acceptor = service
.stream_acceptor(crate::devices::SYNC_ALPN)
.map_err(|err| anyhow::anyhow!("failed to take the device-sync acceptor: {err}"))?;
@@ -788,6 +804,7 @@ impl Federation {
sync_task,
audio_task,
catalog_task,
similarity_task,
device_sync_task,
device_tick_task,
jam_serve_task,
@@ -1147,6 +1164,21 @@ impl Federation {
Ok(FedSearchResults { artists, tracks })
}
/// Bounded fan-out to known peers using the exact model/profile
/// fingerprint carried with the query. No DHT records are written.
pub async fn search_similar(
&self,
query: crate::similarity::QueryVector,
limit: usize,
) -> Result<FedSearchResults> {
anyhow::ensure!(
self.similarity.network_allowed(),
"similarity federation has no consent"
);
let service = self.service().await?;
similarity::search(service, query, limit, Arc::clone(&self.transport_stats)).await
}
/// Resolves a share-link content id to one playable federated track.
///
/// Resolution order: the in-session metadata cache, the DHT content key,
src/federation/similarity.rs
+277
View File
@@ -0,0 +1,277 @@
//! Furumi policy and local-index adapter for the shared similarity protocol.
//!
//! `music_dht::similarity` owns the versioned wire contract and framing. This
//! module owns application policy: consent, peer fan-out, local index access,
//! result conversion, deduplication, and ranking limits.
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use anyhow::{Context as _, Result};
use futures_util::stream::{self, StreamExt as _};
use music_dht::similarity::{self as wire, SimilarityHit, SimilarityRequest, SimilarityResponse};
use music_dht::{ByteStream, EndpointId, ItemId, ItemKind, MusicDhtService, StreamAcceptor};
use crate::federation::{FedSearchResults, FedTrack, TransportStats};
use crate::similarity::{Manager, QueryVector};
pub use music_dht::similarity::SIMILARITY_ALPN;
const MAX_QUERY_PEERS: usize = 16;
const QUERY_CONCURRENCY: usize = 6;
const QUERY_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(
mut acceptor: StreamAcceptor,
similarity: Arc<Manager>,
own: EndpointId,
transport: Arc<TransportStats>,
) {
while let Some(stream) = acceptor.accept().await {
let similarity = Arc::clone(&similarity);
let transport = Arc::clone(&transport);
tokio::spawn(async move {
let peer = stream.peer_id;
if let Err(err) = serve_one(stream, similarity, own, transport).await {
tracing::warn!(peer = %peer, "similarity request failed: {err:#}");
}
});
}
}
async fn serve_one(
mut stream: ByteStream,
similarity: Arc<Manager>,
own: EndpointId,
transport: Arc<TransportStats>,
) -> Result<()> {
super::record_stream_transport(&transport, "similarity", "inbound", "open", &stream);
let request = wire::read_request(&mut stream).await?;
let response = if !similarity.network_allowed() {
SimilarityResponse::refused("similarity federation is disabled or has no privacy consent")?
} else {
let profile = request.profile_id;
let vector = request.vector;
let limit = request.limit;
let matches = tokio::task::spawn_blocking(move || {
similarity.search_vector(&profile, &vector, None, None, limit)
})
.await
.context("local similarity task failed")
.and_then(|result| result);
match matches {
Ok(matches) => {
let hits = matches
.into_iter()
.filter_map(|found| {
let track = found.track;
let hit = SimilarityHit {
score: found.score,
item_id: super::audio::hex_encode(
ItemId::derive(
&own,
ItemKind::Track,
&format!("track:{}", track.id),
)
.as_bytes(),
),
title: track.title,
artist_names: track
.artists
.into_iter()
.map(|artist| artist.name)
.collect(),
featured_artist_names: track
.featured_artists
.into_iter()
.map(|artist| artist.name)
.collect(),
year: track.release_year,
duration_seconds: Some(track.duration_seconds.round() as i64),
content_id: track.content_id,
release_title: Some(track.release_title),
track_number: track.track_number,
disc_number: track.disc_number,
embedding_signature: Some(found.embedding_signature),
};
match hit.validate() {
Ok(()) => Some(hit),
Err(err) => {
tracing::debug!(%err, "invalid local similarity metadata skipped");
None
}
}
})
.collect();
SimilarityResponse::success(hits)?
}
Err(err) => {
SimilarityResponse::refused(format!("similarity query is unavailable: {err:#}"))?
}
}
};
wire::write_response(&mut stream, &response).await?;
stream.send.finish()?;
let _ = stream.send.stopped().await;
super::record_stream_transport(&transport, "similarity", "inbound", "done", &stream);
Ok(())
}
pub async fn search(
service: Arc<MusicDhtService>,
query: QueryVector,
limit: usize,
transport: Arc<TransportStats>,
) -> Result<FedSearchResults> {
let own = service.endpoint_id();
let mut peers = Vec::new();
let mut seen = HashSet::new();
for peer in service
.connected_peers()
.into_iter()
.chain(service.known_peers().into_iter().map(|peer| peer.peer_id))
{
if peer != own && seen.insert(peer) {
peers.push(peer);
}
if peers.len() >= MAX_QUERY_PEERS {
break;
}
}
let query_signature = wire::embedding_signature(&query.vector)?;
let request = Arc::new(SimilarityRequest::new(
query.profile_id,
query.vector,
limit.clamp(1, wire::MAX_SIMILARITY_RESULTS),
)?);
let responses = stream::iter(peers.into_iter().map(|peer| {
let service = Arc::clone(&service);
let request = Arc::clone(&request);
let transport = Arc::clone(&transport);
async move {
tokio::time::timeout(
QUERY_TIMEOUT,
query_peer(service, peer, &request, transport),
)
.await
.map_err(|_| anyhow::anyhow!("similarity peer timed out"))?
}
}))
.buffer_unordered(QUERY_CONCURRENCY)
.collect::<Vec<_>>()
.await;
let mut hits = Vec::new();
for response in responses {
match response {
Ok(peer_hits) => hits.extend(peer_hits),
Err(err) => tracing::debug!(%err, "similarity peer query skipped"),
}
}
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 {
if query
.source_content_id
.as_deref()
.is_some_and(|source| track.content_id.as_deref() == Some(source))
{
continue;
}
let key = track
.content_id
.clone()
.unwrap_or_else(|| format!("{}:{}", track.owner, track.item_id));
if !dedup.insert(key) {
continue;
}
if embedding_signature.is_some_and(|candidate| {
embedding_signatures.iter().any(|existing| {
wire::signature_distance(&candidate, existing)
<= MAX_NEAR_DUPLICATE_SIGNATURE_DISTANCE
})
}) {
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);
if tracks.len() >= limit.min(wire::MAX_SIMILARITY_RESULTS) {
break;
}
}
Ok(FedSearchResults {
artists: Vec::new(),
tracks,
})
}
async fn query_peer(
service: Arc<MusicDhtService>,
owner: EndpointId,
request: &SimilarityRequest,
transport: Arc<TransportStats>,
) -> Result<
Vec<(
FedTrack,
f32,
Option<[u8; wire::SIMILARITY_SIGNATURE_BYTES]>,
)>,
> {
let mut stream = service
.open_stream(owner, SIMILARITY_ALPN)
.await
.map_err(|err| anyhow::anyhow!("cannot reach similarity peer: {err}"))?;
super::record_stream_transport(&transport, "similarity", "outbound", "open", &stream);
let response = wire::exchange(&mut stream, request).await?;
super::record_stream_transport(&transport, "similarity", "outbound", "done", &stream);
anyhow::ensure!(
response.ok,
"peer refused similarity query: {}",
response.error.unwrap_or_default()
);
Ok(response
.hits
.into_iter()
.map(|hit| {
let score = hit.score;
let embedding_signature = hit.embedding_signature;
(
FedTrack {
item_id: hit.item_id,
owner: owner.to_string(),
own: false,
title: hit.title,
artist_names: hit.artist_names,
featured_artist_names: hit.featured_artist_names,
year: hit.year,
duration_seconds: hit.duration_seconds,
content_id: hit.content_id,
release_title: hit.release_title,
track_number: hit.track_number,
disc_number: hit.disc_number,
},
score,
embedding_signature,
)
})
.collect())
}