Merge pull request #112 from tursodatabase/vector

Vector search support

Cargo.toml

@@ -12,7 +12,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-libsql = { version = "0.3.5", features = ["encryption"] }
+libsql = { git = "https://github.com/libsql/libsql/", rev = "84ff2f7578cb634e531ac0075c0286893273238b", features = ["encryption"] }
 tracing = "0.1"
 once_cell = "1.18.0"
 tokio = { version = "1.29.1", features = [ "rt-multi-thread" ] }

examples/vector/package.json

@@ -0,0 +1,16 @@
+{
+  "name": "libsql-examples-vector",
+  "version": "1.0.0",
+  "description": "",
+  "main": "index.js",
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1"
+  },
+  "author": "",
+  "license": "MIT",
+  "dependencies": {
+    "@xenova/transformers": "^2.17.1",
+    "csv-parse": "^5.5.5",
+    "libsql": "../../"
+  }
+}

examples/vector/vector.mjs

@@ -0,0 +1,37 @@
+import { pipeline } from "@xenova/transformers";
+import { createReadStream } from "fs";
+import { parse } from "csv-parse";
+import Database from "libsql";
+
+// Create a embeddings generator.
+const extractor = await pipeline(
+  "feature-extraction",
+  "Xenova/jina-embeddings-v2-small-en",
+  { quantized: false },
+);
+
+// Open a database file.
+const db = new Database("movies.db");
+
+// Create a table for movies with an embedding as a column.
+db.exec("CREATE TABLE movies (title TEXT, year INT, embedding VECTOR(512))");
+
+// Create a vector index on the embedding column.
+db.exec("CREATE INDEX movies_idx USING vector ON movies (embedding)");
+
+// Prepare a SQL `INSERT` statement.
+const stmt = db.prepare(
+  "INSERT INTO movies (title, year, embedding) VALUES (?, ?, vector(?))",
+);
+
+// Process a CSV file of movies generating embeddings for plot synopsis.
+createReadStream("wiki_movie_plots_deduped.csv")
+  .pipe(parse({ columns: true }))
+  .on("data", async (data) => {
+    const title = data.Title;
+    const year = data.Year;
+    const plot = data.Plot;
+    const output = await extractor([plot], { pooling: "mean" });
+    const embedding = output[0].data;
+    stmt.run([title, year, embedding]);
+  });