Use AWS Bedrock rerank API with AI Proxy

echo '
_format_version: "3.0"
plugins:
  - name: ai-proxy
    route: rerank-route
    config:
      llm_format: bedrock
      route_type: llm/v1/chat
      logging:
        log_payloads: false
        log_statistics: true
      auth:
        allow_override: false
        aws_access_key_id: "${{ env "DECK_AWS_ACCESS_KEY_ID" }}"
        aws_secret_access_key: "${{ env "DECK_AWS_SECRET_ACCESS_KEY" }}"
      model:
        provider: bedrock
        name: "${{ env "DECK_AWS_MODEL" }}"
        options:
          bedrock:
            aws_region: "${{ env "DECK_AWS_REGION" }}"
' | deck gateway apply -

cat > bedrock-rerank-demo.py << 'EOF'
#!/usr/bin/env python3
"""Demonstrate AWS Bedrock Rerank for improving RAG retrieval quality"""

import requests
import json

RERANK_URL = "http://localhost:8000/rerank"

print("AWS Bedrock Rerank Demo: RAG Pipeline Improvement")
print("=" * 60)

# Simulate documents retrieved from vector search
query = "What are the health benefits of regular exercise?"
documents = [
    "Regular exercise can improve cardiovascular health and reduce the risk of heart disease.",
    "The Eiffel Tower was completed in 1889 and stands 324 meters tall.",
    "Exercise helps maintain healthy weight by burning calories and building muscle mass.",
    "Python is a high-level programming language known for its simplicity and readability.",
    "Physical activity strengthens bones and muscles, reducing the risk of osteoporosis and falls in older adults."
]

print(f"\nQuery: {query}")
print(f"\nCandidate documents: {len(documents)}")

# Before rerank: show original order
print("\n--- BEFORE RERANK (Original retrieval order) ---")
for idx, doc in enumerate(documents):
    print(f"{idx}. {doc[:80]}...")

# Rerank the documents
print("\n--- RERANKING ---")
try:
    # Build Bedrock rerank request
    sources = []
    for doc in documents:
        sources.append({
            "type": "INLINE",
            "inlineDocumentSource": {
                "type": "TEXT",
                "textDocument": {
                    "text": doc
                }
            }
        })

    response = requests.post(
        RERANK_URL,
        headers={"Content-Type": "application/json"},
        json={
            "queries": [
                {
                    "type": "TEXT",
                    "textQuery": {
                        "text": query
                    }
                }
            ],
            "sources": sources,
            "rerankingConfiguration": {
                "type": "BEDROCK_RERANKING_MODEL",
                "bedrockRerankingConfiguration": {
                    "numberOfResults": 3,
                    "modelConfiguration": {
                        "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0"
                    }
                }
            }
        }
    )

    response.raise_for_status()
    result = response.json()

    print("✓ Reranking complete")

    # After rerank: show reordered results
    print("\n--- AFTER RERANK (Ordered by relevance) ---")
    for item in result['results']:
        idx = item['index']
        score = item['relevanceScore']
        print(f"{idx}. [Relevance: {score:.3f}] {documents[idx][:80]}...")

    # Show the top document that should be sent to LLM
    print("\n--- TOP RESULT FOR LLM CONTEXT ---")
    top_idx = result['results'][0]['index']
    top_score = result['results'][0]['relevanceScore']
    print(f"Relevance Score: {top_score:.3f}")
    print(f"Document: {documents[top_idx]}")

except Exception as e:
    print(f"✗ Failed: {e}")

print("\n" + "=" * 60)
print("Demo complete")
EOF

python3 bedrock-rerank-demo.py

AWS Bedrock Rerank Demo: RAG Pipeline Improvement
============================================================

Query: What are the health benefits of regular exercise?

Candidate documents: 5

--- BEFORE RERANK (Original retrieval order) ---
0. Regular exercise can improve cardiovascular health and reduce the risk of hea...
1. The Eiffel Tower was completed in 1889 and stands 324 meters tall....
2. Exercise helps maintain healthy weight by burning calories and building muscl...
3. Python is a high-level programming language known for its simplicity and read...
4. Physical activity strengthens bones and muscles, reducing the risk of osteopo...

--- RERANKING ---
✓ Reranking complete

--- AFTER RERANK (Ordered by relevance) ---
0. [Relevance: 0.989] Regular exercise can improve cardiovascular health and reduce the risk of hea...
2. [Relevance: 0.876] Exercise helps maintain healthy weight by burning calories and building muscl...
4. [Relevance: 0.823] Physical activity strengthens bones and muscles, reducing the risk of osteopo...

--- TOP RESULT FOR LLM CONTEXT ---
Relevance Score: 0.989
Document: Regular exercise can improve cardiovascular health and reduce the risk of heart disease.

============================================================
Demo complete