Uso de DynamoDB como backend de almacenamiento para Strands Agents

Strands Agents es un SDK de código abierto que adopta un enfoque basado en modelos para crear agentes de IA en unas pocas líneas de código, con compatibilidad de primera clase con Amazon Bedrock y otros proveedores de modelos. Entre sus componentes básicos se encuentra una interfaz de almacenamiento unificada: un contrato orientado a bytes (write, read, delete, list) en el que se basan todos los subsistemas con estado del SDK. El administrador de sesiones conserva las instantáneas de las conversaciones a través de él, el administrador de memoria almacena los recuerdos a largo plazo a través de él y el descargador de contexto y las transcripciones utilizan las mismas operaciones.

El paquete strands-dynamodb-storage implementa ese contrato en una sola tabla de DynamoDB, tanto para Python como para TypeScript. Las claves de almacenamiento separadas por / del SDK se asignan directamente al modelo de claves de DynamoDB: los dos primeros segmentos de una clave se convierten en la clave de partición y el resto en la clave de clasificación, por lo que las operaciones puntuales son llamadas de un solo elemento y mostrar un prefijo es una partición nativa Query, nunca un escaneo de tablas.

Características principales de

Requisitos previos

Instalación

Instale el paquete desde PyPI:

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

O desde npm para TypeScript (el paquete de TypeScript es un espejo de paridad de características):

npm install strands-dynamodb-storage

Creación de la tabla

El paquete no tiene permisos CreateTable y nunca crea infraestructura: se crea la tabla de antemano, aplicando los propios ajustes de etiquetado, copia de seguridad y cifrado. El único requisito es una tabla con una clave de partición de cadenas pk y una clave de clasificación de cadenas 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

Si planea usar la memoria semántica, declare el índice vectorial con este mismo comando: el nombre, las dimensiones y la función de distancia de un índice no se pueden cambiar después de su creación, así que ajuste las dimensiones al modelo de incrustación. El README del paquete muestra la llamada completa que declara el índice.

Sesiones persistentes de agente

Entregue el almacenamiento al administrador de sesiones del agente. El SDK asigna espacios de nombres a las claves, captura una instantánea de la conversación en cada invocación y la restaura cuando vuelve la misma sesión:

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

También puede configurar el almacenamiento una vez en el propio Agent. Cada subsistema que acepta un almacenamiento lo hereda, cada uno con su propio prefijo de clave, por lo que una tabla contiene el estado completo del 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/
)

El contrato de bytes también está disponible directamente. El contrato es asíncrono: dentro de un agente, el SDK lo controla por usted y, en un script simple, encapsule las llamadas en 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())

Memoria semántica a largo plazo con índices vectoriales

La persistencia de la sesión resuelve la mitad del problema de memoria: el agente sobrevive a un reinicio y reanuda la conversación. La parte más difícil es recordar algo que un usuario le contó al agente semanas atrás, en una nueva conversación que no comparte claves con la anterior y que requiere buscar recuerdos por significado y no por clave. Los índices vectoriales de DynamoDB llevan la búsqueda del vecino más cercano a la misma tabla que contiene el estado del agente (consulte Uso de índices vectoriales en DynamoDB).

Los ejemplos siguientes incrustan texto con Amazon Titan Text Embeddings V2 a través de Amazon Bedrock. El modelo devuelve vectores de 1024 dimensiones de forma predeterminada, por lo que el índice de estos ejemplos se crea con 1024 dimensiones. El código siguiente define la función embed() que utilizan los ejemplos 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"]

Al escribir una memoria, se adjunta una incrustación y metadatos opcionales junto con los bytes. El siguiente código crea el almacén con un prefijo por usuario, de modo que la clave memories/m1 llegue a la partición 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"},
)

Recordar por significado es una llamada, con el alcance de la misma partición:

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

Fíjese en el argumento pk. El índice vectorial está dividido de la misma manera que la tabla, y cada búsqueda se limita a una partición, de modo que la búsqueda de un inquilino nunca abarca las memorias de otro inquilino y el trabajo que realiza cada búsqueda rastrea el tamaño de la memoria de ese inquilino y no el de toda la tabla. Tenga en cuenta que el valor de la partición lo proporciona la persona que llama, por lo que se trata del alcance de la consulta y no de un límite de autorización: una entidad principal que tenga dynamodb:SearchVectors en la tabla puede buscar en cualquier partición, y el control de acceso de los inquilinos pertenece a IAM y a la capa de aplicación.

Los resultados devuelven primero los más similares. La dirección de la puntuación bruta sigue la función de distancia del índice: el valor más bajo indica mayor cercanía para la distancia coseno y la distancia euclidiana, y el valor más alto indica mayor similitud para el producto escalar.

Conexión de la memoria con un agente

En un agente real, se quiere que las memorias recuperadas lleguen al modelo automáticamente, y el administrador de memoria del SDK se encarga de ello: recupera las entradas relevantes antes de cada llamada al modelo y las incorpora en la entrada del modelo, además de registrar una herramienta search_memory a la que el modelo puede llamar bajo demanda. El administrador de memoria acepta cualquier objeto que implemente el protocolo MemoryStore del SDK. El paquete no incluye ninguno, por lo que debe definir una clase pequeña en su propia aplicación que se incruste al escribir, se incruste en la búsqueda y devuelva 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
        ]

El siguiente código crea tres memorias y conecta el almacén al agente, de forma que las memorias recuperadas lleguen al modelo sin ningún código de orquestación de su parte:

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

Al ejecutarlo contra una tabla en directo, el agente llama a su herramienta search_memory, compara el viaje a Tokio y la preferencia de asiento junto a la ventana de DynamoDB y recomienda un asiento junto a la ventana para el vuelo. Al pasar add_tool_config=True también se registra una herramienta add_memory, por lo que el modelo puede almacenar nuevos datos en la misma tabla de la que recupera datos.

Permisos de IAM necesarios

En tiempo de ejecución, el paquete emite cuatro operaciones de DynamoDB, más SearchVectors cuando se utiliza la búsqueda semántica, y operaciones de Amazon S3 solo cuando se configura la descarga, por lo que la política de IAM de privilegios mínimos es breve. Sustituya 111122223333 por el ID de Cuenta de AWS y actualice la región para que coincida con el entorno:

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

El README del paquete incluye la política completa, incluidas las instrucciones adicionales para la búsqueda semántica, la descarga de Amazon S3 y la invocación de modelos de incrustación.

Consideraciones

Recursos adicionales