feat(suggestions): improve vector store and add cache management
- Add get_all_related_artists for direct relationship lookup - Add get_related_artists helper method - Add clear_all method to vector store for cache management - Add clear_vector_store Tauri command - Add Artist Vectors cache section to Settings view - Use artist_info tuples (mbid, name) for better track searching
This commit is contained in:
1 parent
9a89074b5f
commit
79c68d292e
4 files changed
+177
-7
No files matched your search
@@ -390,6 +390,94 @@ impl ArtistVectorStore {
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Get related artists for a given artist (from their vector entries)
|
||||
///
|
||||
/// This returns artists that the given artist is connected to via
|
||||
/// MusicBrainz relationships (members, collaborators, etc.)
|
||||
pub fn get_related_artists(&self, mbid: &str) -> Result<Vec<SimilarArtist>, String> {
|
||||
let Some(artist_idx) = self.get_idx(mbid) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut stmt = self
|
||||
.conn
|
||||
.prepare(
|
||||
"SELECT target_idx, SUM(weight) as total_weight
|
||||
FROM vector_entries
|
||||
WHERE artist_idx = ?1
|
||||
GROUP BY target_idx
|
||||
ORDER BY total_weight DESC",
|
||||
)
|
||||
.map_err(|e| format!("Failed to prepare query: {}", e))?;
|
||||
|
||||
let rows = stmt
|
||||
.query_map(params![artist_idx], |row| {
|
||||
Ok((row.get::<_, u32>(0)?, row.get::<_, f32>(1)?))
|
||||
})
|
||||
.map_err(|e| format!("Failed to query relations: {}", e))?;
|
||||
|
||||
let mut results = Vec::new();
|
||||
for row in rows.flatten() {
|
||||
let (target_idx, weight) = row;
|
||||
if let Some(target_mbid) = self.get_mbid(target_idx) {
|
||||
let name = self.get_artist_name(target_mbid);
|
||||
results.push(SimilarArtist {
|
||||
mbid: target_mbid.to_string(),
|
||||
name,
|
||||
similarity: weight,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Get all related artists for multiple source artists, excluding specified MBIDs
|
||||
///
|
||||
/// Returns a deduplicated list sorted by total weight across all sources
|
||||
pub fn get_all_related_artists(
|
||||
&self,
|
||||
source_mbids: &[String],
|
||||
exclude_mbids: &[String],
|
||||
limit: usize,
|
||||
) -> Result<Vec<SimilarArtist>, String> {
|
||||
use std::collections::HashMap;
|
||||
|
||||
let exclude_set: std::collections::HashSet<_> = exclude_mbids.iter().collect();
|
||||
let mut artist_weights: HashMap<String, (Option<String>, f32)> = HashMap::new();
|
||||
|
||||
for mbid in source_mbids {
|
||||
let related = self.get_related_artists(mbid)?;
|
||||
for artist in related {
|
||||
if exclude_set.contains(&artist.mbid) {
|
||||
continue;
|
||||
}
|
||||
let entry = artist_weights
|
||||
.entry(artist.mbid.clone())
|
||||
.or_insert((artist.name.clone(), 0.0));
|
||||
entry.1 += artist.similarity;
|
||||
}
|
||||
}
|
||||
|
||||
let mut results: Vec<_> = artist_weights
|
||||
.into_iter()
|
||||
.map(|(mbid, (name, weight))| SimilarArtist {
|
||||
mbid,
|
||||
name,
|
||||
similarity: weight,
|
||||
})
|
||||
.collect();
|
||||
|
||||
results.sort_by(|a, b| {
|
||||
b.similarity
|
||||
.partial_cmp(&a.similarity)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
results.truncate(limit);
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Get statistics about the store
|
||||
pub fn get_stats(&self) -> Result<StoreStats, String> {
|
||||
let artist_count = self.artist_to_idx.len();
|
||||
@@ -435,6 +523,30 @@ impl ArtistVectorStore {
|
||||
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
/// Clear all data from the store
|
||||
pub fn clear_all(&mut self) -> Result<usize, String> {
|
||||
let deleted = self
|
||||
.conn
|
||||
.execute("DELETE FROM vector_entries", [])
|
||||
.map_err(|e| format!("Failed to delete vector entries: {}", e))?;
|
||||
|
||||
self.conn
|
||||
.execute("DELETE FROM vector_metadata", [])
|
||||
.map_err(|e| format!("Failed to delete metadata: {}", e))?;
|
||||
|
||||
self.conn
|
||||
.execute("DELETE FROM artist_index", [])
|
||||
.map_err(|e| format!("Failed to delete artist index: {}", e))?;
|
||||
|
||||
// Reset in-memory state
|
||||
self.artist_to_idx.clear();
|
||||
self.idx_to_artist.clear();
|
||||
self.next_idx = 0;
|
||||
|
||||
log::info!("Artist vector store cleared: {} entries deleted", deleted);
|
||||
Ok(deleted)
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about the vector store
|
||||
|
||||
@@ -59,8 +59,9 @@ pub async fn get_playlist_suggestions_v2(
|
||||
let skip_network = config.skip_vector_build;
|
||||
|
||||
// Resolve artist names to MBIDs (with caching)
|
||||
// Store (mbid, name) tuples so we can search for tracks by name
|
||||
let resolve_start = Instant::now();
|
||||
let mut artist_mbids = Vec::new();
|
||||
let mut artist_info: Vec<(String, String)> = Vec::new(); // (mbid, name)
|
||||
let mut cached_count = 0;
|
||||
let mut resolved_count = 0;
|
||||
let mut skipped_count = 0;
|
||||
@@ -77,7 +78,7 @@ pub async fn get_playlist_suggestions_v2(
|
||||
if resolved.confidence != MatchConfidence::None
|
||||
&& resolved.confidence != MatchConfidence::Low
|
||||
{
|
||||
artist_mbids.push(mbid);
|
||||
artist_info.push((mbid, artist.name.clone()));
|
||||
cached_count += 1;
|
||||
continue;
|
||||
}
|
||||
@@ -117,7 +118,7 @@ pub async fn get_playlist_suggestions_v2(
|
||||
let cache = mb_state.cache.lock().await;
|
||||
let _ = cache.set_artist(&artist.name, &resolved);
|
||||
}
|
||||
artist_mbids.push(mbid);
|
||||
artist_info.push((mbid, artist.name.clone()));
|
||||
resolved_count += 1;
|
||||
continue;
|
||||
}
|
||||
@@ -148,7 +149,7 @@ pub async fn get_playlist_suggestions_v2(
|
||||
skip_network
|
||||
);
|
||||
|
||||
if artist_mbids.is_empty() {
|
||||
if artist_info.is_empty() {
|
||||
log::warn!("[Suggestions] No MBIDs resolved, returning empty result");
|
||||
return Ok(SuggestionResult {
|
||||
tracks: Vec::new(),
|
||||
@@ -182,11 +183,11 @@ pub async fn get_playlist_suggestions_v2(
|
||||
let exclude_set: HashSet<u64> = input.exclude_track_ids.into_iter().collect();
|
||||
|
||||
// Generate suggestions
|
||||
log::info!("[Suggestions] Calling engine.generate_suggestions with {} MBIDs", artist_mbids.len());
|
||||
log::info!("[Suggestions] Calling engine.generate_suggestions with {} artists", artist_info.len());
|
||||
let engine_start = std::time::Instant::now();
|
||||
|
||||
let result = engine
|
||||
.generate_suggestions(&artist_mbids, &exclude_set, input.include_reasons)
|
||||
.generate_suggestions(&artist_info, &exclude_set, input.include_reasons)
|
||||
.await;
|
||||
|
||||
let engine_elapsed = engine_start.elapsed();
|
||||
@@ -229,3 +230,12 @@ pub async fn cleanup_vector_store(
|
||||
let mut store = store_state.store.lock().await;
|
||||
store.cleanup_expired(max_age_secs)
|
||||
}
|
||||
|
||||
/// Clear all vector store data
|
||||
#[tauri::command]
|
||||
pub async fn clear_vector_store(
|
||||
store_state: State<'_, ArtistVectorStoreState>,
|
||||
) -> Result<usize, String> {
|
||||
let mut store = store_state.store.lock().await;
|
||||
store.clear_all()
|
||||
}
|
||||
@@ -536,6 +536,7 @@ pub fn run() {
|
||||
commands::get_playlist_suggestions_v2,
|
||||
commands::get_vector_store_stats,
|
||||
commands::cleanup_vector_store,
|
||||
commands::clear_vector_store,
|
||||
// Favorites commands
|
||||
commands::get_favorites,
|
||||
commands::add_favorite,
|
||||
|
||||
Reference in new issue
Block a user