Build a protein research copilot with Amazon Bedrock AgentCore

使用 Amazon Bedrock AgentCore 构建蛋白质研究助手 | 人工智能

使用 Amazon Bedrock AgentCore 构建蛋白质研究助手

蛋白质研究人员面临一个耗时的挑战：手动搜索成千上万的肽序列以找到结构相似的候选物，这个过程缓慢、容易出错，并且需要深厚的领域专业知识来解释结果。构建一个蛋白质研究助手可以改变研究人员在大型数据集中搜索结构相似肽的方式——通过一个对话式界面，实现自然语言查询、自动生成嵌入向量和 AI 驱动的结果摘要。

本文将展示如何构建一个对话式蛋白质研究助手，该助手结合了以下三种能力：

• 使用自然语言查询解析提取结构化搜索参数。

• 使用专门的语言模型对蛋白质嵌入向量进行向量相似性搜索。

• 生成搜索结果的 AI 科学摘要。

该系统使用 Strands Agents SDK 在一个代理中协调三个专用工具，部署到 Amazon Bedrock AgentCore 用于生产服务，并将肽嵌入向量存储在 Amazon Aurora PostgreSQL 兼容版中，使用 pgvector 扩展。

到本文结束时，你将构建一个端到端的代理应用程序，演示如何：

• 将自然语言用户输入（如“查找与登革热病毒肽 LPAIVREAI 相似的 10 个肽”）解析为结构化工具参数，使用 Strands Agents SDK 的工具使用模式。

• 将自定义 ML 模型（ESM-C 300M）部署为 Amazon SageMaker AI 无服务器端点，打包权重以实现快速冷启动。

• 在一个查询中结合向量相似性搜索（Amazon Aurora PostgreSQL 上的 pgvector）和元数据过滤。

• 在一个 Bedrock AgentCore 运行时中协调多个专用工具（包括嵌套的 LLM 代理），并生成搜索结果的科学摘要。

先决条件

要跟随本文的步骤，你需要：

• 一个拥有访问 Amazon Bedrock 基础模型（Anthropic Claude Sonnet 4.6）权限的 AWS 账户。

• Python 3.12 或更高版本。

• 配置了适当凭证的 AWS 命令行界面（AWS CLI）。

• Amazon Bedrock、Amazon SageMaker AI、Amazon Aurora、Amazon 弹性容器服务（Amazon ECS）和 AWS CodeBuild 的 IAM 权限。

• 安装了 bedrock-agentcore-starter-toolkit（使用 pip install bedrock-agentcore-starter-toolkit 安装）。

• IEDB 病毒表位数据集。

• 预计部署时间：30–45 分钟；查看 AWS 定价页面以了解 Bedrock、SageMaker AI、Aurora Serverless v2 和 AWS Fargate 的成本估算。

解决方案概述

该助手遵循工具使用模式，其中单个 Strands 代理协调三个专用工具，以处理完整的研究工作流程。当研究人员提交自然语言查询时，代理将其解析为结构化参数，使用蛋白质嵌入向量搜索相似肽，并使用科学背景对结果进行摘要。

以下图表说明了架构：

该架构包含五个组件：

• 在 AWS Fargate 上运行的 Streamlit 前端提供对话式界面。它将查询发送到 AgentCore 运行时，并以结构化格式显示结果，包括可下载的表格。

• 一个 Strands 代理在 Amazon Bedrock AgentCore 运行时中运行，协调工作流程。该代理通过 Bedrock Converse API 使用 Anthropic Claude Sonnet 4.6，并可以访问通过 @tool 装饰器定义的三个工具。

• 一个解析器工具，使用一个专用的 Strands 代理（LLM-as-parser 模式）从自然语言查询中提取结构化搜索参数 —— 序列、物种过滤器、结果限制。

• 一个搜索工具，通过 Amazon SageMaker AI 无服务器端点运行 ESM-C 300M 生成蛋白质嵌入，然后使用 pgvector 对 Amazon Aurora PostgreSQL 进行余弦相似性搜索。

• 一个摘要工具，使用另一个专用的 Strands 代理分析搜索结果，并生成简洁的科学摘要，同时提出进一步研究的建议。

这种单运行时、多工具的设计在保持部署简单的同时，也保持了关注点的清晰分离。每个工具封装了不同的功能，协调代理根据用户的查询决定何时以及如何调用它们。

使用 ESM-C 300M 的蛋白质嵌入

相似性搜索的核心是 ESM-C 300M，这是 EvolutionaryScale（基于 ESM 构建）开发的一种蛋白质语言模型，它生成 960 维的嵌入，能够捕捉氨基酸序列的结构和功能特性。具有相似生物功能的两个肽生成的嵌入在向量空间中是接近的，从而可以在不需要序列比对的情况下进行相似性搜索。

ESM-C 300M 部署为 Amazon SageMaker AI 无服务器端点，空闲时会自动缩放到零，并且在调用之间不会产生任何费用。模型权重被打包到部署工件中，以避免在推理时从 HuggingFace 下载 —— 这对无服务器端点至关重要，因为冷启动延迟是一个关键因素。

推理处理程序直接构建模型架构并加载预打包的权重：

from esm.models.esmc import ESMC
from esm.tokenization import get_esmc_model_tokenizers

def model_fn(model_dir):
    weights_path = os.path.join(model_dir, "weights", "esmc_300m.pt")
    model = ESMC(
        d_model=960,
        n_heads=15,
        n_layers=30,
        tokenizer=get_esmc_model_tokenizers(),
        use_flash_attn=False,
    )
    state_dict = torch.load(weights_path, map_location="cpu")
    model.load_state_dict(state_dict)
    model.eval()
    return model

predict_fn 处理程序接收蛋白质序列，对其进行编码，并返回平均池化嵌入：

def predict_fn(input_data, model):
    sequence = input_data["sequence"]
    protein = ESMProtein(sequence=sequence)
    protein_tensor = model.encode(protein)
    logits_output = model.logits(
        protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)
    )
    embeddings = logits_output.embeddings
    mean_embeddings = embeddings[:, 1:-1, :].mean(dim=1)
    return mean_embeddings[0].detach().cpu().tolist()

该端点以无服务器配置部署，内存为 6144 MB，最大并发数为 5，使用 PyTorch 2.6.0 CPU 推理容器。模型打包脚本通过 from_pretrained 一次下载权重，保存状态字典，并将其与推理代码打包到一个模型.tar.gz 文件中，该文件包含 SageMaker AI 所需的 code/ 目录结构。

使用 Aurora PostgreSQL 和 pgvector 的向量搜索

肽嵌入向量存储在 Amazon Aurora PostgreSQL 兼容版 Serverless v2 中，并使用了 pgvector 扩展。数据库结构非常简单：

CREATE TABLE peptides (
    id SERIAL PRIMARY KEY,
    sequence TEXT NOT NULL,
    embedding vector(960),
    properties JSONB,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

CREATE INDEX peptides_embedding_idx
ON peptides USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

properties JSONB 列存储了生物元数据，包括物种、来源生物体、来源分子、表位位置等，使得可以同时进行向量和元数据过滤。例如，一个查询“查找与登革热病毒中的 LPAIVREAI 相似的肽”会同时在嵌入列上进行余弦相似性搜索，并在 properties->>'species' 上进行过滤。

数据加载管道从 IEDB 病毒表位数据集读取数据，通过 SageMaker AI 端点为每个肽序列生成嵌入向量，并使用 Amazon RDS Data API 将其插入数据库。初始加载会采样 1,000 个线性肽：

def import_peptides(df):
    for i, row in tqdm(df.iterrows(), total=len(df)):
        sequence = row["Epitope_Name"]
        embedding = get_embedding(sequence)  # SageMaker AI 端点调用
        properties = {
            "species": row["Epitope_Species"],
            "source_organism": row["Epitope_Source Organism"],
            "source_molecule": row["Epitope_Source Molecule"],
            # ...其他元数据
        }
        run_statement(
            "INSERT INTO peptides (sequence, embedding, properties) "
            "VALUES (:sequence, :embedding::vector, :properties::jsonb)",
            params=[...]
        )

数据库访问通过 Amazon Relational Database Service (Amazon RDS) Data API 进行，这意味着代理运行时不需要直接连接到数据库的网络连接——它通过 HTTPS 进行通信，简化了 AgentCore 部署的网络需求。

使用 Strands Agents SDK 构建代理

Strands Agents SDK 为构建使用工具的代理提供了清晰的抽象。每个工具都是一个用 @tool 装饰的 Python 函数，代理会自动从函数的文档字符串和类型提示中为 LLM 生成工具描述。

工具定义

解析器工具委托给一个专门的 Strands 代理，该代理充当结构化输出提取器：

from strands import Agent, tool
from strands.models import BedrockModel

parser_agent = Agent(
    model=BedrockModel(model_id="us.anthropic.claude-sonnet-4-6",
                       region_name="us-east-1", streaming=False),
    system_prompt="""你是一个肽查询解析器。从自然语言查询中提取结构化搜索参数。只返回一个有效的 JSON 对象。"""
)

@tool
def parse_peptide_query(query: str) -> str:
    """将自然语言肽查询解析为结构化搜索参数。

    Args:
        query: 用户关于肽的自然语言查询。

    Returns:
        包含提取参数（如序列、物种、限制）的 JSON 字符串。
    """
    result = parser_agent(f"解析此查询: {query}")
    parsed = json.loads(str(result))
    return json.dumps(parsed)

搜索工具结合了 SageMaker AI 嵌入生成与 pgvector 相似性搜索：

@tool
def search_similar_peptides(sequence: str, species: str = "", limit: int = 20) -> str:
    """使用 ESM 嵌入查找与给定序列相似的肽。

    Args:
        sequence: 肽的氨基酸序列（例如，“LPAIVREAI”）。
        species: 可选的物种过滤器（例如，“登革热病毒”）。
        limit: 返回结果的最大数量。

    Returns:
        包含相似肽及其属性列表的 JSON 字符串。
    """
    # 从 SageMaker AI 获取嵌入
    resp = sagemaker_client.invoke_endpoint(
        EndpointName=endpoint, ContentType="application/json",
        Body=json.dumps({"sequence": sequence}))
    embedding = json.loads(resp["Body"].read().decode())["embedding"]

    # 使用可选的元数据过滤器进行向量相似性搜索
    sql = "SELECT sequence, properties, "
    sql += "(embedding <=> :query_embedding::vector) AS cosine_distance "
    sql += "FROM peptides"
    if species:
        sql += " WHERE properties->>'species' = :species"
    sql += " ORDER BY cosine_distance LIMIT :limit"

    results = run_sql(sql, params)
    return json.dumps({"results": peptides, "count": len(peptides)})

摘要工具使用另一个专门的 Strands 代理进行科学分析：

summarizer_agent = Agent(
    model=BedrockModel(model_id="us.anthropic.claude-sonnet-4-6",
                       region_name="us-east-1", streaming=False),
    system_prompt="""你是一名肽研究专家，提供简洁、高层次的摘要。分析搜索结果，并提供一个简短、有洞察力的摘要，重点关注关键发现和进一步研究的想法。"""
)

@tool
def summarize_results(original_query: str, search_results_json: str) -> str:
    """使用科学见解对肽搜索结果进行摘要。

    Args:
        original_query: 原始用户查询。
        search_results_json: 搜索结果的 JSON 字符串。

    Returns:
        搜索结果的简洁科学摘要。
    """
    results = json.loads(search_results_json)
    summary = summarizer_agent(f"原始查询: {original_query}"
                               f"结果: {results}")
    return str(summary)

协调器代理

协调器将所有内容联系在一起。它接收用户的查询并决定调用哪些工具以及调用顺序：

SYSTEM_PROMPT = """你是一名肽研究助手。你有三个工具：
1. parse_peptide_query - 将自然语言查询解析为结构化参数
2. search_similar_peptides - 使用 ESM 嵌入搜索相似的肽
3. summarize_results - 使用科学见解对搜索结果进行摘要

对于每个用户查询，请按照以下工作流程进行：
1. 首先，使用 parse_peptide_query 提取序列和参数
2. 然后，使用 search_similar_peptides 并使用提取的序列
3. 最后，使用 summarize_results 提供见解

始终完成这三个步骤。"""

strands_agent = Agent(
    model=BedrockModel(model_id="us.anthropic.claude-sonnet-4-6",
                       region_name="us-east-1", streaming=False),
    tools=[parse_peptide_query, search_similar_peptides, summarize_results],
    system_prompt=SYSTEM_PROMPT
)

此设计采用“代理作为工具”（agents-as-tools）模式：解析器和摘要器本身是 Strands 代理，但它们被 @tool 装饰器包装，并作为可调用工具暴露给协调器。协调器并不知道或关心这些工具内部使用了 LLM —— 它将它们当作函数调用。这使协调逻辑保持清晰，同时允许每个工具在需要时利用 LLM 的能力。

部署到 Amazon Bedrock AgentCore

Amazon Bedrock AgentCore 提供了一个托管运行时环境，用于托管 AI 代理。代理代码在通过 AWS CodeBuild 构建和部署的容器化环境中运行 —— 不需要本地安装 Docker。

代理入口点

AgentCore 运行时期望一个入口点函数，该函数接收一个负载和上下文：

from bedrock_agentcore.runtime import BedrockAgentCoreApp

app = BedrockAgentCoreApp()

@app.entrypoint
def invoke(payload, context):
    query = payload.get("query") or payload.get("prompt")
    result = strands_agent(query)
    return {
        "status": "success",
        "original_query": query,
        "parsed_query": _tool_outputs.get("parsed_query", {}),
        "search_results": _tool_outputs.get("search_results", []),
        "summary": _tool_outputs.get("summary", str(result)),
        "session_id": context.session_id
    }

if __name__ == '__main__':
    app.run()

入口点函数在一个共享字典中捕获工具的输出，这样响应中将包含结构化数据（解析后的查询、搜索结果表、摘要文本），而不仅仅是代理的最终文本输出。这个结构化响应是 Streamlit 前端用来渲染表格和可展开部分的内容。

基础设施即代码

部署使用 AWS CloudFormation 来管理所有基础设施。VPC 堆栈创建带有 NAT 网关和 VPC 端点的私有子网，这些端点用于 Amazon Bedrock、Amazon RDS Data API 和 AWS Secrets Manager —— 有助于确保代理运行时能够访问所有所需服务，而无需通过公共互联网。

需要使用 Amazon Aurora PostgreSQL 兼容版 Serverless v2 数据库，自动扩展范围为 0.5 到 4 ACUs（1–8 GB 内存）。一个 AWS Lambda 后端的自定义资源在堆栈创建期间初始化 pgvector 扩展并创建 peptides 表：

DBCluster:
  Type: AWS::RDS::DBCluster
  Properties:
    Engine: aurora-postgresql
    EnableHttpEndpoint: true  # Amazon RDS Data API
    ServerlessV2ScalingConfiguration:
      MinCapacity: 0.5
      MaxCapacity: 4

部署解决方案

该解决方案需要以下组件，按顺序部署：

警告：请按顺序完成部署步骤。跳过步骤可能导致部署失败。

• VPC 和网络 —— 带有 NAT 网关和 VPC 端点的私有子网，用于 Amazon Bedrock、Amazon RDS Data API 和 AWS Secrets Manager，使代理运行时能够访问所有所需服务，而无需通过公共互联网。

• Aurora PostgreSQL 数据库 —— 启用 pgvector 扩展并使用 Lambda 后端的 AWS CloudFormation 自定义资源初始化 peptides 表的 Amazon Aurora PostgreSQL 兼容版 Serverless v2 集群。

• SageMaker AI 端点 —— 使用 PyTorch 2.6.0 CPU 推理容器运行的无服务器端点，配置为 ESM-C 300M，内存为 6144 MB，最大并发数为 5。

• 肽数据 — IEDB 病毒表位数据集通过 SageMaker AI 端点为每个序列生成嵌入，并使用 Amazon RDS Data API 插入数据库中。

• AgentCore 运行时和 Streamlit 界面 — Strands 代理通过 AWS CodeBuild 部署到 Amazon Bedrock AgentCore 运行时（无需本地 Docker），Streamlit 前端则部署到 AWS Fargate。

Streamlit 前端

前端是一个轻量级的 Streamlit 应用程序，通过 bedrock-agentcore boto3 客户端与 AgentCore 运行时进行通信。它在 AWS Fargate 上运行，使用一个最小的容器镜像，仅包含 streamlit、pandas 和 boto3 — 不包含任何机器学习库。

client = boto3.client('bedrock-agentcore', region_name='us-east-1')

response = client.invoke_agent_runtime(
    agentRuntimeArn=HOST_RUNTIME_ARN,
    runtimeSessionId=session_id,
    payload=json.dumps({"prompt": query}).encode()
)

UI 在三个部分显示结果：解析后的查询参数（可展开）、一个可排序的相似肽表格（包含余弦距离和元数据），以及 AI 生成的科学摘要。用户可以将结果下载为 CSV 以进行进一步分析。

以下截图展示了搜索查询和结果。

需要考虑的问题

在将此解决方案部署到生产环境之前，请记住以下设计和操作上的权衡：

冷启动延迟。SageMaker AI 无服务器端点在空闲期后首次调用时需要 2-3 分钟，因为容器需要初始化并加载模型权重。在保持活动窗口内的后续调用可在几秒内完成。对于对延迟敏感的工作负载，考虑使用预置端点或在无服务器配置中设置更高的预置并发数。

嵌入模型选择。我们使用 ESM-C 300M，因为它在 CPU 上的嵌入质量与推理速度之间取得了良好的平衡。对于结构相似性任务的更高精度，ESM-C 600M 或 ESM2 模型提供了更大的嵌入维度，但会增加内存和延迟。ESM-C 300M 产生的 960 维嵌入在测试中对肽相似性搜索表现出色。

数据集扩展。初始加载使用了 IEDB 数据集中的 1,000 个采样肽。对于生产环境中更大的数据集，考虑批量加载嵌入、按比例增加 IVFFlat 索引列表参数，并相应地扩展 Aurora ACUs。Amazon RDS Data API 的响应大小限制为 1 MB，因此返回大型结果集的查询可能需要分页。

成本。无服务器组件（SageMaker AI 无服务器端点、Aurora Serverless v2、AgentCore 运行时）在空闲时会扩展到接近零，使这种架构对于具有间歇性使用模式的研究工作负载来说成本效益较高。在活跃使用期间的主要持续成本是 Bedrock LLM 推理（每个查询三次调用：解析器、协调器、摘要器）和 SageMaker AI 端点调用。

清理

为了避免持续收费，请按以下顺序删除资源：

警告：为避免依赖性错误，请按相反顺序删除资源。

• Streamlit 界面 — 通过 AWS CloudFormation 控制台或 AWS CLI 删除 AWS Fargate 堆栈。

• SageMaker AI 端点 — 通过 Amazon SageMaker AI 控制台或 AWS CLI 删除端点、端点配置和模型。

• 数据库 — 通过 AWS CloudFormation 控制台删除 IEDB 数据集，然后删除 Aurora PostgreSQL 数据库堆栈。

• VPC — 通过 AWS CloudFormation 控制台删除 VPC 堆栈。

• AgentCore 运行时 — 通过 Amazon Bedrock AgentCore 控制台删除运行时。

结论

本文展示了如何构建一个蛋白质研究助手，该助手通过单一的对话界面，将蛋白质语言模型嵌入与基于大语言模型（LLM）的分析相结合。

传统上，研究人员需要手动查询序列数据库、运行比对工具，并在多个应用程序中解释结果——这一过程每次搜索可能需要数小时——现在可以简化为一个自然语言查询，可在不到一分钟（或冷启动时的 2-3 分钟）内返回排序和总结后的结果。将解析、基于嵌入的搜索和科学总结整合到一个对话式工作流程中，可以显著加速肽研究和候选筛选的早期阶段。

Strands Agents SDK 的工具使用模式提供了一种清晰的方式，将解析、搜索和总结等专门能力组合成一个连贯的工作流程，而 Amazon Bedrock AgentCore 则处理代理托管和扩展的操作复杂性。

相同的架构可以推广到肽研究以外的领域。在需要在特定嵌入上进行搜索、按结构化元数据进行筛选并综合结果的领域（如基因组学、药物设计和材料科学），可以受益于这种将领域特定的嵌入模型与大语言模型（LLM）编排相结合的模式。使这种模式实用的关键设计决策包括：捆绑模型权重以避免冷启动下载、使用 Amazon RDS Data API 简化网络连接，以及通过基础设施即代码自动部署。

下一步可以考虑探索更大的 ESM 模型以提高嵌入精度、添加对批量查询的支持，或扩展元数据模式以包含来自 IEDB 数据集的更多生物注释。

参考文献

Vita R, Blazeska N, Marrama D; IEDB Curation Team Members; Duesing S, Bennett J, Greenbaum J, De Almeida Mendes M, Mahita J, Wheeler DK, Cantrell JR, Overton JA, Natale DA, Sette A, Peters B. The Immune Epitope Database (IEDB): 2024 update. Nucleic Acids Res. 2025 Jan 6;53(D1):D436-D443. doi: 10.1093/nar/gkae1092. PMID: 39558162 ; PMCID: PMC11701597 .

ESM Team. “ESM Cambrian: Revealing the mysteries of proteins with unsupervised learning.” EvolutionaryScale, 2024. https://evolutionaryscale.ai/blog/esm-cambrian

作者简介

'"`