Amazon DynamoDB ベクトル検索機能をさっそくさわってみた

{
  "name": "dynamodb-vector-test",
  "version": "2.0.1",
  "type": "module",
  "scripts": {
    "setup": "node setup.mjs",
    "search": "node search.mjs",
    "cleanup": "node setup.mjs cleanup"
  },
  "dependencies": {
    "@aws-sdk/client-bedrock-runtime": "^3.658.0",
    "@aws-sdk/client-dynamodb": "^3.1103.0",
    "@aws-sdk/util-dynamodb": "^3.996.0"
  },
  "engines": { "node": ">=18" }
}

// setup.mjs
// DynamoDB ベクトル検索の準備を1ファイルで行う。
//   node setup.mjs            → テーブル作成 ＋ 埋め込み投入
//   node setup.mjs cleanup    → テーブル削除（課金停止用・元に戻せません）

import {
  DynamoDBClient,
  CreateTableCommand,
  DescribeTableCommand,
  PutItemCommand,
  DeleteTableCommand,
} from "@aws-sdk/client-dynamodb";
import { marshall } from "@aws-sdk/util-dynamodb";
import {
  BedrockRuntimeClient,
  InvokeModelCommand,
} from "@aws-sdk/client-bedrock-runtime";

// ===== 設定 =====================================================
export const REGION = process.env.AWS_REGION ?? "ap-northeast-1";
export const TABLE_NAME = process.env.DDB_TABLE ?? "VectorSearchDemo";
export const VECTOR_INDEX = "descriptionVectorIndex";
export const VECTOR_ATTRIBUTE = "descriptionEmbedding";
export const EMBED_MODEL = process.env.BEDROCK_EMBED_MODEL ?? "amazon.titan-embed-text-v2:0";
export const DIMENSIONS = Number(process.env.BEDROCK_EMBED_DIMENSIONS ?? 1024);

const PRODUCTS = [
  { productId: "p001", name: "サマーランニングシューズ", category: "footwear", price: 8900,
    description: "通気性の高いメッシュ素材で、暑い季節のジョギングに最適な軽量スニーカー" },
  { productId: "p002", name: "ウィンターハイキングブーツ", category: "footwear", price: 15800,
    description: "保温性のある裏地付きで、雪道でも滑りにくい防水のトレッキングブーツ" },
  { productId: "p003", name: "マラソンレーシングシューズ", category: "footwear", price: 21000,
    description: "反発性のあるフォームを使った、記録を狙うランナー向けの超軽量シューズ" },
  { productId: "p004", name: "ステンレスボトル", category: "kitchen", price: 3200,
    description: "750ml入る、真空断熱で保冷保温に優れたステンレス製ウォーターボトル" },
  { productId: "p005", name: "キャンバススニーカー", category: "footwear", price: 5400,
    description: "普段使いに合わせやすい、キャンバス地のカジュアルシューズ" },
];

// ===== Bedrock 埋め込み（search.mjs と共通のロジック）===============
const bedrock = new BedrockRuntimeClient({ region: REGION });

export async function embed(text) {
  const isV2 = EMBED_MODEL.startsWith("amazon.titan-embed-text-v2");
  const payload = isV2
    ? { inputText: text, dimensions: DIMENSIONS, normalize: true }
    : { inputText: text };
  const res = await bedrock.send(
    new InvokeModelCommand({
      modelId: EMBED_MODEL,
      contentType: "application/json",
      accept: "application/json",
      body: JSON.stringify(payload),
    })
  );
  return JSON.parse(new TextDecoder().decode(res.body)); // { embedding, inputTextTokenCount }
}

// ===== DynamoDB =================================================
const ddb = new DynamoDBClient({ region: REGION });
const sleep = (ms) => new Promise((r) => setTimeout(r, ms));

async function tableExists() {
  try {
    await ddb.send(new DescribeTableCommand({ TableName: TABLE_NAME }));
    return true;
  } catch (err) {
    if (err?.name === "ResourceNotFoundException") return false;
    throw err;
  }
}

async function createTable() {
  console.log(`テーブル作成中: ${TABLE_NAME} (region: ${REGION})`);
  await ddb.send(
    new CreateTableCommand({
      TableName: TABLE_NAME,
      BillingMode: "PAY_PER_REQUEST",
      // ベクトル属性はキーではないので AttributeDefinitions には含めない
      AttributeDefinitions: [{ AttributeName: "productId", AttributeType: "S" }],
      KeySchema: [{ AttributeName: "productId", KeyType: "HASH" }],
      // ▼ 新機能: ベクトルインデックス（まず最小構成）
      VectorIndexes: [
        {
          IndexName: VECTOR_INDEX,
          VectorAttribute: { AttributeName: VECTOR_ATTRIBUTE },
          Dimensions: DIMENSIONS,
          DistanceFunction: "COSINE", // COSINE / EUCLIDEAN / DOT_PRODUCT
          Projection: { ProjectionType: "ALL" },
          // 絞り込みを使う場合はここに SearchSchema(HASH / INLINE_FILTER)を追加
        },
      ],
    })
  );
}

async function waitActive() {
  process.stdout.write("テーブルが ACTIVE になるのを待機");
  for (let i = 0; i < 60; i++) {
    const { Table } = await ddb.send(new DescribeTableCommand({ TableName: TABLE_NAME }));
    process.stdout.write(".");
    if (Table?.TableStatus === "ACTIVE") return console.log(" OK");
    await sleep(3000);
  }
  throw new Error("テーブルが時間内に ACTIVE になりませんでした");
}

async function seed() {
  console.log("埋め込みを生成して投入中...");
  for (const p of PRODUCTS) {
    const { embedding } = await embed(p.description);
    await ddb.send(
      new PutItemCommand({
        TableName: TABLE_NAME,
        // marshall が number[] を List<Number> ({L:[{N},...]}) に変換
        Item: marshall({ ...p, [VECTOR_ATTRIBUTE]: embedding }, { removeUndefinedValues: true }),
      })
    );
    console.log(`  投入: ${p.productId} ${p.name} (${embedding.length}次元)`);
  }
}

async function runSetup() {
  if (await tableExists()) {
    console.log(`テーブル ${TABLE_NAME} は既に存在します。作成をスキップします。`);
  } else {
    await createTable();
    await waitActive();
  }
  await seed();
  console.log("\n完了しました。次に: npm run search");
}

async function runCleanup() {
  try {
    await ddb.send(new DeleteTableCommand({ TableName: TABLE_NAME }));
    console.log(`削除リクエストを送信しました: ${TABLE_NAME} (${REGION})`);
  } catch (err) {
    if (err?.name === "ResourceNotFoundException") {
      return console.log(`テーブル ${TABLE_NAME} は存在しません（すでに削除済み）。`);
    }
    throw err;
  }
}

// ===== エントリポイント ==========================================
const isDirectRun = import.meta.url === `file://${process.argv[1]}`;
if (isDirectRun) {
  const mode = process.argv[2];
  const task = mode === "cleanup" ? runCleanup() : runSetup();
  task.catch((err) => {
    console.error("\n失敗しました:", err?.name, err?.message);
    console.error("→ フィールド名の不一致が原因の場合、上のメッセージに正しい名前が含まれています。");
    process.exit(1);
  });
}

// search.mjs
// 自然文クエリを埋め込みにして SearchVectors でベクトル検索する。
//   node search.mjs "夏に走るための軽いシューズ" [TopK]

import * as ddbPkg from "@aws-sdk/client-dynamodb";
import { unmarshall } from "@aws-sdk/util-dynamodb";
import { embed, REGION, TABLE_NAME, VECTOR_INDEX } from "./setup.mjs";

const { DynamoDBClient, SearchVectorsCommand } = ddbPkg;

if (!SearchVectorsCommand) {
  console.error(
    "SearchVectorsCommand が見つかりません。SDK を更新してください:\n" +
      "  npm install @aws-sdk/client-dynamodb@latest @aws-sdk/util-dynamodb@latest"
  );
  process.exit(1);
}

const ddb = new DynamoDBClient({ region: REGION });

// COSINE は「距離」(0=完全一致, 2=正反対, 小さいほど近い)。表示用に類似度%へ変換。
const toSimilarityPct = (d) => ((1 - d / 2) * 100).toFixed(1);
// SearchVector は List<Number> の「中身」（Number属性値の配列）。{L:...} では包まない。
const toVectorAttr = (embedding) => embedding.map((n) => ({ N: String(n) }));

async function main() {
  const queryText = process.argv[2] ?? "夏に走るための軽くて涼しいシューズ";
  const topK = Number(process.argv[3] ?? 5);
  console.log(`クエリ: 「${queryText}」  TopK=${topK}\n`);

  const { embedding } = await embed(queryText);

  let response;
  try {
    response = await ddb.send(
      new SearchVectorsCommand({
        TableName: TABLE_NAME,
        IndexName: VECTOR_INDEX,
        SearchVector: toVectorAttr(embedding),
        TopK: topK,
        // 絞り込み(HASHスコープ / INLINE_FILTER)は SearchConditionExpression 等をここに追加
      })
    );
  } catch (err) {
    console.error("検索に失敗しました:", err?.name, err?.message);
    process.exit(1);
  }

  // 生レスポンスを見たいときは DEBUG=1 を付けて実行
  if (process.env.DEBUG) {
    console.log("=== 生レスポンス（デバッグ用）===");
    console.log(JSON.stringify(response, null, 2));
    console.log("================================\n");
  }

  // 実レスポンス: { SearchResults: [{ Item: <marshall済み>, Score: <距離> }] }
  const rows = response.SearchResults ?? [];
  if (rows.length === 0) {
    console.log("結果が空でした。");
    return;
  }

  console.log("類似度が高い順（Score=距離, 小さいほど近い）:");
  rows.forEach((row, i) => {
    const attrs = unmarshall(row.Item);
    const score = row.Score;
    console.log(
      `  ${i + 1}. [dist=${score.toFixed(4)} / 類似度${toSimilarityPct(score)}%] ${attrs.name}`
    );
  });
}

main();