Usar o DynamoDB como backend de armazenamento para o Strands Agents

O Strands Agents é um SDK de código aberto que adota uma abordagem orientada por modelos para criar agentes de IA com poucas linhas de código, com suporte de primeira classe ao Amazon Bedrock e a outros provedores de modelos. Entre seus componentes básicos está uma interface de armazenamento unificada: um contrato orientado a bytes (write, read, delete, list) com o qual todos os subsistemas com estado do SDK trabalham. O Gerenciador de Sessões persiste snapshots de conversas por meio dela, o Gerenciador de Memória armazena memórias de longo prazo por meio dela, e o descarregador de contexto e as transcrições usam as mesmas operações.

O pacote strands-dynamodb-storage implementa esse contrato em uma única tabela do DynamoDB, para Python e TypeScript. As chaves de armazenamento separadas por / do SDK são mapeadas diretamente para o modelo de chaves do DynamoDB: os dois primeiros segmentos de uma chave tornam-se a chave de partição e o restante torna-se a chave de classificação. Assim, as operações pontuais são chamadas para um único item e a listagem de um prefixo é uma Query de partição nativa, e nunca uma varredura da tabela.

Recursos principais

Pré-requisitos

Instalação

Instale o pacote do PyPI:

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

Ou do npm para TypeScript (o pacote TypeScript é um espelho com paridade de recursos):

npm install strands-dynamodb-storage

Criar a tabela

O pacote não tem a permissão CreateTable e nunca cria infraestrutura: crie a tabela antecipadamente, aplicando as próprias configurações de tags, backup e criptografia. Uma tabela com uma chave de partição de string pk e uma chave de classificação de string sk é o único requisito:

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

Se você planeja usar memória semântica, declare o índice vetorial nesse mesmo comando: o nome do índice, as dimensões e a função de distância não podem ser alterados após a criação, portanto dimensione as dimensões de acordo com o modelo de incorporação. O README do pacote mostra a chamada completa que declara o índice.

Persistir sessões do agente

Forneça o armazenamento ao gerenciador de sessões do agente. O SDK cria namespaces para as chaves, captura um snapshot da conversa a cada invocação e o restaura quando a mesma sessão retorna:

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?")

Também é possível definir o armazenamento uma única vez no próprio Agent. Todo subsistema que aceita um armazenamento herda essa configuração, cada um com namespace próprio sob seu próprio prefixo de chave, de modo que uma tabela armazene todo o estado do agente:

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/
)

O contrato orientado a bytes também está disponível diretamente. O contrato é assíncrono: dentro de um agente, o SDK o gerencia para você e, em um script simples, envolva as chamadas em 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())

Memória semântica de longo prazo com índices vetoriais

A persistência de sessões resolve metade do problema de memória: o agente sobrevive a uma reinicialização e retoma a conversa. A parte mais difícil é recuperar algo que um usuário disse ao agente semanas atrás, em uma nova conversa que não compartilha nenhuma chave com a anterior, o que exige pesquisar memórias por significado, e não por chave. Os índices vetoriais do DynamoDB trazem a pesquisa de vizinhos mais próximos para a mesma tabela que armazena o estado do agente (consulte Uso de índices vetoriais no DynamoDB).

Os exemplos a seguir geram incorporações do texto com o Incorporador de Texto do Amazon Titan V2 por meio do Amazon Bedrock. O modelo retorna vetores de 1.024 dimensões por padrão, portanto o índice desses exemplos é criado com 1.024 dimensões. O código a seguir define a função embed() usada pelos exemplos restantes:

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"]

A gravação de uma memória associa uma incorporação e metadados opcionais aos bytes. O código a seguir cria o armazenamento com um prefixo por locatário, de modo que a chave memories/m1 seja armazenada na partição física 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"},
)

A recuperação por significado é feita com uma única chamada, com escopo para a mesma partição:

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

Observe o argumento pk. O índice vetorial é particionado da mesma forma que a tabela, e toda pesquisa tem escopo definido para uma partição. Portanto, a pesquisa de um locatário nunca abrange as memórias de outro locatário, e o trabalho realizado em cada pesquisa acompanha o tamanho da memória desse locatário, e não o da tabela inteira. Lembre-se de que o valor da partição é fornecido pelo chamador, portanto isso define o escopo da consulta, e não uma fronteira de autorização: uma entidade principal com dynamodb:SearchVectors na tabela pode pesquisar qualquer partição, e o controle de acesso do locatário deve ser implementado no IAM e na camada da aplicação.

Os resultados retornam primeiro os mais semelhantes. A direção da pontuação bruta segue a função de distância do índice: valores menores indicam maior proximidade para distância cosseno e distância euclidiana, enquanto valores maiores indicam maior similaridade para produto escalar.

Integrar a memória a um agente

Em um agente real, você quer que as memórias recuperadas cheguem automaticamente ao modelo, e o Gerenciador de Memória do SDK cuida disso: ele recupera entradas relevantes antes de cada chamada ao modelo e as incorpora à entrada do modelo, além de registrar uma ferramenta search_memory que o modelo pode chamar sob demanda. O Gerenciador de Memória aceita qualquer objeto que implemente o protocolo MemoryStore do SDK. O pacote não fornece um desses objetos, portanto defina uma classe pequena na própria aplicação que gere uma incorporação na gravação, gere uma incorporação na pesquisa e retorne valores 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
        ]

O código a seguir inicializa três memórias e integra o armazenamento ao agente, de modo que as memórias recuperadas cheguem ao modelo sem código de orquestração no lado da aplicação:

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?")

Ao executar isso em uma tabela ativa, o agente chama sua ferramenta search_memory, encontra a viagem a Tóquio e a preferência por assento na janela no DynamoDB e recomenda um assento na janela para o voo. Passar add_tool_config=True também registra uma ferramenta add_memory, permitindo que o modelo armazene novos fatos na mesma tabela da qual recupera informações.

Permissões obrigatórias do IAM

No runtime, o pacote emite quatro operações do DynamoDB, além de SearchVectors quando você usa a pesquisa semântica, e operações do Amazon S3 somente quando configura o descarregamento, portanto, a política do IAM com privilégio mínimo é curta. Substitua 111122223333 pelo ID da Conta da AWS e atualize a região para corresponder ao ambiente:

{
  "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"
    }
  ]
}

O README do pacote contém a política completa, incluindo as declarações adicionais para pesquisa semântica, descarregamento para o Amazon S3 e invocação do modelo de incorporação.

Considerações

Recursos adicionais