39 lines
1.0 KiB
Go
39 lines
1.0 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestQueryNotesRanksByCosine(t *testing.T) {
|
|
ctx := context.Background()
|
|
st := newTestStore(t)
|
|
|
|
now := time.Now()
|
|
// 3-dim vectors along distinct axes; query aligns with the "backups" note.
|
|
if _, err := st.WriteNote(ctx, now, "prefer backups at 3am", []float32{1, 0, 0}, "tap:voice"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.WriteNote(ctx, now, "gpu driver fixed the flicker", []float32{0, 1, 0}, "tap:voice"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.WriteNote(ctx, now, "cat likes the window", []float32{0, 0, 1}, "tap:voice"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
got, err := st.QueryNotes(ctx, []float32{0.9, 0.1, 0}, 2)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(got) != 2 {
|
|
t.Fatalf("want 2 notes, got %d", len(got))
|
|
}
|
|
if got[0].Text != "prefer backups at 3am" {
|
|
t.Errorf("nearest = %q, want backups note (score %.3f)", got[0].Text, got[0].Score)
|
|
}
|
|
if got[0].Score <= got[1].Score {
|
|
t.Errorf("scores not descending: %.3f then %.3f", got[0].Score, got[1].Score)
|
|
}
|
|
}
|