sqlite-vec/examples/simple-go-cgo/demo.go

87 lines
1.6 KiB
Go
Raw Normal View History

2024-05-10 20:51:42 -07:00
package main
import (
"database/sql"
"fmt"
"log"
2024-07-16 23:11:26 -07:00
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
2024-05-10 20:51:42 -07:00
_ "github.com/mattn/go-sqlite3"
)
func main() {
sqlite_vec.Auto()
db, err := sql.Open("sqlite3", ":memory:")
if err != nil {
log.Fatal(err)
}
defer db.Close()
var sqliteVersion string
var vecVersion string
err = db.QueryRow("select sqlite_version(), vec_version()").Scan(&sqliteVersion, &vecVersion)
if err != nil {
log.Fatal(err)
}
fmt.Printf("sqlite_version=%s, vec_version=%s\n", sqliteVersion, vecVersion)
2024-05-11 00:21:44 -07:00
_, err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
2024-05-10 20:51:42 -07:00
if err != nil {
log.Fatal(err)
}
items := map[int][]float32{
2024-05-11 00:21:44 -07:00
1: {0.1, 0.1, 0.1, 0.1},
2: {0.2, 0.2, 0.2, 0.2},
3: {0.3, 0.3, 0.3, 0.3},
4: {0.4, 0.4, 0.4, 0.4},
5: {0.5, 0.5, 0.5, 0.5},
2024-05-10 20:51:42 -07:00
}
2024-05-11 00:21:44 -07:00
q := []float32{0.3, 0.3, 0.3, 0.3}
2024-05-10 20:51:42 -07:00
for id, values := range items {
2024-07-31 12:55:20 -07:00
v, err := sqlite_vec.SerializeFloat32(values)
2024-05-10 20:51:42 -07:00
if err != nil {
log.Fatal(err)
}
_, err = db.Exec("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)", id, v)
if err != nil {
log.Fatal(err)
}
}
2024-07-31 12:55:20 -07:00
query, err := sqlite_vec.SerializeFloat32(q)
2024-05-10 20:51:42 -07:00
if err != nil {
log.Fatal(err)
}
rows, err := db.Query(`
SELECT
rowid,
distance
FROM vec_items
WHERE embedding MATCH ?
ORDER BY distance
2024-05-11 00:21:44 -07:00
LIMIT 3
2024-05-10 20:51:42 -07:00
`, query)
if err != nil {
log.Fatal(err)
}
for rows.Next() {
var rowid int64
var distance float64
err = rows.Scan(&rowid, &distance)
if err != nil {
log.Fatal(err)
}
fmt.Printf("rowid=%d, distance=%f\n", rowid, distance)
}
err = rows.Err()
if err != nil {
log.Fatal((err))
}
}