createBedrockKBTool()

La fonction createBedrockKBTool() crée un Tool qui récupère des documents pertinents dans une base de connaissances Amazon Bedrock. Elle prend en charge à la fois la configuration de recherche gérée et la récupération agentique (décomposition des requêtes et reclassement géré), avec un repli automatique vers la récupération standard.

Exemple d'utilisation

import { createBedrockKBTool } from '@mastra/rag'

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

Avec un 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 },
})

Paramètres

knowledgeBaseId:

string

Identifiant de la base de connaissances Amazon Bedrock à interroger.

region?:

string

Région AWS dans laquelle la base de connaissances est déployée. Utilise par défaut la variable d'environnement AWS_REGION ou us-east-1.

numberOfResults?:

number

Nombre maximal de résultats à renvoyer. La valeur par défaut est 5.

useAgenticRetrieval?:

boolean

Utilise AgenticRetrieveStream pour les requêtes complexes avec décomposition des requêtes et reclassement géré. En cas d'échec, se replie sur Retrieve standard. La valeur par défaut est true (désactivez cette fonctionnalité avec la variable d'environnement USE_AGENTIC_RETRIEVAL=false).

userId?:

string

Identifiant d'utilisateur AWS par défaut pour le contrôle d'accès au niveau des documents. Un userId présent dans le contexte de requête Mastra est prioritaire.

Schéma d'entrée

Le Tool accepte l'entrée suivante lorsqu'il est appelé par un Agent :

queryText:

string

Requête de recherche permettant de trouver des documents pertinents dans la base de connaissances.

Schéma de sortie

Le Tool renvoie un objet contenant :

results:

BedrockKBResult[]

Tableau des résultats de la récupération. La récupération standard inclut la source et le score lorsque Bedrock les fournit ; la récupération agentique peut omettre ces champs.

BedrockKBResult

Champ	Type	Description
content	string	Contenu textuel du passage récupéré.
source	string | undefined	URI de la source lorsque Bedrock en fournit un. La récupération agentique inclut ce champ uniquement lorsque les métadonnées du résultat contiennent _source_uri.
score	number | undefined	Score de pertinence renvoyé par la récupération standard. L'API agentique ne renvoie pas de score pour les éléments du résultat.
metadata	Record<string, unknown>	Métadonnées supplémentaires provenant du résultat de la récupération.

Modes de récupération

Récupération agentique (par défaut)

Lorsque useAgenticRetrieval vaut true (par défaut), le Tool utilise AgenticRetrieveStreamCommand, qui :

• Décompose les requêtes complexes en sous-requêtes
• Effectue la récupération en plusieurs passes
• Applique un reclassement géré afin d'améliorer les résultats

Si la récupération agentique échoue (par exemple en raison d'une ancienne version du SDK ou des autorisations), le Tool se replie automatiquement sur la récupération gérée standard.

Récupération gérée standard

Lorsque useAgenticRetrieval vaut false, le Tool utilise RetrieveCommand avec managedSearchConfiguration pour effectuer une récupération directe en une seule passe.

Contrôle d'accès par utilisateur

Définissez userId dans le contexte de requête Mastra afin de le transmettre comme userContext.userId de Bedrock. Cette fonctionnalité permet d'utiliser les bases de connaissances qui appliquent un contrôle d'accès au niveau des documents. La valeur du contexte de requête remplace le userId par défaut configuré sur le Tool.

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

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

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

Autorisations IAM requises

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

Prérequis du SDK

• @aws-sdk/client-bedrock-agent-runtime >= 3.1000 (AgenticRetrieveStreamCommand nécessite la version ~3.1000 ou ultérieure)