Amazon ElastiCache でのセマンティックキャッシュ検証 (Valkey 8.2)

import Valkey from 'iovalkey';
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';

// 設定
const config = {
  elasticacheEndpoint: process.env.ELASTICACHE_ENDPOINT || 'localhost',
  elasticachePort: parseInt(process.env.ELASTICACHE_PORT || '6379'),
  awsRegion: process.env.AWS_REGION || 'ap-northeast-1',
  embeddingModelId: 'amazon.titan-embed-text-v2:0',
  indexName: 'idx:test_cache',
  similarityThreshold: 0.45,
};

// Bedrock クライアント
const bedrockClient = new BedrockRuntimeClient({ region: config.awsRegion });

/**
 * テキストから Embedding を生成
 */
async function generateEmbedding(text) {
  const command = new InvokeModelCommand({
    modelId: config.embeddingModelId,
    contentType: 'application/json',
    accept: 'application/json',
    body: JSON.stringify({ inputText: text }),
  });

  const response = await bedrockClient.send(command);
  const parsed = JSON.parse(new TextDecoder().decode(response.body));
  return parsed.embedding;
}

/**
 * Float32Array を Buffer に変換（Valkey 用）
 */
function embeddingToBuffer(embedding) {
  return Buffer.from(new Float32Array(embedding).buffer);
}

/**
 * メインテスト
 */
async function main() {
  console.log('='.repeat(60));
  console.log('ElastiCache Valkey Vector Search Test');
  console.log('='.repeat(60));
  console.log(`Endpoint: ${config.elasticacheEndpoint}:${config.elasticachePort}`);
  console.log(`Index: ${config.indexName}`);
  console.log(`Similarity Threshold: ${config.similarityThreshold * 100}%`);
  console.log('');

  // Valkey 接続
  console.log('[1] Connecting to Valkey...');
  const valkey = new Valkey({
    host: config.elasticacheEndpoint,
    port: config.elasticachePort,
    tls: config.elasticacheEndpoint !== 'localhost' ? {} : undefined,
    maxRetriesPerRequest: 3,
  });

  valkey.on('error', (err) => {
    console.error('Valkey connection error:', err.message);
  });

  try {
    // 接続確認
    const pong = await valkey.ping();
    console.log(`   PING response: ${pong}`);

    // Valkey バージョン確認
    const info = await valkey.info('server');
    const versionMatch = info.match(/valkey_version:(\S+)/);
    const version = versionMatch ? versionMatch[1] : 'unknown';
    console.log(`   Valkey version: ${version}`);

    if (version !== 'unknown' && parseFloat(version) < 8.2) {
      console.error('   ERROR: Valkey 8.2 or higher is required for vector search');
      process.exit(1);
    }
    console.log('');

    // [2] クリーンアップ & インデックス作成
    console.log('[2] Creating vector index...');
    try {
      await valkey.call('FT.DROPINDEX', config.indexName);
      console.log('   Dropped existing index');
    } catch (err) {
      // インデックスが存在しない場合は無視
    }

    // ベクトルフィールドのみのスキーマ（TEXT フィールドは別途 HSET で保存）
    await valkey.call(
      'FT.CREATE', config.indexName,
      'ON', 'HASH',
      'PREFIX', '1', 'test:',
      'SCHEMA',
      'embedding', 'VECTOR', 'FLAT', '6',
      'TYPE', 'FLOAT32',
      'DIM', '1024',
      'DISTANCE_METRIC', 'COSINE'
    );
    console.log('   Index created successfully');
    console.log('');

    // [3] サンプルデータ登録
    console.log('[3] Storing sample data with embeddings...');
    const sampleData = [
      {
        key: 'test:q1',
        question: '営業時間は何時から何時までですか？',
        answer: '営業時間は平日9時から18時までです。土日祝日はお休みをいただいております。',
      },
      {
        key: 'test:q2',
        question: '返品はできますか？',
        answer: '商品到着後7日以内であれば、未使用品に限り返品を承っております。',
      },
      {
        key: 'test:q3',
        question: '配送料はいくらですか？',
        answer: '5,000円以上のご購入で送料無料です。5,000円未満の場合は全国一律500円です。',
      },
    ];

    for (const item of sampleData) {
      console.log(`   Generating embedding for: "${item.question.substring(0, 30)}..."`);
      const embedding = await generateEmbedding(item.question);
      const embeddingBuffer = embeddingToBuffer(embedding);

      await valkey.hset(item.key, {
        question: item.question,
        answer: item.answer,
        embedding: embeddingBuffer,
      });
      console.log(`   Stored: ${item.key}`);
    }
    console.log('');

    // [4] ベクトル検索テスト
    console.log('[4] Testing vector search...');
    const testQueries = [
      '何時に開いていますか？',           // 類似: 営業時間
      '返品の条件を教えてください',       // 類似: 返品
      '送料はかかりますか？',             // 類似: 配送料
      'クレジットカードは使えますか？',   // 非類似: 新規
    ];

    let cacheHits = 0;

    for (const query of testQueries) {
      console.log(`\n   Query: "${query}"`);

      const queryEmbedding = await generateEmbedding(query);
      const queryBuffer = embeddingToBuffer(queryEmbedding);

      // 今回の検証では SORTBY がエラーになったため、使用しません
      const results = await valkey.call(
        'FT.SEARCH', config.indexName,
        '*=>[KNN 1 @embedding $vec AS score]',
        'PARAMS', '2', 'vec', queryBuffer,
        'DIALECT', '2'
      );

      // 結果のパース
      const numResults = results[0];
      if (numResults > 0) {
        const docId = results[1];
        const docData = results[2];

        // docData は [field, value, field, value, ...] 形式
        const dataMap = {};
        for (let i = 0; i < docData.length; i += 2) {
          if (docData[i] !== 'embedding') {
            dataMap[docData[i]] = docData[i + 1];
          }
        }

        const score = parseFloat(dataMap.score);
        const similarity = 1 - score;  // COSINE 距離から類似度へ変換

        // question と answer を取得（インデックスには含まれないので別途取得）
        const storedData = await valkey.hgetall(docId);

        console.log(`   -> Best match: ${docId}`);
        console.log(`   -> Similarity: ${(similarity * 100).toFixed(2)}%`);
        console.log(`   -> Threshold: ${config.similarityThreshold * 100}%`);

        if (similarity >= config.similarityThreshold) {
          console.log(`   -> CACHE HIT`);
          console.log(`   -> Answer: "${storedData.answer.substring(0, 50)}..."`);
          cacheHits++;
        } else {
          console.log(`   -> CACHE MISS: Similarity below threshold`);
        }
      } else {
        console.log('   -> No results found');
      }
    }
    console.log('');

    // [5] クリーンアップ
    console.log('[5] Cleaning up test data...');
    for (const item of sampleData) {
      await valkey.del(item.key);
    }
    await valkey.call('FT.DROPINDEX', config.indexName);
    console.log('   Test data and index removed');
    console.log('');

    console.log('='.repeat(60));
    console.log('Test completed successfully!');
    console.log(`Cache hits: ${cacheHits} / ${testQueries.length}`);
    console.log('='.repeat(60));

  } catch (error) {
    console.error('\nError:', error.message);
    if (error.stack) {
      console.error(error.stack);
    }
    process.exit(1);
  } finally {
    await valkey.quit();
  }
}

main();

    await valkey.call(
      'FT.CREATE', config.indexName,
      'ON', 'HASH',
      'PREFIX', '1', 'test:',
      'SCHEMA',
      'embedding', 'VECTOR', 'FLAT', '6',
      'TYPE', 'FLOAT32',
      'DIM', '1024',
      'DISTANCE_METRIC', 'COSINE'
    );

function embeddingToBuffer(embedding) {
  return Buffer.from(new Float32Array(embedding).buffer);
}

      await valkey.hset(item.key, {
        question: item.question,
        answer: item.answer,
        embedding: embeddingBuffer,
      });

async function generateEmbedding(text) {
  const command = new InvokeModelCommand({
    modelId: config.embeddingModelId,
    contentType: 'application/json',
    accept: 'application/json',
    body: JSON.stringify({ inputText: text }),
  });

  const response = await bedrockClient.send(command);
  const parsed = JSON.parse(new TextDecoder().decode(response.body));
  return parsed.embedding;
}

      const results = await valkey.call(
        'FT.SEARCH', config.indexName,
        '*=>[KNN 1 @embedding $vec AS score]',
        'PARAMS', '2', 'vec', queryBuffer,
        'DIALECT', '2'
      );

        const score = parseFloat(dataMap.score);
        const similarity = 1 - score;  // COSINE 距離から類似度へ変換