向量搜尋 (Vector Search) 與嵌入技術 (Embeddings) 淺談：通過 Postgres (PGVector) 實作的深入解析

CREATE TABLE items (
    id serial PRIMARY KEY,
    embedding vector(1536) -- 假設每個向量有 1536 維
);

CREATE INDEX ON documents 
USING ivfflat (embedding vector_cosine_ops) 
WITH (lists = 100);

CREATE INDEX ON documents 
USING hnsw (embedding vector_cosine_ops) 
WITH (m = 16, ef_construction = 64);

-- 推薦的生產環境設定
CREATE INDEX ON documents 
USING hnsw (embedding vector_cosine_ops) 
WITH (
    m = 16,                 -- 預設值，適合大多數場景
    ef_construction = 64    -- 預設值，平衡建立速度與品質
);

-- 查詢時動態調整精度（不需重建索引）
SET hnsw.ef_search = 100;  -- 範圍 80-120，數值越高召回率越高但速度越慢

import { SQL } from "bun";
import { faker } from "@faker-js/faker";
import { generateEmbedding } from "./embeddings";
import pgvector from "pgvector";

const sql = new SQL({ url: process.env.DATABASE_URL });

const initDb = async () => {
  /**
   * Create extension if not exists
   */
  await sql`CREATE EXTENSION IF NOT EXISTS vector;`;

  /**
   * Create table if not exists
   */
  await sql`CREATE TABLE IF NOT EXISTS documents (
    id SERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding VECTOR(1536) NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
  );`;
};

const seedDb = async () => {
  const documents = await Promise.all(
    Array.from({ length: 3 }, async () => {
      const content = faker.lorem.paragraph();
      return {
        title: faker.lorem.sentence(),
        content,
        embedding: pgvector.toSql(await generateEmbedding(content)),
      };
    })
  );

  await sql`INSERT INTO documents ${sql(documents)}`;

  console.log("Database seeded successfully");
};

const seedDbWithCapitals = async () => {
  const documents = await Promise.all([
    {
      title: "Taiwan",
      content: "The capital of Taiwan is Taipei",
      embedding: pgvector.toSql(await generateEmbedding("The capital of Taiwan is Taipei")),
    },
    {
      title: "Japan",
      content: "The capital of Japan is Tokyo",
      embedding: pgvector.toSql(await generateEmbedding("The capital of Japan is Tokyo")),
    },
    {
      title: "United States",
      content: "The capital of United States is Washington, D.C.",
      embedding: pgvector.toSql(await generateEmbedding("The capital of United States is Washington, D.C.")),
    },
  ]);

  await sql`INSERT INTO documents ${sql(documents)}`;

  console.log("Database seeded successfully");
};

const searchDb = async (query: string) => {
  const embedding = await generateEmbedding(query);
  if (!embedding) {
    throw new Error("Embedding generation failed");
  }
  const sqlEmbedding = pgvector.toSql(embedding);
  const results = await sql`SELECT * FROM documents ORDER BY embedding <=> ${sqlEmbedding} LIMIT 3`.values();
  return results;
};

const demo = async (search = "What is the capital of Taiwan?", seed?: true | "capitals") => {
  await initDb();
  if (seed) {
    if (seed === "capitals") {
      await seedDbWithCapitals();
    } else {
      await seedDb();
    }
  }
  const results = await searchDb(search);
  console.log(results);
};

const Script = {
  init: async (seed?: true | "capitals") => {
    await initDb();
    if (seed) {
      if (seed === "capitals") {
        await seedDbWithCapitals();
      } else {
        await seedDb();
      }
    }
  },
  initDb,
  seedDb,
  searchDb,
  demo,
};

export default Script;

import Script from "./src/scripts";

await Script.demo(
  "What is the capital of Taiwan?",
  // seed the database or not
  "capitals"
);

[
  [ 4, "Taiwan", "The capital of Taiwan is Taipei", "[0.000814143,-0.019345103,-0.0027284368,...,-0.02137615]",
    2025-05-08T08:29:22.591Z
  ], [ 7, "Japan", "The capital of Japan is Tokyo", "[0.000814143,-0.019345103,-0.0027284368,...,-0.02137615]",
    2025-05-08T08:29:48.303Z
  ], [ 6, "United States", "The capital of United States is Washington, D.C.", "[0.0072439546,-0.01591172,-0.016960844,...,-0.0050426666]",
    2025-05-08T08:29:22.591Z
  ], count: 3, command: "SELECT"
]