PGlite + pgvector で100行で実装するベクトル検索 (node/deno/drizzle)

pglite + pgvector で文章の類似度検索を実装します。

動機

とにかく手っ取り早くローカルにデータを突っ込んでおいて検索する RAG の雛形がほしかったんですが、調べても大規模ストレージを前提とした大掛かりな実装が多いです。

スクリプトを書いたらポンと実行できるセットアップ不要なものがあると、色々と実験ができます。

mastra/rag を読んでたら、簡単にできる気がしたのでやりました。ただ、chunk のドキュメント分割相当のものはまだ作ってません。そこまで難しい概念でもないので、雑に作れそうではあります。

qrdrant も検討しましたが、サーバーを建てるのが面倒でした

https://qdrant.tech/

準備: ベクトル化用の関数

今回は @ai-sdk/openai を使ってベクトル化をします

// OPENAI_API_KEY=
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";
async function generateEmbedding(value: string): number[] {
  const { embedding } = await embed({
    model: openai.embedding("text-embedding-ada-002"),
    value,
  });
  return embedding;
}

この関数は text-embedding-ada-002 で文字列を 1536次元(長)の配列に変換します。

同じ次元のベクトル同士の類似(コサイン類似度)を取ることで、文字列同士の n 次元上の距離を測ることができます。

ここの実装はORAMAでも何でもいいですが、トランスフォーマー上の意味ベクトルの生成なので、同じモデルの同じ次元で embedding を作る必要があります。

pglite + pgvector のセットアップ

deno の node 互換モードで実装したので node でも動くはず。

pglite でデータベースを初期化します。wasmなのでドライバが不要、dataDir が与えられていないのでインメモリにストレージを展開します。

// https://pglite.dev/examples/vector
import { PGlite } from "@electric-sql/pglite";
import { vector } from "@electric-sql/pglite/vector";
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";

// pgvector 拡張を登録
const pglite = new PGlite({
  extensions: { vector },
  // dataDir: "./data"
});
await pglite.exec("CREATE EXTENSION IF NOT EXISTS vector;");

// 検索用 memory テーブルを定義
await pglite.exec(`
  CREATE TABLE IF NOT EXISTS memory (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536)
  );
  CREATE INDEX ON memory USING hnsw (embedding vector_cosine_ops);
`);

CERATE EXTENSION vector で初期化します。text-embedding-ada-002 に合わせて 1536次元で vector 型を初期化します。

セットアップができたので、データを挿入して文書同士の距離を検索します。

// ...
async function generateEmbedding(value: string) {
  const { embedding } = await embed({
    model: openai.embedding("text-embedding-ada-002"),
    value,
  });
  return embedding;
}

// データを挿入
const seedData = [
  "red apple",
  "green tea",
  "yellow banana",
  "pink blossom",
  "black cat",
];
for (const content of seedData) {
  const embedding = await generateEmbedding(content);
  const vec = JSON.stringify(embedding);
  await pglite.exec(
    `INSERT INTO memory (content, embedding) VALUES ('${content}', '${vec}');`
  );
}

// 検索
const queryVec = await generateEmbedding("fruit");
const searchVec = JSON.stringify(queryVec);
const res = await pglite.exec(`
  SELECT
    content,
    embedding <-> '${searchVec}' AS distance
  FROM memory
  ORDER BY distance ASC
  LIMIT 2;`);

console.log(res[0].rows);
// [
//   { content: "red apple", distance: 0.5288196397016931 },
//   { content: "yellow banana", distance: 0.5537656726728964 }
// ]

文書同士の距離が小さいほうが似ていることになります。一応入力から fruit っぽいものが上位にでていますね。

今回は文書が短すぎてあまり意味のある検索ができていませんが、もう少し長い文書を食わせると精度がよくなります。

pglite + pgvector

同じことを drizzle ORM でやります。

ただ、migration をスキップするために直接スキーマを定義します。ここは drizzle-kit でマイグレーションするなら不要です。

// https://orm.drizzle.team/docs/guides/vector-similarity-search
import { PGlite } from "@electric-sql/pglite";
import { vector as pgVector } from "@electric-sql/pglite/vector";
import { index, integer, pgTable, vector, text } from "drizzle-orm/pg-core";
import { drizzle, type PgliteDatabase } from "drizzle-orm/pglite";
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";
import { cosineDistance, sql, desc, gt } from "drizzle-orm";

// openai embedding
async function generateEmbedding(value: string) {
  const { embedding } = await embed({
    model: openai.embedding("text-embedding-ada-002"),
    value,
  });
  return embedding;
}

// schema
export const memory = pgTable(
  "memory",
  {
    id: integer().primaryKey().generatedAlwaysAsIdentity(),
    content: text("content").notNull(),
    embedding: vector("embedding", { dimensions: 1536 }),
  },
  (table) => [
    index("embeddingIndex").using(
      "hnsw",
      table.embedding.op("vector_cosine_ops")
    ),
  ]
);

const pglite = new PGlite({
  extensions: { vector: pgVector },
  // dataDir: "./data",
});
await pglite.exec("CREATE EXTENSION IF NOT EXISTS vector;");

// ここは drizzle-kit で migration するなら不要
await pglite.exec(`
  CREATE TABLE IF NOT EXISTS memory (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536)
  );
  CREATE INDEX ON memory USING hnsw (embedding vector_cosine_ops);
`);

const db = drizzle({
  client: pglite,
  schema: { memory },
});

async function insert(content: string) {
  const embedding = await generateEmbedding(content);
  await db.insert(memory).values({
    content: content,
    embedding,
  });
}
async function query(
  text: string,
  opts: {
    threshold?: number;
    limit?: number;
  }
) {
  const embedding = await generateEmbedding(text);
  const similarity = sql<number>`1 - (${cosineDistance(
    memory.embedding,
    embedding
  )})`;
  return db
    .select({
      id: memory.id,
      content: memory.content,
      similarity,
    })
    .from(memory)
    .where(gt(similarity, opts.threshold ?? 0.7))
    .orderBy((t) => desc(t.similarity))
    .limit(opts.limit ?? 5);
}

// run
await insert("red apple");
await insert("green tea");
await insert("yellow banana");
await insert("pink blossom");
await insert("black cat");

// search
const result = await query("fruit", {
  threshold: 0.7,
  limit: 5,
});
console.log(result);
/**
 * [
  { id: 1, content: "red apple", similarity: 0.8601749430732128 },
  { id: 3, content: "yellow banana", similarity: 0.846671173174945 },
  { id: 4, content: "pink blossom", similarity: 0.8398418116852516 },
  { id: 2, content: "green tea", similarity: 0.8273870081086043 },
  { id: 5, content: "black cat", similarity: 0.7991470334119698 }
]
 */

(1 - 距離) なので similarity が高いほど意味が類似しています。

Drizzle上だとここがちょっとテクニカルですが、うまく型が付いてます。偉い
。

  const similarity = sql<number>`1 - (${cosineDistance(
    memory.embedding,
    embedding
  )})`;
  return db
    .select({
      id: memory.id,
      content: memory.content,
      similarity,
    })
    .from(memory)
    .where(gt(similarity, opts.threshold ?? 0.7))
    .orderBy((t) => desc(t.similarity))
    .limit(opts.limit ?? 5);

おわり

最終的にORM付きの 100 行のコードになりました。自分は今後これを使い回すと思います。

pglite から別の postgres adapter に切り替えればプロダクション移行も簡単そうではあります。

とにかくストレージのセットアップが不要なのが嬉しく、ドライバ不要の pglite のインメモリでローカルのテストが完結します。

https://zenn.dev/mizchi/articles/deno-drizzle-pglite