mirror of https://gitee.com/godoos/godoos.git
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
1.7 KiB
100 lines
1.7 KiB
package main
|
|
|
|
import (
|
|
_ "embed"
|
|
"log"
|
|
|
|
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/ncruces"
|
|
"github.com/ncruces/go-sqlite3"
|
|
)
|
|
|
|
var Db *sqlite3.Conn
|
|
|
|
func main() {
|
|
db, err := sqlite3.Open(":memory:")
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
stmt, _, err := db.Prepare(`SELECT sqlite_version(), vec_version()`)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
stmt.Step()
|
|
|
|
log.Printf("sqlite_version=%s, vec_version=%s\n", stmt.ColumnText(0), stmt.ColumnText(1))
|
|
stmt.Close()
|
|
|
|
err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
items := map[int][]float32{
|
|
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},
|
|
}
|
|
q := []float32{0.3, 0.3, 0.3, 0.3}
|
|
|
|
stmt, _, err = db.Prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
for id, values := range items {
|
|
v, err := sqlite_vec.SerializeFloat32(values)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
stmt.BindInt(1, id)
|
|
stmt.BindBlob(2, v)
|
|
err = stmt.Exec()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
stmt.Reset()
|
|
}
|
|
stmt.Close()
|
|
|
|
stmt, _, err = db.Prepare(`
|
|
SELECT
|
|
rowid,
|
|
distance
|
|
FROM vec_items
|
|
WHERE embedding MATCH ?
|
|
ORDER BY distance
|
|
LIMIT 3
|
|
`)
|
|
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
query, err := sqlite_vec.SerializeFloat32(q)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
stmt.BindBlob(1, query)
|
|
|
|
for stmt.Step() {
|
|
rowid := stmt.ColumnInt64(0)
|
|
distance := stmt.ColumnFloat(1)
|
|
log.Printf("rowid=%d, distance=%f\n", rowid, distance)
|
|
}
|
|
if err := stmt.Err(); err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
err = stmt.Close()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
err = db.Close()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
|