DynamoDB를 Strands Agents의 스토리지 백엔드로 사용

Strands Agents는 Amazon Bedrock 및 기타 모델 공급자를 최우선으로 지원하며, 몇 줄의 코드로 AI 에이전트를 구축하는 모델 기반 접근 방식을 취하는 오픈 소스 SDK입니다. 구성 요소 중에는 SDK의 모든 상태 저장 하위 시스템이 사용하는 바이트 단위 계약(write, read, delete, list)인 통합 스토리지 인터페이스가 있습니다. 세션 관리자는 이를 통해 대화 스냅샷을 유지하고, 메모리 관리자는 이를 통해 장기 메모리를 저장하며, 컨텍스트 오프로더와 트랜스크립트는 동일한 작업을 사용합니다.

strands-dynamodb-storage 패키지는 Python 및 TypeScript 모두에 대해 단일 DynamoDB 테이블에서 해당 계약을 구현합니다. SDK의 /로 구분된 스토리지 키는 DynamoDB 키 모델에 직접 매핑됩니다. 키의 처음 두 세그먼트는 파티션 키가 되고 나머지는 정렬 키가 되므로, 포인트 작업은 단일 항목 호출이고 접두사 나열은 테이블 스캔이 아닌 기본 파티션 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 자체에 스토리지를 한 번 설정할 수도 있습니다. 그런 다음 스토리지를 수락하는 모든 하위 시스템은 스토리지를 상속합니다. 각 하위 시스템은 자체 키 접두사로 네임스페이스되므로 테이블 하나에 전체 에이전트의 상태가 포함됩니다.

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"},
)

의미로 호출하는 것은 동일한 파티션으로 범위가 지정된 한 번의 직접 호출입니다.

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 인수를 확인합니다. 벡터 인덱스는 테이블과 동일한 방식으로 분할되며, 모든 검색은 한 파티션으로 범위가 지정되므로 테넌트의 검색은 다른 테넌트의 메모리 범위를 벗어나지 않으며, 각 검색이 수행하는 작업은 전체 테이블이 아닌 해당 테넌트의 메모리 크기를 추적합니다. 파티션 값은 호출자가 제공하므로 이는 권한 부여 경계가 아닌 쿼리 범위 지정입니다. 테이블에 대해 dynamodb:SearchVectors를 보유한 위탁자는 모든 파티션을 검색할 수 있으며, 테넌트 액세스 제어는 IAM 및 애플리케이션 계층에 속합니다.

결과는 가장 유사한 항목을 먼저 반환합니다. 원시 점수의 방향은 인덱스의 거리 함수를 따릅니다. 코사인 및 유클리드 거리의 경우 낮을수록 더 가깝고, 내적의 경우 높을수록 더 유사합니다.

에이전트에 메모리 연결

실제 에이전트에서는 검색된 메모리가 모델에 자동으로 전달되도록 하며, SDK의 Memory Manager가 이를 처리합니다. 즉, 각 모델 직접 호출 전에 관련 항목을 검색하여 모델 입력에 통합하고, 모델이 온디맨드로 직접적으로 호출할 수 있는 search_memory 도구를 등록합니다. Memory Manager는 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
        ]

다음 코드는 세 개의 메모리를 시드하고 저장소를 에이전트에 연결하므로 검색된 메모리는 사용자 측에서 오케스트레이션 코드 없이 모델에 도달합니다.

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 오프로드 및 임베딩 모델 간접 호출에 대한 추가 문을 포함한 전체 정책이 포함되어 있습니다.

고려 사항

추가 리소스