feat(recommendation): add LoadCandidatesFromSimilarity (5-source candidate pool)
Implements M4c's similarity-driven candidate pool: CandidateSourceLimits, DefaultCandidateSourceLimits, and LoadCandidatesFromSimilarity consuming LoadRadioCandidatesV2 (5-way UNION + max-score dedup). Adds 10 integration tests covering all sources, exclusions, dedup, and edge cases. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -58,6 +58,95 @@ func LoadCandidates(
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// CandidateSourceLimits controls per-source K values for the M4c
|
||||
// similarity-driven pool. Defaults via DefaultCandidateSourceLimits().
|
||||
type CandidateSourceLimits struct {
|
||||
LBSimilar int
|
||||
SimilarArtist int
|
||||
TagOverlap int
|
||||
LikesOverlap int
|
||||
RandomFill int
|
||||
}
|
||||
|
||||
// DefaultCandidateSourceLimits returns the v1 hardcoded constants per spec.
|
||||
func DefaultCandidateSourceLimits() CandidateSourceLimits {
|
||||
return CandidateSourceLimits{
|
||||
LBSimilar: 30,
|
||||
SimilarArtist: 30,
|
||||
TagOverlap: 20,
|
||||
LikesOverlap: 20,
|
||||
RandomFill: 30,
|
||||
}
|
||||
}
|
||||
|
||||
// LoadCandidatesFromSimilarity is M4c's primary candidate-pool loader.
|
||||
// 5-way SQL UNION (LB-similar / similar-artist tracks / MB-tag overlap /
|
||||
// likes-overlap / random fill) + dedup-by-max sim_score. Returns
|
||||
// []Candidate (same shape as LoadCandidates) so Shuffle is unchanged.
|
||||
//
|
||||
// Caller (radio handler) falls back to LoadCandidates on error.
|
||||
func LoadCandidatesFromSimilarity(
|
||||
ctx context.Context,
|
||||
q *dbq.Queries,
|
||||
userID, seedID pgtype.UUID,
|
||||
recentlyPlayedHours int,
|
||||
currentVector SessionVector,
|
||||
exclude []pgtype.UUID,
|
||||
limits CandidateSourceLimits,
|
||||
) ([]Candidate, error) {
|
||||
if exclude == nil {
|
||||
exclude = []pgtype.UUID{}
|
||||
}
|
||||
rows, err := q.LoadRadioCandidatesV2(ctx, dbq.LoadRadioCandidatesV2Params{
|
||||
UserID: userID,
|
||||
ID: seedID,
|
||||
Column3: int64(recentlyPlayedHours),
|
||||
Column4: exclude,
|
||||
Limit: int32(limits.LBSimilar),
|
||||
Limit_2: int32(limits.SimilarArtist),
|
||||
Limit_3: int32(limits.TagOverlap),
|
||||
Limit_4: int32(limits.LikesOverlap),
|
||||
Limit_5: int32(limits.RandomFill),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
likes, err := loadContextualLikesByTrack(ctx, q, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := make([]Candidate, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
var lpt *time.Time
|
||||
if r.LastPlayedAt.Valid {
|
||||
t := r.LastPlayedAt.Time
|
||||
lpt = &t
|
||||
}
|
||||
// sqlc returns SimilarityScore as interface{} (couldn't infer the
|
||||
// type through max(...) over a UNION). Type-assert; default to 0
|
||||
// on the (impossible-but-defensive) case where it's nil/wrong type.
|
||||
var simScore float64
|
||||
if v, ok := r.SimilarityScore.(float64); ok {
|
||||
simScore = v
|
||||
}
|
||||
ctxScore := ContextualMatchScore(currentVector, likes[r.Track.ID], DefaultSimilarityWeights)
|
||||
out = append(out, Candidate{
|
||||
Track: r.Track,
|
||||
Inputs: ScoringInputs{
|
||||
IsGeneralLiked: r.IsLiked,
|
||||
LastPlayedAt: lpt,
|
||||
PlayCount: int(r.PlayCount),
|
||||
SkipCount: int(r.SkipCount),
|
||||
ContextualMatchScore: ctxScore,
|
||||
SimilarityScore: simScore,
|
||||
},
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// loadContextualLikesByTrack fetches the user's active contextual_likes in
|
||||
// one query and groups them by track_id. Rows whose session_vector fails
|
||||
// to unmarshal are skipped with no error (don't poison scoring over one
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
package recommendation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
|
||||
"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
|
||||
)
|
||||
|
||||
// helperLBSimilarity inserts a track_similarity row.
|
||||
func helperLBSimilarity(t *testing.T, f fixture, a, b pgtype.UUID, score float64) {
|
||||
t.Helper()
|
||||
if _, err := f.pool.Exec(context.Background(),
|
||||
`INSERT INTO track_similarity (track_a_id, track_b_id, score, source) VALUES ($1, $2, $3, 'listenbrainz')`,
|
||||
a, b, score); err != nil {
|
||||
t.Fatalf("insert track_similarity: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// helperArtistSimilarity inserts an artist_similarity row.
|
||||
func helperArtistSimilarity(t *testing.T, f fixture, a, b pgtype.UUID, score float64) {
|
||||
t.Helper()
|
||||
if _, err := f.pool.Exec(context.Background(),
|
||||
`INSERT INTO artist_similarity (artist_a_id, artist_b_id, score, source) VALUES ($1, $2, $3, 'listenbrainz')`,
|
||||
a, b, score); err != nil {
|
||||
t.Fatalf("insert artist_similarity: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// helperSetTrackGenre updates a track's genre column. Used to retrofit
|
||||
// genres onto the fixture's auto-created tracks (fixture creates tracks
|
||||
// with NULL genre).
|
||||
func helperSetTrackGenre(t *testing.T, f fixture, trackID pgtype.UUID, genre string) {
|
||||
t.Helper()
|
||||
if _, err := f.pool.Exec(context.Background(),
|
||||
`UPDATE tracks SET genre = $1 WHERE id = $2`, genre, trackID); err != nil {
|
||||
t.Fatalf("set genre: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func defaultLimits() CandidateSourceLimits {
|
||||
return DefaultCandidateSourceLimits()
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_LBSimilarSourceContributes(t *testing.T) {
|
||||
f := newFixture(t, 5)
|
||||
seed := f.tracks[0]
|
||||
target := f.tracks[1]
|
||||
helperLBSimilarity(t, f, seed.ID, target.ID, 0.85)
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
var found *Candidate
|
||||
for i := range got {
|
||||
if got[i].Track.ID == target.ID {
|
||||
found = &got[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if found == nil {
|
||||
t.Fatal("LB-similar target missing from candidates")
|
||||
}
|
||||
if found.Inputs.SimilarityScore < 0.84 || found.Inputs.SimilarityScore > 0.86 {
|
||||
t.Errorf("LB-similar SimilarityScore = %v, want ~0.85", found.Inputs.SimilarityScore)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_SimilarArtistTracksContribute(t *testing.T) {
|
||||
f := newFixture(t, 1) // creates 1 artist + 1 album + 1 track (the seed)
|
||||
seed := f.tracks[0]
|
||||
// seed_info CTE only emits artist_id rows when genre is non-NULL (it
|
||||
// uses a LATERAL regexp_split_to_table join filtered by trim(g) <> '').
|
||||
// Give the seed a genre so that seed_info is populated and the
|
||||
// similar_artists JOIN can find artist_a_id = seed.ArtistID.
|
||||
helperSetTrackGenre(t, f, seed.ID, "Rock")
|
||||
// Add a SECOND artist + track in that artist; relate the two artists via artist_similarity.
|
||||
otherArtist, _ := f.q.UpsertArtist(context.Background(), dbq.UpsertArtistParams{Name: "OtherArtist", SortName: "OtherArtist"})
|
||||
otherAlbum, _ := f.q.UpsertAlbum(context.Background(), dbq.UpsertAlbumParams{Title: "OtherAlbum", SortTitle: "OtherAlbum", ArtistID: otherArtist.ID})
|
||||
otherTrack, _ := f.q.UpsertTrack(context.Background(), dbq.UpsertTrackParams{
|
||||
Title: "OtherTrack", AlbumID: otherAlbum.ID, ArtistID: otherArtist.ID,
|
||||
FilePath: "/tmp/other.flac", DurationMs: 180_000,
|
||||
})
|
||||
helperArtistSimilarity(t, f, seed.ArtistID, otherArtist.ID, 0.8)
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == otherTrack.ID {
|
||||
// 0.8 × 0.5 = 0.4
|
||||
if c.Inputs.SimilarityScore < 0.39 || c.Inputs.SimilarityScore > 0.41 {
|
||||
t.Errorf("similar-artist SimilarityScore = %v, want ~0.4 (0.8 × 0.5)", c.Inputs.SimilarityScore)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("similar-artist track missing from candidates")
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_TagOverlapContributes(t *testing.T) {
|
||||
f := newFixture(t, 2)
|
||||
seed := f.tracks[0]
|
||||
target := f.tracks[1]
|
||||
helperSetTrackGenre(t, f, seed.ID, "Rock; Pop")
|
||||
helperSetTrackGenre(t, f, target.ID, "Rock")
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == target.ID {
|
||||
// Seed has 2 tags; target shares 1 → jaccard 1/2 = 0.5.
|
||||
if c.Inputs.SimilarityScore < 0.49 || c.Inputs.SimilarityScore > 0.51 {
|
||||
t.Errorf("tag-overlap SimilarityScore = %v, want ~0.5", c.Inputs.SimilarityScore)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("tag-overlap target missing from candidates")
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_LikesOverlapContributes(t *testing.T) {
|
||||
f := newFixture(t, 2)
|
||||
seed := f.tracks[0]
|
||||
liked := f.tracks[1]
|
||||
helperSetTrackGenre(t, f, seed.ID, "Rock")
|
||||
helperSetTrackGenre(t, f, liked.ID, "Rock")
|
||||
if _, err := f.q.LikeTrack(context.Background(), dbq.LikeTrackParams{UserID: f.user, TrackID: liked.ID}); err != nil {
|
||||
t.Fatalf("like: %v", err)
|
||||
}
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == liked.ID {
|
||||
// Both tracks tagged "Rock" → jaccard 1/1 = 1.0 from tag-overlap.
|
||||
// likes-overlap = 0.6. Max wins = 1.0.
|
||||
if c.Inputs.SimilarityScore < 0.59 {
|
||||
t.Errorf("likes-overlap candidate SimilarityScore = %v, want ≥ 0.6", c.Inputs.SimilarityScore)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("liked track with shared tag missing from candidates")
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_RandomFillReturnsTracks(t *testing.T) {
|
||||
f := newFixture(t, 10) // 10 tracks; no similarity data
|
||||
seed := f.tracks[0]
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if len(got) == 0 {
|
||||
t.Error("random fill returned 0 candidates; expected at least some")
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Inputs.SimilarityScore != 0 {
|
||||
t.Errorf("random-fill track %s has SimilarityScore = %v, want 0", c.Track.Title, c.Inputs.SimilarityScore)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_ExcludeListRespected(t *testing.T) {
|
||||
f := newFixture(t, 5)
|
||||
seed := f.tracks[0]
|
||||
excluded := f.tracks[1].ID
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true},
|
||||
[]pgtype.UUID{excluded}, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == excluded {
|
||||
t.Error("excluded track appeared in candidates")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_SeedAlwaysExcluded(t *testing.T) {
|
||||
f := newFixture(t, 5)
|
||||
seed := f.tracks[0]
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == seed.ID {
|
||||
t.Error("seed track appeared in candidates")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_RecentlyPlayedExcluded(t *testing.T) {
|
||||
f := newFixture(t, 5)
|
||||
seed := f.tracks[0]
|
||||
recent := f.tracks[1].ID
|
||||
var sessionID pgtype.UUID
|
||||
if err := f.pool.QueryRow(context.Background(),
|
||||
`INSERT INTO play_sessions (user_id, started_at, last_event_at, client_id)
|
||||
VALUES ($1, now() - interval '5 minutes', now(), 'test') RETURNING id`,
|
||||
f.user).Scan(&sessionID); err != nil {
|
||||
t.Fatalf("session: %v", err)
|
||||
}
|
||||
if _, err := f.pool.Exec(context.Background(),
|
||||
`INSERT INTO play_events (user_id, track_id, session_id, started_at, ended_at, duration_played_ms, completion_ratio, was_skipped)
|
||||
VALUES ($1, $2, $3, now() - interval '30 minutes', now() - interval '20 minutes', 200000, 0.9, false)`,
|
||||
f.user, recent, sessionID); err != nil {
|
||||
t.Fatalf("play_event: %v", err)
|
||||
}
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
for _, c := range got {
|
||||
if c.Track.ID == recent {
|
||||
t.Error("recently-played track appeared in candidates")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_DedupTakesMaxScore(t *testing.T) {
|
||||
f := newFixture(t, 2)
|
||||
seed := f.tracks[0]
|
||||
target := f.tracks[1]
|
||||
helperSetTrackGenre(t, f, seed.ID, "Rock")
|
||||
helperSetTrackGenre(t, f, target.ID, "Rock") // jaccard 1/1 = 1.0 from tag-overlap
|
||||
helperLBSimilarity(t, f, seed.ID, target.ID, 0.5) // weaker LB signal
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
count := 0
|
||||
for _, c := range got {
|
||||
if c.Track.ID == target.ID {
|
||||
count++
|
||||
// tag-overlap (1.0) wins over LB (0.5) per max() — expect ≥ 0.99
|
||||
if c.Inputs.SimilarityScore < 0.99 {
|
||||
t.Errorf("dedup max SimilarityScore = %v, want ≥ 0.99 (tag-overlap should win)", c.Inputs.SimilarityScore)
|
||||
}
|
||||
}
|
||||
}
|
||||
if count != 1 {
|
||||
t.Errorf("target appeared %d times, want 1 (dedup failed)", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCandidatesFromSimilarity_EmptyLibrary_NoError(t *testing.T) {
|
||||
f := newFixture(t, 1) // just the seed
|
||||
seed := f.tracks[0]
|
||||
got, err := LoadCandidatesFromSimilarity(
|
||||
context.Background(), f.q, f.user, seed.ID, 1, SessionVector{Seed: true}, nil, defaultLimits(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
// Only the seed exists; it's excluded → 0 candidates.
|
||||
if len(got) != 0 {
|
||||
t.Errorf("got %d candidates from seed-only library, want 0", len(got))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user