Zenn
Closed8

Turso + Drizzle で Vector Search

Ikuma SudoIkuma Sudo

https://docs.turso.tech/sdk/ts/orm/drizzle を参考にDrizzleをセットアップします．Bunを使います．

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"
  }
}

document-searchという名前でデータベースを作成します．

turso db create document-search

データベースのURLを確認します．

turso db show --url document-search

データベースに接続するための認証トークンを取得します．

turso db tokens create document-search

忘れないうちに.envに書いておきます．

.env
TURSO_DATABASE_URL=
TURSO_AUTH_TOKEN=

データベースに接続します．

db/index.ts
import { drizzle } from "drizzle-orm/libsql";
import { createClient } from "@libsql/client";

const turso = createClient({
  url: process.env.TURSO_DATABASE_URL!,
  authToken: process.env.TURSO_AUTH_TOKEN!,
});

export const db = drizzle(turso);

drizzle.config.tsを作成します．

drizzle.config.ts
require("dotenv").config();

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;

Schemaを作成します．

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
})

db/index.tsにインデックスを作成するSQLを追加します．

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 vector_table(vector)
+  USING vector_cosine(3)
+`);

マイグレーションファイルを生成，適用します．

bun run db:generate
bun run db:migrate
Ikuma SudoIkuma Sudo

テーブルが作成されているかどうか確認します．Drizzle Studioを起動します．

bun run db:studio

https://local.drizzle.studio/ にアクセスします．document_tableが作成されているはずです．

ちなみにTursoを使っている場合はDrizzle Studioとほぼ同じものがDashboardから利用できます．

Ikuma SudoIkuma Sudo

OpenAIの埋め込みモデルを使って動作確認用のデータを作成します．OpenAIの埋め込みモデル利用するためにVercel AI SDKを利用します．
https://sdk.vercel.ai/docs/getting-started/nodejs

bun add ai @ai-sdk/openai

.envにOPENAIのAPIキーを追加します．

TURSO_DATABASE_URL=
TURSO_AUTH_TOKEN=
+OPENAI_API_KEY=
Ikuma SudoIkuma Sudo

適当なテキストでベクトル検索を試してみます．

テキストを埋め込み，その結果をデータベースに保存するコード．

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)})`,
  }))
)

クエリするコード

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}, vector(${JSON.stringify(embedding)}))`

const res = await db.select({
  document: documentTable.document,
  similarity: similarity
}).from(documentTable).orderBy(desc(similarity))

console.log(res)

bun run insert.ts

「今日の仙台市の天気は？」でクエリしてみます．

bun run query.ts
output
[
  {
    document: "今日の仙台の天気は曇りです．",
    similarity: 0.8126901388168335,
  }, {
    document: "今日の仙台の気温は14℃です．",
    similarity: 0.7196863889694214,
  }, {
    document: "私の好きな動物はペンギンです．",
    similarity: 0.10916465520858765,
  }, {
    document: "機能はキムチ鍋を食べました．",
    similarity: 0.09915316104888916,
  }
]
このスクラップは1ヶ月前にクローズされました