DynamoDBに追加されたベクトル検索を、Streamsでベクトル化してStrands Agentsから使ってみた

import json
import os

import boto3

TABLE_NAME = os.environ["TABLE_NAME"]
EMBED_MODEL_ID = os.environ["EMBED_MODEL_ID"]

dynamodb = boto3.client("dynamodb")
bedrock = boto3.client("bedrock-runtime")

def embed(text: str) -> list[float]:
    response = bedrock.invoke_model(
        modelId=EMBED_MODEL_ID,
        body=json.dumps({"inputText": text, "dimensions": 1024, "normalize": True}),
    )
    return json.loads(response["body"].read())["embedding"]

def handler(event, context):
    for record in event["Records"]:
        if record["eventName"] not in ("INSERT", "MODIFY"):
            continue

        new_image = record["dynamodb"].get("NewImage", {})
        old_image = record["dynamodb"].get("OldImage", {})
        text = new_image.get("text", {}).get("S")
        if not text:
            continue

        # 自分自身のembedding書き込みで再度発火した場合はスキップ（無限ループ防止）
        if "embedding" in new_image and old_image.get("text", {}).get("S") == text:
            continue

        embedding = embed(text)
        dynamodb.update_item(
            TableName=TABLE_NAME,
            Key={"doc_id": new_image["doc_id"]},
            UpdateExpression="SET embedding = :emb",
            ExpressionAttributeValues={
                ":emb": {"L": [{"N": str(v)} for v in embedding]}
            },
        )
        print(f"embedded: {new_image['doc_id']['S']} ({len(embedding)} dims)")

    return {"statusCode": 200}

import time

import boto3

TABLE_NAME = "agent-knowledge"
INDEX_NAME = "embedding-index"

dynamodb = boto3.client("dynamodb", region_name="ap-northeast-1")

def main():
    dynamodb.update_table(
        TableName=TABLE_NAME,
        AttributeDefinitions=[
            {"AttributeName": "category", "AttributeType": "S"}
        ],
        VectorIndexUpdates=[
            {
                "Create": {
                    "IndexName": INDEX_NAME,
                    "VectorAttribute": {"AttributeName": "embedding"},
                    "Dimensions": 1024,
                    "DistanceFunction": "COSINE",
                    "SearchSchema": [
                        {
                            "AttributeName": "category",
                            "SearchSchemaElementType": "INLINE_FILTER",
                        }
                    ],
                    "Projection": {"ProjectionType": "ALL"},
                }
            }
        ],
    )

    while True:
        table = dynamodb.describe_table(TableName=TABLE_NAME)["Table"]
        indexes = table.get("VectorIndexes", [])
        status = indexes[0]["IndexStatus"] if indexes else "NOT_FOUND"
        print(f"IndexStatus: {status}")
        if status == "ACTIVE":
            break
        time.sleep(10)

    print("ベクトルインデックスの作成が完了しました")

if __name__ == "__main__":
    main()

IndexStatus: CREATING
IndexStatus: CREATING
...
IndexStatus: ACTIVE
ベクトルインデックスの作成が完了しました

import json
import sys

import boto3
from strands import Agent, tool
from strands.models import BedrockModel

REGION = "ap-northeast-1"
TABLE_NAME = "agent-knowledge"
INDEX_NAME = "embedding-index"
EMBED_MODEL_ID = "amazon.titan-embed-text-v2:0"

bedrock_runtime = boto3.client("bedrock-runtime", region_name=REGION)
dynamodb = boto3.client("dynamodb", region_name=REGION)

@tool
def search_knowledge(query: str) -> list[dict]:
    """社内ナレッジベースをベクトル検索して、関連するドキュメントを返す。

    経費、勤怠、IT関連の社内規程・手続きに関する質問に回答する際に使用する。

    Args:
        query: 検索したい内容（自然言語）
    """
    response = bedrock_runtime.invoke_model(
        modelId=EMBED_MODEL_ID,
        body=json.dumps({"inputText": query, "dimensions": 1024, "normalize": True}),
    )
    embedding = json.loads(response["body"].read())["embedding"]

    results = dynamodb.search_vectors(
        TableName=TABLE_NAME,
        IndexName=INDEX_NAME,
        SearchVector=[{"N": str(v)} for v in embedding],
        TopK=3,
    )
    return [
        {
            "title": r["Item"]["title"]["S"],
            "text": r["Item"]["text"]["S"],
            "category": r["Item"]["category"]["S"],
            "score": r["Score"],
        }
        for r in results["SearchResults"]
    ]

def create_agent() -> Agent:
    return Agent(
        model=BedrockModel(
            model_id="jp.anthropic.claude-haiku-4-5-20251001-v1:0",
            region_name=REGION,
        ),
        tools=[search_knowledge],
        system_prompt=(
            "あなたは社内ヘルプデスクのアシスタントです。"
            "社内規程に関する質問には必ず search_knowledge ツールで検索し、"
            "検索結果に基づいて回答してください。"
            "検索結果にない情報は推測せず、わからないと答えてください。"
        ),
    )

if __name__ == "__main__":
    agent = create_agent()
    agent(sys.argv[1] if len(sys.argv) > 1 else "経費精算の締め日はいつ？")

from agent import create_agent
from bedrock_agentcore.runtime import BedrockAgentCoreApp

app = BedrockAgentCoreApp()

@app.entrypoint
def invoke(payload, context):
    agent = create_agent()
    result = agent(payload.get("prompt", ""))
    return {"result": str(result)}

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

FROM public.ecr.aws/docker/library/python:3.13-slim

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY agent.py main.py ./

EXPOSE 8080
CMD ["python", "main.py"]

strands-agents
bedrock-agentcore
boto3>=1.43.66

git clone https://github.com/yuu551/agentcore-strands-dynamo.git
cd agentcore-strands-dynamo
pnpm install
uv sync

npx cdk deploy

uv run python scripts/create_vector_index.py

uv run python scripts/seed_items.py

put: doc-001 経費精算の締め日
put: doc-002 出張旅費の上限
put: doc-003 リモートワーク規程
put: doc-004 有給休暇の申請方法
put: doc-005 貸与PCの交換サイクル
put: doc-006 社内Wi-Fiの接続方法

import boto3

TABLE_NAME = "agent-knowledge"

dynamodb = boto3.client("dynamodb", region_name="ap-northeast-1")

def main():
    paginator = dynamodb.get_paginator("scan")
    for page in paginator.paginate(
        TableName=TABLE_NAME, ProjectionExpression="doc_id, embedding"
    ):
        for item in sorted(page["Items"], key=lambda i: i["doc_id"]["S"]):
            emb = item.get("embedding", {}).get("L")
            status = f"あり({len(emb)}次元)" if emb else "なし"
            print(f"{item['doc_id']['S']}: embedding={status}")

if __name__ == "__main__":
    main()

uv run python scripts/check_embeddings.py

doc-001: embedding=あり(1024次元)
doc-002: embedding=あり(1024次元)
doc-003: embedding=あり(1024次元)
doc-004: embedding=あり(1024次元)
doc-005: embedding=あり(1024次元)
doc-006: embedding=あり(1024次元)

def embed(text: str) -> list[float]:
    response = bedrock_runtime.invoke_model(
        modelId=EMBED_MODEL_ID,
        body=json.dumps({"inputText": text, "dimensions": 1024, "normalize": True}),
    )
    return json.loads(response["body"].read())["embedding"]

for query in queries:
    vector = [{"N": str(v)} for v in embed(query)]
    results = dynamodb.search_vectors(
        TableName=TABLE_NAME,
        IndexName=INDEX_NAME,
        SearchVector=vector,
        TopK=2,
    )
    hits = [
        (round(r["Score"], 4), r["Item"]["title"]["S"])
        for r in results["SearchResults"]
    ]
    print(f"{query} -> {hits}")

uv run python scripts/search_test.py

経費の締め切りはいつ？ -> [(0.4065, '経費精算の締め日'), (0.9086, '出張旅費の上限')]
ホテル代はいくらまで出る？ -> [(0.6383, '出張旅費の上限'), (0.913, 'リモートワーク規程')]
パソコンが壊れたらどうすればいい？ -> [(0.6795, '貸与PCの交換サイクル'), (0.8679, 'リモートワーク規程')]

vector = [{"N": str(v)} for v in embed("社内ネットワークにつなぎたい")]
results = dynamodb.search_vectors(
    TableName=TABLE_NAME,
    IndexName=INDEX_NAME,
    SearchVector=vector,
    TopK=3,
    SearchConditionExpression="category = :c",
    ExpressionAttributeValues={":c": {"S": "IT"}},
)

--- category=IT フィルタ ---
0.5044 社内Wi-Fiの接続方法
0.8222 貸与PCの交換サイクル

import json
import sys

import boto3

client = boto3.client("bedrock-agentcore", region_name="ap-northeast-1")

def main():
    runtime_arn = sys.argv[1]
    prompt = sys.argv[2] if len(sys.argv) > 2 else "経費精算の締め日はいつ？"

    response = client.invoke_agent_runtime(
        agentRuntimeArn=runtime_arn,
        qualifier="DEFAULT",
        payload=json.dumps({"prompt": prompt}),
    )
    print(json.loads(response["response"].read())["result"])

if __name__ == "__main__":
    main()

uv run python scripts/invoke_runtime.py \
  arn:aws:bedrock-agentcore:ap-northeast-1:xxxxxxxxxxxx:runtime/knowledge_agent-xxxxxxxxxx \
  "出張のホテル代っていくらまで経費で落とせる？新幹線のグリーン車は乗っていい？"

社内規程に基づいて回答いたします。

## 出張費用について

**ホテル代（宿泊費）**
- 国内出張の場合、**1泊12,000円まで**が上限となります。

**新幹線のグリーン車**
- 新幹線は**普通車指定席まで**利用可能です。
- グリーン車の利用は認められていません。

なお、海外出張の場合は事前に部門長の承認が必要となりますので、ご注意ください。

table.put_item(Item={
    "doc_id": "doc-007",
    "category": "福利厚生",
    "title": "書籍購入補助制度",
    "text": "業務に関連する書籍は月5,000円まで会社負担で購入できます。購入申請は総務部のフォームから行い、電子書籍も対象です。",
})

uv run python scripts/invoke_runtime.py \
  arn:aws:bedrock-agentcore:ap-northeast-1:xxxxxxxxxxxx:runtime/knowledge_agent-xxxxxxxxxx \
  "技術書を会社のお金で買いたいんだけど、何か制度ある？"

はい、ございます！

**書籍購入補助制度**があります。

## 制度の内容
- **上限額**: 月5,000円まで会社負担
- **対象**: 業務に関連する書籍
- **申請方法**: 総務部のフォームから申請
- **電子書籍**: 対象に含まれます

技術書は業務に関連する書籍として対象になりますので、総務部のフォームから購入申請を行ってください。

document = {
    "doc_id": "doc-008",
    "category": "IT",
    "title": "情報セキュリティ規程",
    "chunks": [
        "パスワードは12文字以上で、英大文字・英小文字・数字・記号のうち3種類以上を...",
        "社外への資料送付は、機密区分がconfidential以上の場合、上長の承認と...",
        "セキュリティインシデントを発見した場合は、30分以内に情報システム部の...",
    ],
}

for i, chunk in enumerate(document["chunks"]):
    table.put_item(
        Item={
            "doc_id": f"{document['doc_id']}#chunk-{i}",
            "category": document["category"],
            "title": document["title"],
            "text": chunk,
        }
    )

uv run python scripts/seed_chunked_doc.py

put: doc-008#chunk-0 (97文字)
put: doc-008#chunk-1 (84文字)
put: doc-008#chunk-2 (90文字)

パスワードは何文字必要？ -> [(0.3302, 'doc-008#chunk-0'), (0.8678, 'doc-005')]
セキュリティ事故を見つけたらどこに連絡する？ -> [(0.5501, 'doc-008#chunk-2'), (0.7785, 'doc-006')]