AWS MemoryDB のセマンティックキャッシュによる LLM アプリケーションの高速化検証

[ユーザー] 
    ↓ 1. 質問送信
[Amazon API Gateway] 
    ↓ 2. Lambda呼び出し
[AWS Lambda] 
    ↓ 3. 埋め込み生成
[Amazon Bedrock - Titan Embeddings V2]
    ↓ 4. ベクトル検索
[Amazon MemoryDB]
    ↓ 5. 回答生成 (キャッシュミス時)
[Amazon Bedrock - Claude Sonnet 4]

mkdir semantic-cache-lambda
cd semantic-cache-lambda
npm init -y

npm install @aws-sdk/client-bedrock-runtime iovalkey

semantic-cache-lambda/
├── index.js              # Lambda ハンドラー
├── semantic-cache.js     # セマンティックキャッシュクラス
├── package.json
└── node_modules/

// semantic-cache.js - セマンティックキャッシュクラス
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';
import Redis from 'iovalkey';

export class SemanticCache {
    constructor() {
        this.client = null;
        this.bedrockClient = new BedrockRuntimeClient({
            region: process.env.AWS_REGION || 'ap-northeast-1'
        });
        this.threshold = 0.4; // 類似度の閾値 (0.4 = 40% 以上の類似度でキャッシュヒット)
        this.indexName = 'semantic_cache_idx';
    }

    async connect() {
        if (!this.client) {
            // MemoryDB に接続
            this.client = new Redis({
                host: process.env.MEMORYDB_ENDPOINT,
                port: 6379,
                tls: { rejectUnauthorized: false }
            });
            await this.setupIndex();
        }
    }

    // ベクトル検索用のインデックスを作成
    async setupIndex() {
        try {
            await this.client.call('FT.INFO', this.indexName);
            console.log('✅ ベクトルインデックスは既に存在します');
        } catch {
            console.log('🔧 ベクトルインデックスを作成中...');
            await this.client.call(
                'FT.CREATE', this.indexName,
                'ON', 'HASH',
                'PREFIX', '1', 'cache:',
                'SCHEMA',
                'embed', 'VECTOR', 'FLAT', '6',
                'TYPE', 'FLOAT32',
                'DIM', '1024',
                'DISTANCE_METRIC', 'COSINE',
                'question', 'TEXT',
                'answer', 'TEXT'
            );
            console.log('✅ ベクトルインデックスを作成しました');
        }
    }

    // Titan Embeddings V2 でテキストをベクトル化
    async generateEmbedding(text) {
        const command = new InvokeModelCommand({
            modelId: 'amazon.titan-embed-text-v2:0',
            body: JSON.stringify({
                inputText: text,
                dimensions: 1024
            })
        });
        const response = await this.bedrockClient.send(command);
        const result = JSON.parse(new TextDecoder().decode(response.body));
        return result.embedding;
    }

    // MemoryDB の KNN 検索でセマンティック検索
    async searchCache(queryEmbedding) {
        const queryVector = Buffer.from(new Float32Array(queryEmbedding).buffer);

        // KNN 検索で類似するベクトルを検索
        const results = await this.client.call(
            'FT.SEARCH', this.indexName,
            `*=>[KNN 1 @embed $vec AS score]`,
            'PARAMS', '2', 'vec', queryVector,
            'RETURN', '3', 'question', 'answer', 'score',
            'SORTBY', 'score',
            'DIALECT', '2'
        );

        // 検索結果を解析
        if (results[0] > 0) {
            const fields = results[2];
            const doc = {};
            for (let i = 0; i < fields.length; i += 2) {
                doc[fields[i]] = fields[i + 1];
            }

            // 類似度の閾値チェック
            const similarity = 1 - parseFloat(doc.score);
            if (similarity >= this.threshold) {
                return doc;
            }
        }
        return null;
    }

    // 質問と回答をキャッシュに保存
    async saveToCache(question, answer, embedding) {
        const key = `cache:${Date.now()}`;
        const vectorBytes = Buffer.from(new Float32Array(embedding).buffer);

        // ハッシュとして保存
        await this.client.hset(key, {
            question,
            answer,
            embed: vectorBytes
        });

        // 24時間の TTL を設定
        await this.client.expire(key, 86400);
    }
}

// index.js - Lambda ハンドラー
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';
import { SemanticCache } from './semantic-cache.js';

// Lambda コンテナ再利用のためにグローバルで初期化
const cache = new SemanticCache();
const bedrockClient = new BedrockRuntimeClient({
    region: process.env.AWS_REGION || 'ap-northeast-1'
});

let isInitialized = false;

// Claude Sonnet 4 で回答を生成
async function generateAnswer(question) {
    const command = new InvokeModelCommand({
        modelId: 'apac.anthropic.claude-sonnet-4-20250514-v1:0',
        body: JSON.stringify({
            messages: [{
                role: "user",
                content: `簡潔で分かりやすく答えてください: ${question}`
            }],
            max_tokens: 300,
            anthropic_version: "bedrock-2023-05-31"
        })
    });

    const response = await bedrockClient.send(command);
    const result = JSON.parse(new TextDecoder().decode(response.body));
    return result.content[0].text;
}

export const handler = async (event) => {
    try {
        // 初回実行時のみ MemoryDB に接続
        if (!isInitialized) {
            console.log('🚀 セマンティックキャッシュを初期化中...');
            await cache.connect();
            isInitialized = true;
            console.log('✅ 初期化完了');
        }

        // リクエストから質問を取得
        const body = typeof event.body === 'string' ? JSON.parse(event.body) : event.body;
        const { question } = body;

        if (!question) {
            return {
                statusCode: 400,
                body: JSON.stringify({ error: 'question パラメータが必要です' })
            };
        }

        const startTime = Date.now();
        console.log(`❓ 質問: ${question}`);

        // 1. Titan Embeddings V2 で質問をベクトル化
        const embedding = await cache.generateEmbedding(question);

        // 2. MemoryDB でセマンティック検索
        const cached = await cache.searchCache(embedding);

        let answer, isFromCache = false, llmTime = 0;

        if (cached) {
            // キャッシュヒット - 既存の回答を返す
            answer = cached.answer;
            isFromCache = true;
            const similarity = 1 - parseFloat(cached.score);
            console.log(`✅ キャッシュヒット (類似度: ${similarity.toFixed(3)})`);
            console.log(`   類似質問: ${cached.question}`);
        } else {
            // キャッシュミス - Claude で新しい回答を生成
            console.log('❌ キャッシュミス - Claude で回答生成中...');
            const llmStart = Date.now();
            answer = await generateAnswer(question);
            llmTime = Date.now() - llmStart;

            // 新しい回答をキャッシュに保存
            await cache.saveToCache(question, answer, embedding);
        }

        const totalTime = Date.now() - startTime;
        console.log(`💬 回答: ${answer}`);
        console.log(`⏱️  レスポンス時間: ${totalTime}ms (キャッシュ: ${isFromCache ? 'ヒット' : 'ミス'})`);

        return {
            statusCode: 200,
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({
                success: true,
                data: {
                    question,
                    answer,
                    isFromCache,
                    responseTime: totalTime,
                    llmTime
                }
            })
        };

    } catch (error) {
        console.error('❌ エラー:', error);
        return {
            statusCode: 500,
            body: JSON.stringify({ success: false, error: error.message })
        };
    }
};

{
  "name": "semantic-cache-lambda",
  "version": "1.0.0",
  "type": "module",
  "main": "index.js",
  "dependencies": {
    "@aws-sdk/client-bedrock-runtime": "^3.x.x",
    "iovalkey": "^5.x.x"
  }
}

zip -r function.zip .

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ec2:CreateNetworkInterface",
                "ec2:DescribeNetworkInterfaces",
                "ec2:DeleteNetworkInterface"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "bedrock:InvokeModel"
            ],
            "Resource": [
                "arn:aws:bedrock:*:*:foundation-model/amazon.titan-embed-text-v2:0",
                "arn:aws:bedrock:*:*:foundation-model/anthropic.claude-sonnet-4-20250514-v1:0",
                "arn:aws:bedrock:*:*:foundation-model/apac.anthropic.claude-sonnet-4-20250514-v1:0",
                "arn:aws:bedrock:*:*:inference-profile/apac.anthropic.claude-sonnet-4-20250514-v1:0"
            ]
        }
    ]
}

{"body": "{\"question\": \"Python とは何ですか？\"}"}

{"body": "{\"question\": \"Pythonって何？\"}"}

{"body": "{\"question\": \"MySQL について教えて\"}"}

{"body": "{\"question\": \"mysqlとは？\"}"}

{"body": "{\"question\": \"AWS とは何ですか？\"}"}

{"body": "{\"question\": \"AmazonWebServiceって何？\"}"}

{"body": "{\"question\": \"機械学習とは何ですか？\"}"}

{"body": "{\"question\": \"ML って何？\"}"}

{"body": "{\"question\": \"HTML とは何ですか？\"}"}

{"body": "{\"question\": \"エイチティーエムエルって何？\"}"}