Amazon ElastiCache for Valkey + Amazon Bedrock AgentCoreでセマンティックキャッシュを実装してみた

git clone https://github.com/yuu551/agentcore-elastic-cache-blog-sample.git
cd agentcore-elastic-cache-blog-sample
pnpm install

import os
from dataclasses import dataclass

@dataclass
class Config:
    # Valkey 接続設定
    valkey_host: str = os.getenv('VALKEY_HOST', 'localhost')
    valkey_port: int = int(os.getenv('VALKEY_PORT', '6379'))
    valkey_ssl: bool = os.getenv('VALKEY_SSL', 'true').lower() == 'true'

    # Amazon Bedrock モデル設定
    embedding_model_id: str = os.getenv('EMBEDDING_MODEL_ID', 'amazon.titan-embed-text-v2:0')
    llm_model_id: str = os.getenv('LLM_MODEL_ID', 'us.anthropic.claude-haiku-4-5-20251001-v1:0')

    # セマンティックキャッシュ設定
    similarity_threshold: float = float(os.getenv('SIMILARITY_THRESHOLD', '0.8'))

    # AWS / ベクトル検索設定
    aws_region: str = os.getenv('AWS_REGION', 'us-west-2')
    vector_dims: int = int(os.getenv('VECTOR_DIMS', '1024'))

import json
import struct
import time
from hashlib import md5
from typing import Optional

import boto3
import valkey
from .config import Config

class SemanticCache:
    INDEX_NAME = 'semantic_cache_idx'
    KEY_PREFIX = 'cache:'

    def __init__(self, config: Config):
        self.config = config

        # Bedrock クライアント（埋め込みベクトル生成用）
        self.bedrock_client = boto3.client('bedrock-runtime', region_name=config.aws_region)

        # Valkey クライアント
        self.client = valkey.Valkey(
            host=config.valkey_host,
            port=int(config.valkey_port),
            decode_responses=False,
            ssl=config.valkey_ssl,
        )
        self.client.ping()
        self._setup_index()

    def _setup_index(self) -> None:
        """ベクトル検索インデックスをセットアップ"""
        try:
            self.client.execute_command('FT.INFO', self.INDEX_NAME)
        except Exception as e:
            if 'unknown index name' in str(e).lower() or 'no such index' in str(e).lower():
                self.client.execute_command(
                    'FT.CREATE', self.INDEX_NAME,
                    'ON', 'HASH',
                    'PREFIX', '1', self.KEY_PREFIX,
                    'SCHEMA',
                    'query', 'TAG',
                    'embedding', 'VECTOR', 'HNSW', '6',
                    'TYPE', 'FLOAT32',
                    'DIM', str(self.config.vector_dims),
                    'DISTANCE_METRIC', 'COSINE',
                )

    def _get_embedding(self, text: str) -> list[float]:
        """Titan Text Embeddings V2でベクトル生成"""
        response = self.bedrock_client.invoke_model(
            modelId=self.config.embedding_model_id,
            contentType='application/json',
            accept='application/json',
            body=json.dumps({'inputText': text}),
        )
        return json.loads(response['body'].read())['embedding']

    def search(self, query: str, k: int = 3, min_similarity: Optional[float] = None) -> Optional[dict]:
        """類似クエリを検索"""
        threshold = min_similarity if min_similarity is not None else self.config.similarity_threshold

        # クエリをベクトルに変換
        embedding = self._get_embedding(query)
        embedding_bytes = struct.pack(f'{len(embedding)}f', *embedding)

        # FT.SEARCH で KNN ベクトル検索を実行
        results = self.client.execute_command(
            'FT.SEARCH', self.INDEX_NAME,
            f'*=>[KNN {k} @embedding $vec AS score]',
            'PARAMS', '2', 'vec', embedding_bytes,
            'RETURN', '1', 'score',
        )

        if not results or results[0] == 0:
            return None

        # ドキュメントキーとスコアを取得
        doc_key = results[1]
        if isinstance(doc_key, bytes):
            doc_key = doc_key.decode('utf-8')

        top_fields = results[2]
        score_value = top_fields[1]
        if isinstance(score_value, bytes):
            score_value = score_value.decode('utf-8')

        # COSINE距離を類似度に変換（0=完全一致 → 類似度1.0）
        distance = float(score_value)
        similarity = 1 - (distance / 2)

        if similarity < threshold:
            return None

        # 回答本文はインデックスに載せていないため、HGETで取得
        answer = self.client.hget(doc_key, 'answer')
        if isinstance(answer, bytes):
            answer = answer.decode('utf-8')

        return {
            'answer': answer,
            'score': similarity,
            'timing_ms': { ... },  # 処理時間の計測結果（省略）
        }

    def store(self, query: str, answer: str) -> bool:
        """クエリと回答をキャッシュに保存"""
        key = f"cache:{md5(query.encode('utf-8')).hexdigest()}"
        embedding = self._get_embedding(query)
        embedding_bytes = struct.pack(f'{len(embedding)}f', *embedding)

        self.client.hset(key, mapping={
            'query': query,
            'answer': answer,
            'embedding': embedding_bytes,
        })
        return True

import logging
import time
from typing import Any

from bedrock_agentcore.runtime import BedrockAgentCoreApp
from strands import Agent

from .config import Config
from .semantic_cache import SemanticCache

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

config = Config()
cache = SemanticCache(config=config)
agent = Agent(model=config.llm_model_id)

app = BedrockAgentCoreApp()

@app.entrypoint
def invoke(payload: dict) -> dict[str, Any]:
    user_message = payload.get('prompt', 'Hello! How can I help you today?')
    threshold = payload.get('min_similarity', config.similarity_threshold)

    request_start = time.perf_counter()

    # 1. セマンティックキャッシュを検索
    cache_search_start = time.perf_counter()
    cached = cache.search(query=user_message, min_similarity=threshold)
    cache_search_ms = (time.perf_counter() - cache_search_start) * 1000

    if cached:
        total_ms = (time.perf_counter() - request_start) * 1000
        return {
            'response': cached['answer'],
            'cache_hit': True,
            'similarity_score': cached['score'],
            'timing_ms': {
                'cache_search': round(cache_search_ms, 1),
                'total': round(total_ms, 1),
                **cached.get('timing_ms', {}),
            },
        }

    # 2. キャッシュミス → Strands Agent で回答を生成
    agent_start = time.perf_counter()
    result = agent(user_message)
    agent_ms = (time.perf_counter() - agent_start) * 1000

    # result.message は辞書型のため、テキスト部分を抽出
    message = result.message
    if isinstance(message, dict):
        content = message.get('content', [])
        answer = content[0].get('text', str(result)) if content else str(result)
    else:
        answer = str(message)

    # 3. 生成した回答をキャッシュに保存
    cache_store_start = time.perf_counter()
    cache.store(user_message, answer)
    cache_store_ms = (time.perf_counter() - cache_store_start) * 1000

    total_ms = (time.perf_counter() - request_start) * 1000
    return {
        'response': answer,
        'cache_hit': False,
        'model_id': config.llm_model_id,
        'timing_ms': {
            'cache_search': round(cache_search_ms, 1),
            'agent_invocation': round(agent_ms, 1),
            'cache_store': round(cache_store_ms, 1),
            'total': round(total_ms, 1),
        },
    }

if __name__ == '__main__':
    app.run()

pnpm dlx cdk deploy SemanticCacheStack

 ✅  SemanticCacheStack-dev

✨  Deployment time: 617.23s

Outputs:
SemanticCacheStack-dev.AgentCoreConstructRuntimeArnC90D93E6 = arn:aws:bedrock-agentcore:us-west-2:xxxxxxxxxxxx:runtime/semanticCacheAgent-xxxxxxxxxx
SemanticCacheStack-dev.AgentCoreConstructRuntimeId94C4621A = semanticCacheAgent-xxxxxxxxxx
SemanticCacheStack-dev.ElastiCacheConstructClusterEndpointCDBE4A08 = master.semantic-cache-dev-cluster.xxxxx.usw2.cache.amazonaws.com:6379
SemanticCacheStack-dev.StackEnvironment = dev

uv init
uv add boto3

uv run invoke.py "セマンティックキャッシュとは何ですか？"

━━━ リクエスト ━━━
  プロンプト: セマンティックキャッシュとは何ですか？
  セッション ID: eb3b6afe-bd02-43fb-b550-933e07c9507a

━━━ レスポンス (キャッシュミス ✗) ━━━

# セマンティックキャッシュについて

## 概要

セマンティックキャッシュは、意味的に類似したクエリに対して、キャッシュされた結果を再利用するキャッシング技術です。

（中略）

── メタ情報 ──
  使用モデル: us.anthropic.claude-haiku-4-5-20251001-v1:0
── 処理時間 ──
  キャッシュ検索: 145.2 ms
  LLM 呼び出し: 4802.0 ms
  キャッシュ保存: 119.9 ms
  合計: 5067.4 ms

uv run invoke.py "セマンティックキャッシュとは何ですか？"

━━━ リクエスト ━━━
  プロンプト: セマンティックキャッシュとは何ですか？
  セッション ID: ddd93382-9a59-4959-82d1-8d29c0c04607

━━━ レスポンス (キャッシュヒット ✓) ━━━

# セマンティックキャッシュについて

（先ほどと同じ回答が返却される）

── メタ情報 ──
  類似度スコア: 1.0000
── 処理時間 ──
  キャッシュ検索: 161.1 ms
  合計: 161.1 ms
  埋め込み生成: 141.8 ms
  ベクトル検索: 18.2 ms

uv run invoke.py "セマンティックキャッシュって難しい？"

━━━ リクエスト ━━━
  プロンプト: セマンティックキャッシュって難しい？
  セッション ID: fcaf3ed5-a2d5-4934-b4b1-c84219724786

━━━ レスポンス (キャッシュヒット ✓) ━━━

# セマンティックキャッシュについて

（先ほどと同じ回答が返却される）

── メタ情報 ──
  類似度スコア: 0.9020
── 処理時間 ──
  キャッシュ検索: 143.0 ms
  合計: 143.0 ms
  埋め込み生成: 124.4 ms
  ベクトル検索: 17.2 ms

uv run invoke.py "チーズナンって美味しい？"

━━━ リクエスト ━━━
  プロンプト: チーズナンって美味しい？
  セッション ID: 0e8f9a58-d7b8-4793-aaf3-86b17fef66fb

━━━ レスポンス (キャッシュミス ✗) ━━━

# チーズナンについて

（中略）

── メタ情報 ──
  使用モデル: us.anthropic.claude-haiku-4-5-20251001-v1:0
── 処理時間 ──
  キャッシュ検索: 152.2 ms
  LLM 呼び出し: 3443.9 ms
  キャッシュ保存: 123.0 ms
  合計: 3719.4 ms