Amazon S3 Vectors (3) : PDFに対するベクトル検索の実装とデータの一括削除サンプルスクリプト

import boto3

s3vectors = boto3.client("s3vectors", region_name="us-east-1")

bucket_name = "vectorbucket"
index_name = "vectorindex"

# 1. ベクトルキーの一覧を取得
paginator = s3vectors.get_paginator("list_vectors")
pages = paginator.paginate(vectorBucketName=bucket_name, indexName=index_name)

all_keys = []
for page in pages:
    for vector in page.get("vectors", []):
        all_keys.append(vector["key"])

# 2. 各キーを1件ずつ削除（keys=[key] の形式）
for key in all_keys:
    print(f"Deleting vector: {key}")
    s3vectors.delete_vectors(
        vectorBucketName=bucket_name,
        indexName=index_name,
        keys=[key]  
    )

all_keys = []
for page in pages:
    for vector in page.get("vectors", []):
        all_keys.append(vector["key"])

    s3vectors.delete_vectors(
        vectorBucketName=bucket_name,
        indexName=index_name,
        keys=[key]  
    )

import boto3
import json
from PyPDF2 import PdfReader

# ===== 設定 =====
region = "us-east-1"
pdf_path = "example.pdf"
vector_bucket_name = "vectorbucket"
vector_index_name = "vectorindex"
chunk_size = 500
max_metadata_bytes = 1800  # UTF-8バイトでの制限（安全マージン付き）

def split_text_by_length(text, max_chars=2000):
    return [text[i:i+max_chars] for i in range(0, len(text), max_chars)]

def trim_to_max_bytes(s, max_bytes):
    """指定されたUTF-8バイト数以内に文字列を収める"""
    encoded = s.encode("utf-8")
    if len(encoded) <= max_bytes:
        return s
    # 切り詰め
    trimmed = encoded[:max_bytes]
    # バイト列が途中のマルチバイト文字で終わっていないか確認し、復元
    while True:
        try:
            return trimmed.decode("utf-8")
        except UnicodeDecodeError:
            trimmed = trimmed[:-1]

def extract_chunks_from_pdf(pdf_path, chunk_size=2000):
    reader = PdfReader(pdf_path)
    chunks = []
    for page_number, page in enumerate(reader.pages, start=1):
        text = page.extract_text()
        if text:
            split_chunks = split_text_by_length(text.strip(), chunk_size)
            for chunk_index, chunk in enumerate(split_chunks):
                chunks.append({
                    "text": chunk,
                    "page": page_number,
                    "chunk": chunk_index + 1
                })
    return chunks

# ===== クライアント初期化 =====
bedrock = boto3.client("bedrock-runtime", region_name=region)
s3vectors = boto3.client("s3vectors", region_name=region)

# ===== チャンク抽出 =====
chunks = extract_chunks_from_pdf(pdf_path, chunk_size)

# ===== ベクトル化 & 登録 =====
vectors = []
for i, chunk in enumerate(chunks):
    response = bedrock.invoke_model(
        modelId="amazon.titan-embed-text-v2:0",
        body=json.dumps({"inputText": chunk["text"]})
    )

    response_body = json.loads(response["body"].read())
    embedding = response_body["embedding"]

    source_text_trimmed = trim_to_max_bytes(chunk["text"], max_metadata_bytes)

    vector = {
        "key": f"page-{chunk['page']}-chunk-{chunk['chunk']}",
        "data": {"float32": embedding},
        "metadata": {
            "source_text": source_text_trimmed,
            "page_number": chunk["page"],
            "chunk_number": chunk["chunk"]
        }
    }
    vectors.append(vector)

# ===== S3 Vectors に登録 =====
s3vectors.put_vectors(
    vectorBucketName=vector_bucket_name,
    indexName=vector_index_name,
    vectors=vectors
)

print(f"{len(vectors)} 個のベクトルを S3 Vectors に登録しました。")

chunk_size = 500
max_metadata_bytes = 1800  # UTF-8バイトでの制限（安全マージン付き）

import boto3
import json

# ===== 設定 =====
region = "us-east-1"
vector_bucket_name = "vectorbucket"
vector_index_name = "vectorindex"
input_text = "個人情報とは？"

# ===== クライアント初期化 =====
bedrock = boto3.client("bedrock-runtime", region_name=region)
s3vectors = boto3.client("s3vectors", region_name=region)

# ===== 入力テキストをベクトルに変換 =====
response = bedrock.invoke_model(
    modelId="amazon.titan-embed-text-v2:0",
    body=json.dumps({"inputText": input_text})
)
embedding = json.loads(response["body"].read())["embedding"]

# ===== ベクトル検索実行 =====
response = s3vectors.query_vectors(
    vectorBucketName=vector_bucket_name,
    indexName=vector_index_name,
    queryVector={"float32": embedding},
    topK=1,
    returnDistance=True,
    returnMetadata=True
)

# ===== 結果を日本語で正しく表示 =====
print(json.dumps(response["vectors"], indent=2, ensure_ascii=False))