75 lines
1.7 KiB
Go
75 lines
1.7 KiB
Go
|
package index
|
||
|
|
|
||
|
|
import (
|
||
|
|
"math"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestInMemoryVectorIndex_AddSearch(t *testing.T) {
|
||
|
|
idx := NewInMemoryVectorIndex()
|
||
|
|
|
||
|
|
idx.Add("a", []float64{1, 0, 0})
|
||
|
|
idx.Add("b", []float64{0, 1, 0})
|
||
|
|
idx.Add("c", []float64{0, 0, 1})
|
||
|
|
|
||
|
|
if idx.Len() != 3 {
|
||
|
|
t.Fatalf("Len = %d, want 3", idx.Len())
|
||
|
|
}
|
||
|
|
|
||
|
|
hits := idx.Search([]float64{1, 0, 0}, 3)
|
||
|
|
if len(hits) != 3 {
|
||
|
|
t.Fatalf("Search: got %d hits, want 3", len(hits))
|
||
|
|
}
|
||
|
|
if hits[0].ID != "a" {
|
||
|
|
t.Errorf("top hit = %q, want %q", hits[0].ID, "a")
|
||
|
|
}
|
||
|
|
if math.Abs(hits[0].Score-1.0) > 1e-9 {
|
||
|
|
t.Errorf("top score = %f, want 1.0", hits[0].Score)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestInMemoryVectorIndex_Remove(t *testing.T) {
|
||
|
|
idx := NewInMemoryVectorIndex()
|
||
|
|
idx.Add("a", []float64{1, 0})
|
||
|
|
idx.Add("b", []float64{0, 1})
|
||
|
|
|
||
|
|
idx.Remove("a")
|
||
|
|
if idx.Len() != 1 {
|
||
|
|
t.Fatalf("Len after remove = %d, want 1", idx.Len())
|
||
|
|
}
|
||
|
|
|
||
|
|
hits := idx.Search([]float64{0, 1}, 5)
|
||
|
|
if len(hits) != 1 || hits[0].ID != "b" {
|
||
|
|
t.Errorf("expected only 'b' in results")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestInMemoryVectorIndex_RemoveByPrefix(t *testing.T) {
|
||
|
|
idx := NewInMemoryVectorIndex()
|
||
|
|
idx.Add("note1_0", []float64{1, 0})
|
||
|
|
idx.Add("note1_1", []float64{0.9, 0.1})
|
||
|
|
idx.Add("note2_0", []float64{0, 1})
|
||
|
|
|
||
|
|
idx.RemoveByPrefix("note1_")
|
||
|
|
if idx.Len() != 1 {
|
||
|
|
t.Fatalf("Len after RemoveByPrefix = %d, want 1", idx.Len())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestCosine(t *testing.T) {
|
||
|
|
tests := []struct {
|
||
|
|
a, b []float64
|
||
|
|
want float64
|
||
|
|
}{
|
||
|
|
{[]float64{1, 0}, []float64{1, 0}, 1.0},
|
||
|
|
{[]float64{1, 0}, []float64{0, 1}, 0.0},
|
||
|
|
{[]float64{1, 0}, []float64{-1, 0}, -1.0},
|
||
|
|
}
|
||
|
|
for _, tt := range tests {
|
||
|
|
got := cosine(tt.a, tt.b)
|
||
|
|
if math.Abs(got-tt.want) > 1e-9 {
|
||
|
|
t.Errorf("cosine(%v, %v) = %f, want %f", tt.a, tt.b, got, tt.want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|