Build agent memory with NVIDIA NeMo Agent Toolkit and Amazon S3 Vectors

import boto3

REGION = "us-west-2"
VECTOR_BUCKET = "amzn-s3-demo-research-agent-memory"
INDEX_NAME = "agent-long-term-memory"

s3vectors = boto3.client("s3vectors", region_name=REGION)

# 创建向量存储桶
s3vectors.create_vector_bucket(vectorBucketName=VECTOR_BUCKET)

# 创建索引（1024维度匹配Amazon Titan Text Embeddings V2）
s3vectors.create_index(
    vectorBucketName=VECTOR_BUCKET,
    indexName=INDEX_NAME,
    dataType="float32",
    dimension=1024,
    distanceMetric="cosine",
    metadataConfiguration={"nonFilterableMetadataKeys": ["content"]},
)

import boto3
import json
import uuid
from datetime import datetime, timezone
from nat.plugin_api import MemoryBaseConfig, MemoryEditor, MemoryItem, register_memory
from nat.builder.builder import Builder

class S3VectorsMemoryConfig(MemoryBaseConfig, name="s3vectors_memory"):
    """Amazon S3 Vectors的NAT内存提供者配置。"""
    vector_bucket: str
    index_name: str
    aws_region: str = "us-west-2"
    default_top_k: int = 5

class S3VectorsMemoryEditor(MemoryEditor):
    """由Amazon S3 Vectors支持的NAT MemoryEditor。"""

def __init__(self, config: S3VectorsMemoryConfig):
    self._vector_bucket = config.vector_bucket
    self._index_name = config.index_name
    self._default_top_k = config.default_top_k
    self._s3vectors = boto3.client('s3vectors', region_name=config.aws_region)
    self._bedrock = boto3.client('bedrock-runtime', region_name=config.aws_region)

def _get_embedding(self, text: str) -> list[float]:
    """使用Amazon Titan Text Embeddings V2生成嵌入向量。"""
    response = self._bedrock.invoke_model(
        modelId='amazon.titan-embed-text-v2:0',
        contentType='application/json',
        accept='application/json',
        body=json.dumps({
            'inputText': text,
            'dimensions': 1024,
            'normalize': True
        })
    )
    return json.loads(response['body'].read())['embedding']

async def add_items(self, items: list[MemoryItem], **kwargs) -> None:
    """将记忆项作为向量存储在Amazon S3 Vectors中。"""
    vectors = []
    for item in items:
        # 从记忆内容构建要嵌入的文本
        text = item.memory or json.dumps(item.conversation)
        embedding = self._get_embedding(text)
        # 使用uuid4防止同一秒内出现键冲突
        key = f"mem_{item.user_id}_{uuid.uuid4().hex[:12]}"
        # 注意：S3 Vectors元数据值有大小限制
        # 在生产环境中，将内容截断为1024个字符
        content_for_metadata = text[:1024]
        mem_metadata = {
            'user_id': item.user_id,
            'memory_type': item.metadata.get('memory_type', 'episodic'),
            'agent_id': item.metadata.get('agent_id', ''),
            'team_id': item.metadata.get('team_id', ''),
            'task_id': item.metadata.get('task_id', ''),
            'confidence': item.metadata.get('confidence', 0.8),
            'created_at_epoch': int(datetime.now(timezone.utc).timestamp()),
            'is_shared': item.metadata.get('is_shared', True),
            'source': item.metadata.get('source', 'agent'),
            'content': content_for_metadata,
        }
        # 如果存在，添加领域特定元数据
        if 'ticker' in item.metadata:
            mem_metadata['ticker'] = item.metadata['ticker']
        vectors.append({
            'key': key,
            'data': {'float32': embedding},
            'metadata': mem_metadata
        })
    self._s3vectors.put_vectors(
        vectorBucketName=self._vector_bucket,
        indexName=self._index_name,
        vectors=vectors
    )

    async def search(self, query: str, top_k: int = None, **kwargs) -> list[MemoryItem]:
        """从 Amazon S3 Vectors 中检索语义相关记忆"""
        query_embedding = self._get_embedding(query)
        effective_top_k = top_k or self._default_top_k
        # 从 kwargs 构建元数据过滤器
        filter_expr = {}
        for field in ('agent_id', 'memory_type', 'ticker', 'team_id', 'user_id'):
            if field in kwargs:
                filter_expr[field] = {'$eq': kwargs[field]}
        # 支持 is_shared 的布尔过滤
        if 'is_shared' in kwargs:
            filter_expr['is_shared'] = {'$eq': kwargs['is_shared']}
        response = self._s3vectors.query_vectors(
            vectorBucketName=self._vector_bucket,
            indexName=self._index_name,
            queryVector={'float32': query_embedding},
            topK=effective_top_k,
            filter=filter_expr if filter_expr else None,
            returnMetadata=True
        )
        results = []
        for vec in response.get('vectors', []):
            results.append(MemoryItem(
                conversation=[],
                tags=[vec['metadata'].get('memory_type', '')],
                metadata=vec['metadata'],
                user_id=vec['metadata'].get('user_id', ''),
                memory=vec['metadata'].get('content', '')
            ))
        return results

    async def remove_items(self, **kwargs) -> None:
        """从 Amazon S3 Vectors 中移除记忆项"""
        keys = kwargs.get('keys', [])
        if keys:
            self._s3vectors.delete_vectors(
                vectorBucketName=self._vector_bucket,
                indexName=self._index_name,
                keys=keys
            )

# 注册记忆提供者以便 NAT 可以发现它
@register_memory(config_type=S3VectorsMemoryConfig)
async def build_s3vectors_memory(config: S3VectorsMemoryConfig, builder: Builder):
    yield S3VectorsMemoryEditor(config)

# config.yml - 带 Amazon S3 Vectors 记忆的 NAT 代理配置
memory:
  agent_memory:
    _type: s3vectors_memory
    vector_bucket: "amzn-s3-demo-research-agent-memory"
    index_name: "agent-long-term-memory"
    aws_region: "us-west-2"

functions:
  add_memory:
    _type: add_memory
    memory: agent_memory
    description: |
      存储重要的发现、模式或事实到长期记忆中。
      在研究过程中发现新信息后使用此功能。
  get_memory:
    _type: get_memory
    memory: agent_memory
    description: |
      在开始研究前检索相关先验知识。
      使用当前研究主题进行查询以回忆相关发现。
  web_search:
    _type: web_search
  financial_data:
    _type: financial_data_api

workflow:
  _type: auto_memory_agent
  inner_agent_name: research_agent
  memory_name: agent_memory
  llm_name: bedrock_llm
  save_user_messages_to_memory: true
  retrieve_memory_for_every_response: true
  save_ai_messages_to_memory: true

llm:
  bedrock_llm:
    _type: bedrock
    model_id: "anthropic.claude-sonnet-4-20250514"
    temperature: 0.3

# 分析代理检索研究代理的共享发现
research_findings = await memory_client.search(
    query=f"Recent research findings for {ticker}",
    top_k=10,
    team_id='investment-research',
    is_shared=True
)

# 综合代理查询所有团队知识
all_team_knowledge = await memory_client.search(
    query=f"Complete analysis and research for {ticker} Q2 2026",
    top_k=20,
    team_id='investment-research'
)

async def consolidate_memories(ticker: str, memory_client, workflow):
    """将情景记忆提炼为持久的语义知识。"""
    episodes = await memory_client.search(
        query=f"All findings about {ticker}",
        top_k=50,
        memory_type='episodic',
        ticker=ticker
    )
    episode_texts = [e.memory for e in episodes if e.memory]
    consolidation_prompt = f"""根据以下关于{ticker}的{len(episode_texts)}个观察结果
识别持久模式和通用知识。

{chr(10).join(episode_texts)}

## 在 Amazon EKS 上部署代理工作负载

对于需要对代理工作负载进行完全操作控制的团队来说，Amazon EKS 是一个明确的选择。使用 Amazon EKS，您可以控制扩展、网络和生命周期管理，并且可以与 AWS Identity and Access Management (IAM) 本机集成以访问 Amazon S3 向量。如果您更倾向于完全托管的运行时，Amazon Bedrock AgentCore 提供了一个托管替代方案，可以消除这种操作开销。您使用 nat serve 将 NAT 代理作为容器化服务运行。每种代理类型都是一个 Kubernetes 部署，通过服务账户的 IAM 角色（IRSA）授予 IAM 访问权限。

### 步骤 1：构建代理容器

以下 Dockerfile 将代理与其配置和内存插件打包在一起：

### 步骤 2：使用 Kubernetes 清单部署

以下清单使用自动扩展部署研究代理：

### S3 向量访问的 IAM 策略

以下 IAM 策略授予对内存层的限定访问权限。通过 IRSA 将其附加到与 agent-sa 服务账户关联的 IAM 角色：

### 部署与验证

构建并推送容器镜像后，部署代理：

S3 Vectors 提供强写一致性。Pod 和代理类型可以立即看到代理 Pod 存储的记忆，无需缓存失效。HorizontalPodAutoscaler 根据 CPU 使用情况扩展代理副本。每个副本通过 IRSA 连接到相同的 S3 Vectors 索引，无论哪个 Pod 处理请求，均可实现一致的记忆访问。

## 评估记忆影响

使用 NAT 的评估工具（nat eval），可以量化记忆对代理性能的影响。使用相同的数据集配置两次评估运行：一次启用记忆，一次不启用：

基于该设计的架构特性，比较这些运行时可以预期以下定性结果。这些是方向性预期，而非基准测量：

- 可靠性提高 – 回忆的记忆为代理提供了可验证的来源上下文，代理可以引用这些内容，而不仅仅是依赖参数化知识。

- 令牌使用量减少 – 代理跳过重新推导已在先前会话中确定的事实。减少的程度取决于工作流中重复上下文的数量。

- 延迟适度增加 – 每次记忆回忆都会增加一个 S3 Vectors 查询（标准访问模式下为亚秒级），通常只占整体 LLM 推理时间的一小部分。

- 重复工作减少 – 通过共享记忆，代理可以避免重复已完成的队友的研究或分析。

这些影响的幅度取决于您的工作负载：会话之间重复的上下文数量、协调的代理数量以及检索预算（top-k）的大小。NAT 的超参数优化器可以系统地调整 top_k 和相似度阈值，以找到最适合您用例的成本-质量平衡点。

## 清理资源

为了避免本演练中创建的资源产生持续费用，请删除以下内容：

- 删除 Amazon EKS 部署及相关资源：kubectl delete namespace agent-team