Feature Request: Add Amazon Bedrock support for indexing and retrieval

[dgallitelli]

Summary

Request to add Amazon Bedrock as a supported LLM provider for both the indexing (tree generation) and retrieval phases.

Motivation

Many enterprise users operate within AWS environments and would benefit from using Amazon Bedrock for:

• Data residency: Keep document processing within AWS regions
• Unified billing: Consolidate LLM costs under existing AWS accounts
• Model choice: Access to Claude (Anthropic), Amazon Nova, Llama, and other models via a single API
• Enterprise compliance: Leverage existing AWS security controls and IAM policies

Current State

• Upstream PageIndex only supports OpenAI for tree generation
• A community fork adds Bedrock support for queries only
• PR Feat: Add multi-provider LLM support (OpenAI + Gemini) #43 introduces a multi-provider abstraction layer (OpenAI + Gemini) that could be extended

Proposed Implementation

Extend the LLMProvider abstraction (from PR #43) to support Amazon Bedrock:

class BedrockProvider(LLMProvider):
    def __init__(self, model: str, region: str = "us-east-1"):
        self.client = boto3.client('bedrock-runtime', region_name=region)
        self.model = model

    def call(self, prompt: str) -> str:
        response = self.client.converse(
            modelId=self.model,
            messages=[{"role": "user", "content": [{"text": prompt}]}],
            inferenceConfig={"temperature": 0, "maxTokens": 4096}
        )
        return response['output']['message']['content'][0]['text']

    def call_with_finish_reason(self, prompt: str, chat_history=None) -> tuple:
        # Map Bedrock stop reasons to existing format
        # end_turn -> stop, max_tokens -> length
        ...

    async def call_async(self, prompt: str) -> str:
        # Wrap synchronous boto3 in asyncio executor
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.call, prompt)

Suggested Bedrock Models

Model	Model ID	Use Case
Claude Sonnet 4	us.anthropic.claude-sonnet-4-20250514-v1:0	High accuracy
Claude Haiku	us.anthropic.claude-haiku-4-5-20251001-v1:0	Cost-efficient
Amazon Nova Pro	us.amazon.nova-pro-v1:0	AWS-native
Amazon Nova Lite	us.amazon.nova-lite-v1:0	Fast/cheap

Usage Example

# With Bedrock
python run_pageindex.py --pdf_path doc.pdf --provider bedrock --model us.anthropic.claude-sonnet-4-20250514-v1:0

# With Bedrock (environment-based)
export PAGEINDEX_PROVIDER=bedrock
export AWS_REGION=us-east-1
python run_pageindex.py --pdf_path doc.pdf

Key Implementation Considerations

1. Dependencies: Add boto3 to requirements.txt
2. Authentication: Support IAM roles, credentials file, and environment variables
3. Stop reason mapping: Bedrock uses end_turn/max_tokens vs OpenAI's stop/length
4. Message format: Bedrock Converse API uses {"content": [{"text": "..."}]} structure
5. Async support: boto3 is synchronous; wrap in asyncio.run_in_executor()

Related

• Issue Support custom models #90: Support custom models
• Issue Ollama #27: Ollama support
• PR Feat: Add multi-provider LLM support (OpenAI + Gemini) #43: Multi-provider LLM support (OpenAI + Gemini)
• PR Feat: Cost Tracking and support for AzureOpenAI #37: Azure OpenAI support

Happy to contribute a PR if this feature is welcome!

[karunarora]

Hi, I'd like to work on this. Should I build the Bedrock provider on top of PR #43's LLMProvider abstraction, or implement it independently against main? Happy to go either way.

[KylinMountain]

Thanks for the detailed writeup! Good news — Bedrock is supported out of the box via the LiteLLM integration (#168), so a custom BedrockProvider isn't needed. Use a bedrock/... model name with standard AWS credentials:

# config.yaml (or --model)
model: "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0"

AWS auth uses the usual chain (AWS_REGION, credentials/IAM). One extra step: pip install boto3 (LiteLLM's Bedrock path needs it). Docs: https://docs.litellm.ai/docs/providers/bedrock

@karunarora so there's no need for a separate provider PR — appreciate the offer though! If you hit rough edges on the LiteLLM path (e.g. stop-reason mapping), a small fix on top is very welcome.

Closing as covered by #168.