BedrockKnowledgeBaseRetriever

Retrieve documents from a knowledge base hosted on Amazon Bedrock based on a query.

Amazon Bedrock is a fully managed service that makes state-of-the-art language models available for use through a unified API. To learn more, see Amazon Bedrock documentation.

Unlike retrievers that work against a document store you manage yourself, BedrockKnowledgeBaseRetriever queries a knowledge base that Amazon Bedrock hosts and manages for you. This lets you reuse a knowledge base you've already set up in AWS directly inside your pipeline.

Key Features

• Retrieves documents from an Amazon Bedrock-hosted knowledge base.
• Identifies the knowledge base with knowledge_base_id, or with the AWS_KNOWLEDGE_BASE_ID environment variable if you omit the ID.
• Lets you cap the number of returned documents with number_of_results, and override that limit at query time with top_k.
• Tries Amazon Bedrock's agentic retrieval first, then falls back to the standard Retrieve API if agentic retrieval isn't available.
• Returns standard Haystack Document objects with content, a relevance score, and source metadata, so you can connect the output to rankers, prompt builders, or other document-processing components like any other retriever.

Configuration

To use this component, connect Haystack Platform with Amazon Bedrock first. You'll need:

• The region name
• Access key ID
• Secret access key

For detailed explanation, see Use Amazon Bedrock and SageMaker Models.

1. Drag the BedrockKnowledgeBaseRetriever component onto the canvas from the Component Library.
2. Click the component to open the configuration panel.
3. On the General tab, set knowledge_base_id to the ID of the Amazon Bedrock knowledge base you want to query. You can find this ID in the Amazon Bedrock console.
4. Go to the Advanced tab to configure the AWS credentials, number_of_results, and use_agentic_retrieval.

Connections

BedrockKnowledgeBaseRetriever accepts a query string and an optional top_k value as inputs. It outputs a list of Document objects retrieved from the knowledge base.

Connect the pipeline's query input to the query input. Connect the documents output to a ranker, a prompt builder, or another component that processes retrieved documents.

Source Code

To check this component's source code, open knowledge_base_retriever.py in the Haystack Core Integrations repository.

Usage Example

Basic Component Configuration

  bedrock_kb_retriever:
    type: haystack_integrations.components.retrievers.amazon_bedrock.knowledge_base_retriever.AmazonBedrockKnowledgeBaseRetriever
    init_parameters:
      knowledge_base_id: YOUR_KNOWLEDGE_BASE_ID
      number_of_results: 10
      use_agentic_retrieval: true
      aws_access_key_id:
        type: env_var
        env_vars:
        - AWS_ACCESS_KEY_ID
        strict: false
      aws_secret_access_key:
        type: env_var
        env_vars:
        - AWS_SECRET_ACCESS_KEY
        strict: false
      aws_session_token:
        type: env_var
        env_vars:
        - AWS_SESSION_TOKEN
        strict: false
      aws_region_name:
        type: env_var
        env_vars:
        - AWS_DEFAULT_REGION
        strict: false
      aws_profile_name:
        type: env_var
        env_vars:
        - AWS_PROFILE
        strict: false

Using the Component in a Pipeline

This is an example of a document search pipeline that uses BedrockKnowledgeBaseRetriever to fetch documents from a knowledge base and rank them:

# haystack-pipeline
components:
  bedrock_kb_retriever: # Retrieves documents from an Amazon Bedrock knowledge base
    type: haystack_integrations.components.retrievers.amazon_bedrock.knowledge_base_retriever.AmazonBedrockKnowledgeBaseRetriever
    init_parameters:
      knowledge_base_id: YOUR_KNOWLEDGE_BASE_ID
      number_of_results: 20
      use_agentic_retrieval: true
      aws_access_key_id:
        type: env_var
        env_vars:
        - AWS_ACCESS_KEY_ID
        strict: false
      aws_secret_access_key:
        type: env_var
        env_vars:
        - AWS_SECRET_ACCESS_KEY
        strict: false
      aws_session_token:
        type: env_var
        env_vars:
        - AWS_SESSION_TOKEN
        strict: false
      aws_region_name:
        type: env_var
        env_vars:
        - AWS_DEFAULT_REGION
        strict: false
      aws_profile_name:
        type: env_var
        env_vars:
        - AWS_PROFILE
        strict: false

  AmazonBedrockRanker:
    type: haystack_integrations.components.rankers.amazon_bedrock.ranker.AmazonBedrockRanker
    init_parameters:
      model: cohere.rerank-v3-5:0
      top_k: 10
      aws_access_key_id:
        type: env_var
        env_vars:
        - AWS_ACCESS_KEY_ID
        strict: false
      aws_secret_access_key:
        type: env_var
        env_vars:
        - AWS_SECRET_ACCESS_KEY
        strict: false
      aws_session_token:
        type: env_var
        env_vars:
        - AWS_SESSION_TOKEN
        strict: false
      aws_region_name:
        type: env_var
        env_vars:
        - AWS_DEFAULT_REGION
        strict: false
      aws_profile_name:
        type: env_var
        env_vars:
        - AWS_PROFILE
        strict: false

connections:  # Defines how the components are connected
- sender: bedrock_kb_retriever.documents
  receiver: AmazonBedrockRanker.documents

inputs:  # Define the inputs for your pipeline
  query:  # These components will receive the query as input
  - "bedrock_kb_retriever.query"
  - "AmazonBedrockRanker.query"

outputs:  # Defines the output of your pipeline
  documents: "AmazonBedrockRanker.documents" # The output of the pipeline is the ranked, retrieved documents

max_runs_per_component: 100

metadata: {}

Parameters

Inputs

Parameter	Type	Description
query	str	The search query used to retrieve documents from the knowledge base.
top_k	Optional[int]	Maximum number of results. Overrides number_of_results if you pass it at query time.

Outputs

Parameter	Type	Description
documents	List[Document]	Documents retrieved from the Amazon Bedrock knowledge base. Each document includes content, a relevance score, and metadata with source, knowledge_base_id, and knowledge_base_type. source is the origin of the retrieved content, for example, an S3 URI, a web URL, or a Confluence, Salesforce, SharePoint, or custom document location.

Init Parameters

These are the parameters you can configure in Pipeline Builder:

Parameter	Type	Default	Description
knowledge_base_id	Optional[str]	None	The ID of the Amazon Bedrock knowledge base to query. If you omit this, the retriever uses the AWS_KNOWLEDGE_BASE_ID environment variable.
aws_access_key_id	Optional[Secret]	Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False)	AWS access key ID.
aws_secret_access_key	Optional[Secret]	Secret.from_env_var(["AWS_SECRET_ACCESS_KEY"], strict=False)	AWS secret access key.
aws_session_token	Optional[Secret]	Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False)	AWS session token.
aws_region_name	Optional[Secret]	Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False)	AWS region name. You can also pass this as a plain string.
aws_profile_name	Optional[Secret]	Secret.from_env_var(["AWS_PROFILE"], strict=False)	AWS profile name.
number_of_results	int	5	The default maximum number of documents to return. You can override this at query time with top_k.
use_agentic_retrieval	Optional[bool]	None	If True, the retriever first tries Amazon Bedrock's agentic retrieval, then falls back to the standard Retrieve API if agentic retrieval isn't available. Defaults to True unless the USE_AGENTIC_RETRIEVAL environment variable is set to false.

Run Method Parameters

These are the parameters you can configure for the component's run() method. This means you can pass these parameters at query time through the API, in Playground, or when running a job. For details, see Modify Pipeline Parameters at Query Time.

Parameter	Type	Default	Description
query	str		The search query used to retrieve documents from the knowledge base.
top_k	Optional[int]	None	The maximum number of documents you want the retriever to return. Overrides the init-time number_of_results value.

Related Information

• Use Amazon Bedrock and SageMaker Models
• Add Integrations
• AmazonBedrockRanker