S3 Vectors と Bedrock で、サーバーレスなお手軽セマンティック検索基盤を構築してみた

{"id": "abc123def456ghi789jk01", "title": "dbt platform の Advanced CI で event_time による最適化機能を試してみた", "author": "yasuhara-tomoki", "slug": "dbt-platform-advanced-ci-event-time", "published_at": 1736920200, "language": "ja", "summary": "dbt platform Advanced CI機能は、Enterprise以上のプランで利用可能なCI最適化機能です。..."}
{"id": "xyz987wvu654tsr321qp02", "title": "CloudWatch Application Signals で複数の異なるルックバックウィンドウを持つバーンレートアラームを組み合わせつつ CDK で実装してみる", "author": "masukawa-kentaro", "slug": "cloudwatch-application-signals-burn-rate-alarm-cdk", "published_at": 1733797500, "language": "ja", "summary": "CloudWatch Application Signals でバーンレートアラームを複数のルックバックウィンドウで設..."}

import json

with open('articles.jsonl') as f_in, open('embed_input.jsonl', 'w') as f_out:
    for line in f_in:
        article = json.loads(line)
        text = f"{article['title']}\n{article['summary']}"
        record = {
            'recordId': article['id'],
            'modelInput': {
                'schemaVersion': 'nova-multimodal-embed-v1',
                'taskType': 'SINGLE_EMBEDDING',
                'singleEmbeddingParams': {
                    'embeddingPurpose': 'GENERIC_INDEX',
                    'embeddingDimension': 1024,
                    'text': {'truncationMode': 'END', 'value': text}
                }
            }
        }
        f_out.write(json.dumps(record, ensure_ascii=False) + '\n')

aws s3 cp embed_input.jsonl s3://my-batch-bucket/input-rag/embed_input.jsonl

import boto3

bedrock = boto3.client('bedrock', region_name='us-east-1')

resp = bedrock.create_model_invocation_job(
    modelId='amazon.nova-2-multimodal-embeddings-v1:0',
    jobName='rag-vectorize-demo',
    roleArn='arn:aws:iam::<ACCOUNT_ID>:role/BedrockBatchInferenceRole',
    inputDataConfig={
        's3InputDataConfig': {
            's3Uri': 's3://my-batch-bucket/input-rag/embed_input.jsonl'
        }
    },
    outputDataConfig={
        's3OutputDataConfig': {
            's3Uri': 's3://my-batch-bucket/output-rag/'
        }
    }
)
job_arn = resp['jobArn']
print(f'ジョブARN: {job_arn}')

import time

while True:
    status = bedrock.get_model_invocation_job(jobIdentifier=job_arn)['status']
    print(f'status: {status}')
    if status in ('Completed', 'Failed', 'Stopped'):
        break
    time.sleep(60)

s3vectors = boto3.client('s3vectors', region_name='<YOUR_REGION>')

s3vectors.create_index(
    vectorBucketName='my-vectors-bucket',
    indexName='my-blog-index',
    dataType='float32',
    dimension=1024,
    distanceMetric='cosine',
    metadataConfiguration={
        'nonFilterableMetadataKeys': ['title']
    }
)

import json
import boto3

s3 = boto3.client('s3', region_name='us-east-1')
s3vectors = boto3.client('s3vectors', region_name='<YOUR_REGION>')

# articles.jsonl からメタデータをロード
article_meta = {}
with open('articles.jsonl') as f:
    for line in f:
        a = json.loads(line)
        article_meta[a['id']] = {
            'title': a['title'],
            'author': a['author'],
            'slug': a['slug'],
            'published_at': a['published_at'],
            'language': a['language']
        }

# バッチ出力を読み込み、メタデータと結合して登録
resp = s3.get_object(
    Bucket='my-batch-bucket',
    Key='output-rag/<JOB_ID>/embed_input.jsonl.out'
)

vectors = []
total = 0
for line in resp['Body'].iter_lines():
    result = json.loads(line)
    article_id = result['recordId']
    meta = article_meta.get(article_id, {})

    vectors.append({
        'key': article_id,
        'data': {'float32': result['modelOutput']['embeddings'][0]['embedding']},
        'metadata': {
            'title': meta.get('title', ''),
            'author': meta.get('author', ''),
            'slug': meta.get('slug', ''),
            'published_at': meta.get('published_at', 0),  # epoch秒（int）で格納
            'language': meta.get('language', '')
        }
    })

    if len(vectors) >= 100:
        s3vectors.put_vectors(
            vectorBucketName='my-vectors-bucket',
            indexName='my-blog-index',
            vectors=vectors
        )
        total += len(vectors)
        vectors = []

if vectors:
    s3vectors.put_vectors(
        vectorBucketName='my-vectors-bucket',
        indexName='my-blog-index',
        vectors=vectors
    )
    total += len(vectors)

print(f'登録完了: {total}件')

# 特定の著者の記事のみ検索
result = s3vectors.query_vectors(
    vectorBucketName=BUCKET_NAME,
    indexName=INDEX_NAME,
    queryVector={'float32': query_vector},
    topK=10,
    returnDistance=True,
    returnMetadata=True,
    filter={"author": {"$eq": "wakatsuki-ryuta"}}
)

from datetime import datetime, timezone

# 2026年以降の記事のみ
epoch_2026 = int(datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp())
result = s3vectors.query_vectors(
    vectorBucketName=BUCKET_NAME,
    indexName=INDEX_NAME,
    queryVector={'float32': query_vector},
    topK=10,
    returnDistance=True,
    returnMetadata=True,
    filter={"published_at": {"$gte": epoch_2026}}
)

# 2025年の記事のみ（範囲指定）
epoch_2025 = int(datetime(2025, 1, 1, tzinfo=timezone.utc).timestamp())
result = s3vectors.query_vectors(
    ...,
    filter={"published_at": {"$gte": epoch_2025, "$lt": epoch_2026}}
)

# リクエストボディを作成
cat > /tmp/embed_request.json << 'EOF'
{
  "schemaVersion": "nova-multimodal-embed-v1",
  "taskType": "SINGLE_EMBEDDING",
  "singleEmbeddingParams": {
    "embeddingPurpose": "GENERIC_INDEX",
    "embeddingDimension": 1024,
    "text": {
      "truncationMode": "END",
      "value": "Lambda コールドスタート 対策"
    }
  }
}
EOF

# Bedrock でベクトル化（日本語はファイル経由で渡す）
aws bedrock-runtime invoke-model \
  --model-id amazon.nova-2-multimodal-embeddings-v1:0 \
  --region us-east-1 \
  --content-type application/json \
  --body fileb:///tmp/embed_request.json \
  /tmp/query_vector.json

# ベクトル配列を抽出
python3 -c "
import json
d = json.load(open('/tmp/query_vector.json'))
print(json.dumps(d['embeddings'][0]['embedding']))
" > /tmp/query_vec_array.json

# S3 Vectors で類似検索
aws s3vectors query-vectors \
  --vector-bucket-name my-vectors-bucket \
  --index-name my-blog-index \
  --query-vector "float32=$(cat /tmp/query_vec_array.json)" \
  --top-k 5 \
  --return-distance \
  --return-metadata \
  --region <YOUR_REGION>

{
    "vectors": [
        {
            "distance": 0.1290004849433899,
            "key": "article-001",
            "metadata": {
                "title": "Lambdaのコールドスタート問題と対策について整理する",
                "published_at": 1732491853,
                "language": "ja",
                "author": "manabe-kenji",
                "slug": "lambda-coldstart-measures"
            }
        },
        {
            "distance": 0.1369316577911377,
            "key": "article-002",
            "metadata": {
                "language": "ja",
                "slug": "lambda-cold-start-avoid-hack",
                "published_at": 1559084450,
                "title": "VPC Lambdaのコールドスタートにお悩みの方へ捧ぐコールドスタート予防のハック Lambdaを定期実行するならメモリの割り当ては1600Mがオススメ？！",
                "author": "iwata-tomoya"
            }
        },
        {
            "distance": 0.14367318153381348,
            "key": "article-003",
            "metadata": {
                "title": "LambdaのProvisioned Concurrencyを使って、コールドスタート対策をしてみた #reinvent",
                "language": "ja",
                "published_at": 1576128882,
                "author": "sato-naoya",
                "slug": "lambda-provisioned-concurrency-coldstart"
            }
        },
        {
            "distance": 0.15859538316726685,
            "key": "article-004",
            "metadata": {
                "language": "ja",
                "title": "[速報]コールドスタート対策のLambda定期実行とサヨナラ！！ LambdaにProvisioned Concurrencyの設定が追加されました　 #reinvent",
                "author": "iwata-tomoya",
                "slug": "lambda-support-provisioned-concurrency",
                "published_at": 1575418456
            }
        },
        {
            "distance": 0.16454929113388062,
            "key": "article-005",
            "metadata": {
                "language": "ja",
                "published_at": 1669858986,
                "slug": "session-lamba-snapstart",
                "author": "hamada-koji",
                "title": "Lambdaのコールドスタートを解決するLambda SnapStartのセッションに参加してきた(SVS320)　#reinvent"
            }
        }
    ],
    "distanceMetric": "cosine"
}