createBedrockKBTool()

createBedrockKBTool() 函数会创建一个 Tool，用于从 Amazon Bedrock Knowledge Base 检索相关文档。它既支持托管搜索配置，也支持 Agentic Retrieval（查询分解和托管重排序），并可自动回退到标准检索。

使用示例

import { createBedrockKBTool } from '@mastra/rag'

const kbTool = createBedrockKBTool({
  knowledgeBaseId: 'YOUR_KB_ID',
  region: 'us-west-2',
  numberOfResults: 5,
  useAgenticRetrieval: true,
})

与 Agent 配合使用

import { Agent } from '@mastra/core/agent'
import { createBedrockKBTool } from '@mastra/rag'

const kbTool = createBedrockKBTool({
  knowledgeBaseId: 'YOUR_KB_ID',
})

const agent = new Agent({
  name: 'KnowledgeAssistant',
  instructions: 'Use the knowledge base tool to answer questions.',
  model: myModel,
  tools: { kb: kbTool },
})

参数

knowledgeBaseId:

string

要查询的 Amazon Bedrock Knowledge Base ID。

region?:

string

部署 Knowledge Base 的 AWS 区域。默认为 AWS_REGION 环境变量的值或 us-east-1。

numberOfResults?:

number

返回结果的最大数量。默认为 5。

useAgenticRetrieval?:

boolean

对复杂查询使用 AgenticRetrieveStream，以支持查询分解和托管重排序。失败时回退到标准 Retrieve。默认为 true（可通过 USE_AGENTIC_RETRIEVAL=false 环境变量禁用）。

userId?:

string

用于文档级访问控制的默认 AWS 用户 ID。Mastra 请求上下文中的 userId 优先级更高。

输入 Schema

由 Agent 调用时，该 Tool 接受以下输入：

queryText:

string

用于在 Knowledge Base 中查找相关文档的搜索查询。

输出 Schema

该 Tool 返回一个包含以下内容的对象：

results:

BedrockKBResult[]

检索结果数组。Bedrock 提供来源和分数时，标准检索会包含这些字段；Agentic Retrieval 可能省略这些字段。

BedrockKBResult

字段	类型	描述
content	string	检索到的段落文本内容。
source	string | undefined	Bedrock 提供的来源 URI。仅当结果元数据包含 _source_uri 时，Agentic Retrieval 才会包含此字段。
score	number | undefined	标准检索返回的相关性分数。Agentic API 不会为结果项返回分数。
metadata	Record<string, unknown>	检索结果中的其他元数据。

检索模式

Agentic Retrieval（默认）

当 useAgenticRetrieval 为 true（默认）时，该 Tool 使用 AgenticRetrieveStreamCommand，它会：

• 将复杂查询分解为子查询
• 通过多轮检索获取结果
• 应用托管重排序以获得更好的结果

如果 Agentic Retrieval 失败（例如 SDK 版本较旧或权限不足），它会自动回退到标准托管检索。

标准托管检索

当 useAgenticRetrieval 为 false 时，该 Tool 使用带有 managedSearchConfiguration 的 RetrieveCommand 进行直接的单轮检索。

基于用户的访问控制

在 Mastra 请求上下文中设置 userId，以将其作为 Bedrock userContext.userId 转发。这适用于实施文档级访问控制的 Knowledge Base。请求上下文中的值会覆盖在该 Tool 上配置的默认 userId。

import { RequestContext } from '@mastra/core/request-context'

const requestContext = new RequestContext()
requestContext.set('userId', 'user-123')

await agent.generate('Find my private documents', { requestContext })

所需 IAM 权限

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": ["bedrock:Retrieve", "bedrock:AgenticRetrieveStream"],
      "Resource": "arn:aws:bedrock:*:*:knowledge-base/*"
    }
  ]
}

SDK 要求

• @aws-sdk/client-bedrock-agent-runtime >= 3.1000（AgenticRetrieveStreamCommand 需要 ~3.1000+）