sqlite-vec/examples/simple-node/demo.mjs

53 lines
1.1 KiB
JavaScript
Raw Normal View History

2024-05-10 20:51:42 -07:00
import * as sqliteVec from "sqlite-vec";
import Database from "better-sqlite3";
const db = new Database(":memory:");
2024-05-10 22:05:49 -07:00
sqliteVec.load(db);
2024-05-10 20:51:42 -07:00
const { sqlite_version, vec_version } = db
.prepare(
"select sqlite_version() as sqlite_version, vec_version() as vec_version;"
)
.get();
console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`);
2024-05-11 00:21:44 -07:00
const items = [
[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]],
];
const query = [0.3, 0.3, 0.3, 0.3];
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])");
2024-05-10 20:51:42 -07:00
const insertStmt = db.prepare(
2024-05-11 00:21:44 -07:00
"INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)"
2024-05-10 20:51:42 -07:00
);
const insertVectors = db.transaction((items) => {
for (const [id, vector] of items) {
2024-05-11 00:21:44 -07:00
insertStmt.run(BigInt(id), new Float32Array(vector));
2024-05-10 20:51:42 -07:00
}
});
2024-05-11 00:21:44 -07:00
insertVectors(items);
2024-05-10 20:51:42 -07:00
const rows = db
.prepare(
`
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
`
)
2024-05-11 00:21:44 -07:00
.all(new Float32Array(query));
2024-05-10 20:51:42 -07:00
console.log(rows);