🔍
Turso + Drizzleでベクトル検索
Turso (libsql) ではベクトル検索がデフォルトで利用できます．この記事では Drizzle から Turso のベクトル検索機能を利用する方法を紹介します．
コードは GitHub にあります．
※元のスクラップ
Drizzle のセットアップ
Bun を使って Drizzle をセットアップします．
bun add drizzle-orm @libsql/client
bun add -D drizzle-kit
※ Bun を使っているので dotenv は不要です．
package.json にスクリプトを追加します．
package.json
{
"scripts": {
"db:generate": "drizzle-kit generate",
"db:migrate": "drizzle-kit migrate",
"db:studio": "drizzle-kit studio"
}
}
Turso データベースの作成
document-search という名前でデータベースを作成します．
turso db create document-search
データベースの URL を確認します．
turso db show --url document-search
データベースに接続するための認証トークンを取得します．
turso db tokens create document-search
.env にデータベース URL と認証トークンを書き込んでおきます．
.env
TURSO_DATABASE_URL=<your_database_url>
TURSO_AUTH_TOKEN=<your_auth_token>
データベース接続設定
データベースに接続するためのコード db/index.ts を作成します．
db/index.ts
import { drizzle } from "drizzle-orm/libsql";
import { createClient } from "@libsql/client";
import { sql } from "drizzle-orm";
const turso = createClient({
url: process.env.TURSO_DATABASE_URL!,
authToken: process.env.TURSO_AUTH_TOKEN!,
});
export const db = drizzle(turso);
await db.run(sql`
CREATE INDEX IF NOT EXISTS vector_index
ON document_table(embedding)
USING vector_cosine(1536)
`);
drizzle.config.ts を作成します．
drizzle.config.ts
import type { Config } from "drizzle-kit";
export default {
schema: "./db/schema.ts",
out: "./migrations",
dialect: "turso",
dbCredentials: {
url: process.env.TURSO_DATABASE_URL!,
authToken: process.env.TURSO_AUTH_TOKEN,
},
} satisfies Config;
スキーマ定義
スキーマを定義する db/schema.ts を作成します．
db/schema.ts
import { sql } from "drizzle-orm";
import { customType, integer, sqliteTable, text } from "drizzle-orm/sqlite-core";
const float32Array = customType<{
data: number[];
config: { dimensions: number };
configRequired: true;
driverData: Buffer;
}>({
dataType(config) {
return `F32_BLOB(${config.dimensions})`;
},
fromDriver(value: Buffer) {
return Array.from(new Float32Array(value.buffer));
},
toDriver(value: number[]) {
return sql`vector32(${JSON.stringify(value)})`;
},
});
export const documentTable = sqliteTable("document_table", {
id: integer("id").primaryKey(),
document: text("document").notNull(),
embedding: float32Array("embedding", { dimensions: 1536 }) // openai text-embedding-3-small model output
})
マイグレーション
マイグレーションファイルを生成し，適用します．
bun run db:generate
bun run db:migrate
Drizzle Studio で確認
テーブルが作成されているか確認するために Drizzle Studio を起動します．
bun run db:studio
ブラウザで https://local.drizzle.studio/ にアクセスします．document_table が作成されているはずです．Turso を使用している場合は，Turso Dashboard からも同様の確認ができます．
OpenAI の埋め込みモデルを使ったデータ作成
OpenAI の埋め込みモデルを使って動作確認用のデータを作成します．OpenAI の埋め込みモデルを利用するために Vercel AI SDK を利用します．
bun add ai @ai-sdk/openai
.env に OpenAI の API キーを追加します．
.env
TURSO_DATABASE_URL=<your_database_url>
TURSO_AUTH_TOKEN=<your_auth_token>
+OPENAI_API_KEY=<your_api_key>
データ挿入
テキストを埋め込み，その結果をデータベースに保存するコード insert.ts を作成します．
insert.ts
import { openai } from '@ai-sdk/openai';
import { embedMany } from 'ai';
import { db } from './db';
import { documentTable } from './db/schema';
import { sql } from 'drizzle-orm';
const documents = [
"今日の仙台の気温は14℃です．",
"今日の仙台の天気は曇りです．",
"私の好きな動物はペンギンです．",
"昨日はキムチ鍋を食べました．"
];
const { embeddings } = await embedMany({
model: openai.embedding('text-embedding-3-small'),
values: documents,
});
await db.insert(documentTable).values(
embeddings.map((embedding, i) => ({
document: documents[i],
embedding: sql`vector32(${JSON.stringify(embedding)})`,
}))
);
console.log("Inserted data successfully.");
bun run insert.ts
ベクトル検索
クエリを実行するコード query.ts を作成します．
query.ts
import { desc, sql } from "drizzle-orm";
import { db } from "./db";
import { documentTable } from "./db/schema";
import { embed } from "ai";
import { openai } from "@ai-sdk/openai";
const query = "今日の仙台市の天気は？";
const { embedding } = await embed({
model: openai.embedding("text-embedding-3-small"),
value: query,
});
const similarity = sql<number>`1 - vector_distance_cos(${documentTable.embedding}, vector32(${JSON.stringify(embedding)}))`;
const res = await db.select({
document: documentTable.document,
similarity: similarity,
}).from(documentTable).orderBy(desc(similarity));
console.log(res);
実行します．
bun run query.ts
output
[
{
document: "今日の仙台の天気は曇りです．",
similarity: 0.8126901388168335,
}, {
document: "今日の仙台の気温は14℃です．",
similarity: 0.7196863889694214,
}, {
document: "私の好きな動物はペンギンです．",
similarity: 0.10916465520858765,
}, {
document: "機能はキムチ鍋を食べました．",
similarity: 0.09915316104888916,
}
]
確かに似た意味の文を取得できています．
Discussion