Amazon DynamoDB ベクトル検索機能 (2) 文字列による絞り込みとのハイブリッド検索

// search.mjs
// 自然文クエリを埋め込みにして SearchVectors で「絞り込み付き」ベクトル検索する。
//   node search.mjs "夏に走る軽いシューズ" mp=JP cat=footwear k=5
//     mp  … marketplace（HASH。必須スコープ。省略時 JP）
//     cat … category（INLINE_FILTER。任意の絞り込み）
//     k   … TopK（省略時 5）
//
// 設定・埋め込みロジックは setup.mjs から再利用（import しても表は作られません）。
// 生レスポンスを見たいときは DEBUG=1 を付けて実行してください。

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

// 引数をパース: 最初の非フラグ語をクエリ、mp=/cat=/k= をオプションとして拾う
//   例) node search.mjs "夏に走る軽いシューズ" mp=JP cat=footwear k=5
function parseArgs(argv) {
  const opts = { mp: "JP", cat: undefined, k: 5 };
  let query;
  for (const a of argv.slice(2)) {
    const m = a.match(/^(mp|cat|k)=(.*)$/);
    if (m) {
      if (m[1] === "k") opts.k = Number(m[2]);
      else opts[m[1]] = m[2];
    } else if (query === undefined) {
      query = a;
    }
  }
  return { query: query ?? "夏に走るための軽くて涼しいシューズ", ...opts };
}

async function main() {
  const { query, mp, cat, k } = parseArgs(process.argv);

  // 絞り込み条件を組み立てる
  //   HASH(marketplace) は必須なので常に指定。INLINE_FILTER(category) は任意。
  const names = { "#mp": "marketplace" };
  const values = { ":mp": { S: mp } };
  let condition = "#mp = :mp";
  if (cat) {
    names["#cat"] = "category";
    values[":cat"] = { S: cat };
    condition += " AND #cat = :cat";
  }

  console.log(
    `クエリ: 「${query}」  marketplace=${mp}` +
      (cat ? ` category=${cat}` : " (カテゴリ絞り込みなし)") +
      `  TopK=${k}\n`
  );

  const { embedding } = await embed(query);
  const topK = k;

  let response;
  try {
    response = await ddb.send(
      new SearchVectorsCommand({
        TableName: TABLE_NAME,
        IndexName: VECTOR_INDEX,
        SearchVector: toVectorAttr(embedding),
        TopK: topK,
        SearchConditionExpression: condition,
        ExpressionAttributeNames: names,
        ExpressionAttributeValues: values,
      })
    );
  } 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}  <${attrs.marketplace}/${attrs.category}>`
    );
  });
}

main();

// setup.mjs
// DynamoDB ベクトル検索の準備を1ファイルで行う。
//   node setup.mjs            → テーブル作成 ＋ 埋め込み投入
//   node setup.mjs cleanup    → テーブル削除（課金停止用・元に戻せません）
//
// 注意:
//   - 実際に AWS 上に課金対象のテーブルが作られます（PAY_PER_REQUEST）。
//   - この機能は 2026-08 に GA されたばかりで、CloudFormation ではまだ作れません。
//   - VectorIndexes のネスト構造は最新 SDK で変わる可能性があります。エラー時は
//     例外メッセージに正しいフィールド名が示されるので、それに合わせて調整してください。

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

// marketplace: ベクトルインデックスの HASH（検索時に必須のスコープ）
// category:    INLINE_FILTER（任意の絞り込み）
const PRODUCTS = [
  { productId: "p001", marketplace: "JP", name: "サマーランニングシューズ", category: "footwear", price: 8900,
    description: "通気性の高いメッシュ素材で、暑い季節のジョギングに最適な軽量スニーカー" },
  { productId: "p002", marketplace: "JP", name: "ウィンターハイキングブーツ", category: "footwear", price: 15800,
    description: "保温性のある裏地付きで、雪道でも滑りにくい防水のトレッキングブーツ" },
  { productId: "p003", marketplace: "JP", name: "マラソンレーシングシューズ", category: "footwear", price: 21000,
    description: "反発性のあるフォームを使った、記録を狙うランナー向けの超軽量シューズ" },
  { productId: "p004", marketplace: "JP", name: "ステンレスボトル", category: "kitchen", price: 3200,
    description: "750ml入る、真空断熱で保冷保温に優れたステンレス製ウォーターボトル" },
  { productId: "p005", marketplace: "JP", name: "キャンバススニーカー", category: "footwear", price: 5400,
    description: "普段使いに合わせやすい、キャンバス地のカジュアルシューズ" },
  // 別マーケット(US)の商品。JP にスコープすると意味的に近くても除外されることを確認するため
  { productId: "p006", marketplace: "US", name: "US トレイルランニングシューズ", category: "footwear", price: 19800,
    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",
      // ベクトル属性(List)は不要だが、SearchSchema の HASH / INLINE_FILTER 属性は
      // AttributeDefinitions への定義が必須。
      AttributeDefinitions: [
        { AttributeName: "productId", AttributeType: "S" },
        { AttributeName: "marketplace", AttributeType: "S" }, // SearchSchema: HASH
        { AttributeName: "category", AttributeType: "S" },    // SearchSchema: INLINE_FILTER
      ],
      KeySchema: [{ AttributeName: "productId", KeyType: "HASH" }],
      // ▼ 新機能: ベクトルインデックス（絞り込み対応）
      VectorIndexes: [
        {
          IndexName: VECTOR_INDEX,
          VectorAttribute: { AttributeName: VECTOR_ATTRIBUTE },
          Dimensions: DIMENSIONS,
          DistanceFunction: "COSINE", // COSINE / EUCLIDEAN / DOT_PRODUCT
          Projection: { ProjectionType: "ALL" },
          // 絞り込み用スキーマ
          //   HASH          … 検索を分割するパーティションキー。検索時に = で必須指定
          //   INLINE_FILTER … 任意の絞り込み属性。= のほか比較・範囲演算子も使える
          SearchSchema: [
            { AttributeName: "marketplace", SearchSchemaElementType: "HASH" },
            { AttributeName: "category", SearchSchemaElementType: "INLINE_FILTER" },
          ],
        },
      ],
      // 注: SearchSchema の HASH / INLINE_FILTER 属性は、上の AttributeDefinitions に
      //     型付きで定義しておく必要があります（未定義だと ValidationException になります）。
    })
  );
}

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