-
Notifications
You must be signed in to change notification settings - Fork 0
/
sqlite_test.go
80 lines (61 loc) · 1.5 KB
/
sqlite_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package vector
import (
"context"
"log/slog"
"net/url"
"strconv"
"testing"
)
func TestSQLiteDatabase(t *testing.T) {
slog.SetLogLoggerLevel(slog.LevelDebug)
ctx := context.Background()
type settings struct {
compression string
expected_results []int
max_distance float32
}
tests := []settings{
settings{
compression: "none",
expected_results: []int{1, 1, 0, 0},
max_distance: 4,
},
settings{
compression: "quantize",
expected_results: []int{1, 0, 0, 0},
max_distance: 1,
},
settings{
compression: "matroyshka",
expected_results: []int{1, 1, 0, 0},
max_distance: 0.5,
},
}
for _, settings := range tests {
expected_results := settings.expected_results
q := url.Values{}
q.Set("dsn", "{tmp}.db")
q.Set("embedder-uri", "ollama://?model=mxbai-embed-large")
q.Set("dimensions", "1024")
q.Set("max-results", "10")
q.Set("max-distance", strconv.FormatFloat(float64(settings.max_distance), 'g', -1, 32))
q.Set("compression", settings.compression)
u := url.URL{}
u.Scheme = "sqlite"
u.RawQuery = q.Encode()
db_uri := u.String()
slog.Debug("Create new vector database", "uri", db_uri)
db, err := NewDatabase(ctx, db_uri)
if err != nil {
t.Fatalf("Failed to create database for '%s', %v", db_uri, err)
}
err = testDatabaseWithLocations(ctx, db, expected_results)
if err != nil {
t.Fatal(err)
}
err = db.Close(ctx)
if err != nil {
t.Fatalf("Failed to close database, %v", err)
}
}
}