Strands Agents 用ストレージバックエンドとしての DynamoDB の使用

Strands Agents はオープンソース SDK で、数行のコードで AI エージェントを構築するモデル駆動型のアプローチを採用しており、Amazon Bedrock やその他のモデルプロバイダーをファーストクラスでサポートしています。その構成要素の 1 つに、統合ストレージインターフェイスがあります。これは、SDK 内のすべてのステートフルサブシステムが使用するバイト指向のコントラクト (write、read、delete、list) です。セッションマネージャーはこのインターフェイスを通じて会話スナップショットを永続化し、メモリマネージャーは長期メモリを保存します。コンテキストオフローダーとトランスクリプトも同じオペレーションを使用します。

strands-dynamodb-storage パッケージは、Python と TypeScript の両方で、単一の DynamoDB テーブル上にそのコントラクトを実装します。SDK の / 区切りストレージキーは DynamoDB キーモデルに直接マッピングされます。キーの最初の 2 つのセグメントがパーティションキーになり、残りがソートキーになるため、ポイントオペレーションは単一項目の呼び出しとなり、プレフィックスのリスト取得はネイティブパーティション Query で行われ、テーブルスキャンは不要です。

主な特徴

前提条件

インストール

PyPI からパッケージをインストールします。

pip install strands-dynamodb-storage "strands-agents>=1.48.0"

または TypeScript の場合は npm からインストールします (TypeScript パッケージは機能的に同等のミラーです)。

npm install strands-dynamodb-storage

テーブルを作成する

パッケージには CreateTable アクセス許可がなく、インフラストラクチャも作成されません。事前にテーブルを作成し、独自のタグ付け、バックアップ、暗号化設定を適用します。必要なのは、文字列パーティションキー pk と文字列ソートキー sk を持つテーブルのみです。

aws dynamodb create-table \
  --table-name agent-storage \
  --attribute-definitions AttributeName=pk,AttributeType=S AttributeName=sk,AttributeType=S \
  --key-schema AttributeName=pk,KeyType=HASH AttributeName=sk,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST

セマンティックメモリを使用する場合は、同じコマンドでベクトルインデックスを宣言します。インデックスの名前、ディメンション、距離関数は作成後に変更できないため、埋め込みモデルに合わせてディメンションを設定してください。パッケージ README には、インデックスを宣言する完全な呼び出しが示されています。

エージェントセッションを永続化する

ストレージをエージェントのセッションマネージャーに渡します。SDK はキーを名前空間で管理し、呼び出しのたびに会話のスナップショットを作成し、同じセッションが再開されると復元します。

from strands import Agent
from strands.session import SnapshotSessionManager
from strands_dynamodb_storage import DynamoDBStorage

storage = DynamoDBStorage("agent-storage", region_name="us-east-1")
session = SnapshotSessionManager(session_id="user-42", storage=storage)

agent = Agent(session_manager=session)
agent("Where did we leave off?")

Agent 自体にストレージを 1 回設定することもできます。ストレージを受け入れるすべてのサブシステムはそれを継承し、それぞれが独自のキープレフィックスの下に名前空間化されるため、1 つのテーブルでエージェントの状態全体を管理できます。

from strands.vended_plugins.context_offloader import ContextOffloader

agent = Agent(
    storage=storage,
    session_manager=SnapshotSessionManager(),  # persists under session/
    plugins=[ContextOffloader()],              # offloads oversized tool results under offloader/
)

バイトコントラクトは直接利用することもできます。コントラクトは非同期です。エージェント内では SDK によって自動的に実行され、プレーンスクリプトでは asyncio.run で呼び出しをラップします。

import asyncio


async def main():
    await storage.write("session/user-42/notes", b"prefers aisle seats")
    keys = await storage.list("session/user-42/")


asyncio.run(main())

ベクトルインデックスを使用したセマンティック長期メモリ

セッション永続化でメモリ問題の半分が解決します。エージェントは再起動後も存続し、会話を再開できます。問題の残りの半分は、ユーザーが数週間前にエージェントに伝えた内容を、以前の会話とキーを共有しない新しい会話の中で思い出すことです。そのためには、キーではなく意味に基づいてメモリを検索する必要があります。DynamoDB ベクトルインデックスを使用すると、エージェントの状態を保持する同じテーブルに対して最近傍検索が可能になります (「DynamoDB でのベクトルインデックスの使用」を参照)。

次の例では、Amazon Bedrock を介して Amazon Titan Text Embeddings V2 でテキストを埋め込みます。モデルはデフォルトで 1,024 次元のベクトルを返すため、これらの例のインデックスは 1,024 次元で作成されます。次のコードは、残りの例で使用する embed() 関数を定義します。

import json

import boto3

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


def embed(text):
    response = bedrock.invoke_model(
        modelId="amazon.titan-embed-text-v2:0",
        body=json.dumps({"inputText": text, "dimensions": 1024}),
    )
    return json.loads(response["body"].read())["embedding"]

メモリを書き込むと、埋め込みとオプションのメタデータがバイトデータとともにアタッチされます。次のコードは、テナントごとのプレフィックスを使用してストアを構築するため、キー memories/m1 は物理パーティション user/u1 に配置されます。

from strands_dynamodb_storage import DynamoDBStorage, SearchQuery

storage = DynamoDBStorage("agent-storage", region_name="us-east-1", prefix="user/u1")

await storage.write(
    "memories/m1",
    b"prefers window seats on long flights",
    vector=embed("prefers window seats on long flights"),
    metadata={"kind": "preference"},
)

意味に基づく呼び出しは、同じパーティションにスコープされた 1 回の呼び出しです。

results = await storage.search(SearchQuery(
    vector=embed("what are this user's seating preferences?"),
    top_k=5,
    pk="user/u1",  # the physical partition: the full key's first two segments
    filter={"kind": "preference"},
))

pk 引数に注目してください。ベクトルインデックスはテーブルと同じ方法でパーティション分割され、すべての検索は 1 つのパーティションにスコープされるため、あるテナントの検索が別のテナントのメモリを参照することはありません。各検索が実行する処理は、テーブル全体ではなく、そのテナントのメモリのサイズを追跡します。パーティション値は呼び出し元によって指定されるため、これは認可境界ではなくクエリスコープです。テーブルに対して dynamodb:SearchVectors を保持するプリンシパルは任意のパーティションを検索できるため、テナントのアクセスコントロールは IAM とアプリケーションレイヤーで管理する必要があります。

結果は類似度の高い順に返されます。raw スコアの方向はインデックスの距離関数に従います。コサイン距離およびユークリッド距離では値が低いほど近く、ドット積では値が高いほど類似度が高くなります。

エージェントへのメモリの接続

実際のエージェントでは、取得したメモリがモデルに自動的に渡されるようにする必要があります。SDK のメモリマネージャーがこれを処理します。各モデル呼び出しの前に関連するエントリを取得してモデル入力に組み込み、モデルがオンデマンドで呼び出せる search_memory ツールを登録します。メモリマネージャーは、SDK の MemoryStore プロトコルを実装するすべてのオブジェクトを受け入れます。パッケージにはメモリマネージャーが付属していないため、書き込み時に埋め込み、検索時に埋め込み、MemoryEntry 値を返す小さなクラスを独自のアプリケーションに定義します。

import uuid

from strands.memory import MemoryEntry
from strands_dynamodb_storage import DynamoDBStorage, SearchQuery


class DynamoDBMemoryStore:
    def __init__(self, storage, partition, embed):
        self.storage = storage
        self.partition = partition
        self.embed = embed  # the embed() function defined earlier
        self.name = "dynamodb"
        self.description = "Long-term memories in DynamoDB, searched by meaning"
        self.max_search_results = 3
        self.writable = True
        self.extraction = None

    async def add(self, content, metadata=None):
        await self.storage.write(
            f"memories/{uuid.uuid4().hex[:8]}",
            content.encode(),
            vector=self.embed(content),
            metadata=metadata,
        )

    async def search(self, query, options=None):
        results = await self.storage.search(SearchQuery(
            vector=self.embed(query),
            top_k=self.max_search_results,
            pk=self.partition,
            include_values=True,
        ))
        return [
            MemoryEntry(content=r.data.decode(), metadata=r.metadata)
            for r in results
            if r.data is not None
        ]

次のコードは 3 つのメモリをシードし、ストアをエージェントに接族します。これにより、取得したメモリはオーケストレーションコードなしでモデルに到達します。

import asyncio

from strands import Agent
from strands.memory import MemoryManager

storage = DynamoDBStorage("agent-storage", region_name="us-east-1", prefix="user/u1")
store = DynamoDBMemoryStore(storage, partition="user/u1", embed=embed)


async def seed():
    await store.add("Prefers window seats on long flights")
    await store.add("Planning a trip to Tokyo in December")
    await store.add("Allergic to peanuts")


asyncio.run(seed())

memory = MemoryManager(stores=[store], add_tool_config=True)
agent = Agent(memory_manager=memory)
agent("Book me a flight seat for my December trip. Which seat should I pick?")

これをライブテーブルに対して実行すると、エージェントは search_memory ツールを呼び出し、DynamoDB から東京への旅行と窓側座席の希望を照合して、フライトの窓側座席を推奨します。add_tool_config=True を渡すと add_memory ツールも登録されるため、モデルは呼び出し元と同じテーブルを通じて新しい事実を保存できます。

必要な IAM 許可

実行時に、パッケージは 4 つの DynamoDB オペレーションを発行し、セマンティック検索を使用する場合は SearchVectors をさらに実行します。また、オフロードを設定する場合にのみ Amazon S3 オペレーションを発行するため、最小特権の IAM ポリシーは短くなります。111122223333 を AWS アカウント ID に置き換え、環境に合わせてリージョンを更新してください。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "dynamodb:PutItem",
        "dynamodb:GetItem",
        "dynamodb:DeleteItem",
        "dynamodb:Query"
      ],
      "Resource": "arn:aws:dynamodb:us-east-1:111122223333:table/agent-storage"
    }
  ]
}

パッケージ README には、セマンティック検索、Amazon S3 オフロード、埋め込みモデルの呼び出しに関する追加ステートメントを含む、完全なポリシーが記載されています。

考慮事項

その他のリソース