RAG Querying Workflows

Welcome to the final lesson of Managing Data for GenAI with Bedrock Knowledge Bases! Having successfully built your vector storage infrastructure, populated it with document embeddings, and created a Bedrock Knowledge Base, you have now arrived at the culmination of your GenAI data management journey: implementing powerful querying capabilities that transform your stored knowledge into actionable intelligence.

In this fourth and final lesson, we'll explore two distinct yet complementary approaches to retrieving information from your carefully constructed data pipeline. You'll master both direct vector similarity search using S3 Vectors for precise mathematical matching and the complete RAG (Retrieval-Augmented Generation) workflow, which combines intelligent document retrieval with response generation. By comparing these approaches side by side, you'll develop a nuanced understanding of when to use each method and how they work together to create truly intelligent question-answering systems that can provide both precise document matches and contextually rich, human-readable responses.

Let's begin by understanding the fundamental distinction between direct vector similarity search and full RAG workflows. These two approaches serve different purposes in the GenAI ecosystem and offer complementary strengths for various use cases.

Direct vector similarity search operates at the mathematical level, transforming your query into an embedding vector and finding the most similar document vectors through geometric distance calculations. This approach excels when you need precise, deterministic matches and want to examine the raw similarity scores between your query and stored documents. It's particularly valuable for debugging your vector index, analyzing the quality of your embeddings, or building custom retrieval logic where you need direct access to similarity metrics.

Full RAG workflows, in contrast, orchestrate multiple AI operations to provide complete question-answering experiences. The RAG process begins with the same vector similarity search but then feeds the retrieved documents into a large language model that synthesizes the information into coherent, contextual responses. This approach transforms raw document matches into human-readable answers complete with proper citations, making it ideal for end-user applications, customer support systems, or any scenario where you need intelligent, conversational responses rather than just document references.

It's important to note that RAG workflows are significantly more expensive than direct vector search due to the additional large language model inference required for response generation. While direct search only incurs the computational cost of vector similarity calculations, RAG workflows also consume substantial tokens for both input context (retrieved documents) and output generation (synthesized responses), making cost optimization an important consideration when choosing between these approaches for your specific use case.

Our implementation begins by establishing the necessary AWS service clients and defining the configuration parameters that will govern both our direct search and RAG operations, like we did in previous lessons:

import os
import json
import boto3

# Create AWS clients
s3_vectors_client = boto3.client("s3vectors")
bedrock_runtime_client = boto3.client("bedrock-runtime")
bedrock_agent_runtime_client = boto3.client("bedrock-agent-runtime")

# AWS region name
REGION_NAME = os.getenv("AWS_REGION", "us-east-1")

# Knowledge Base ID
KNOWLEDGE_BASE_ID = os.getenv("KNOWLEDGE_BASE_ID")

# S3 Vectors configuration
VECTOR_BUCKET_NAME = "bedrock-vector-bucket"
VECTOR_INDEX_NAME = "bedrock-vector-index"

# Set the embedding model ID to use
EMBEDDING_MODEL_ID = "amazon.titan-embed-text-v2:0"

# Set the number of dimensions for the embeddings
EMBEDDING_DIMENSIONS = 1024

# RAG generation model configuration
GENERATION_MODEL_ID = "amazon.nova-pro-v1:0"

This configuration establishes three distinct AWS service clients, each serving a specific role in our querying architecture. The s3_vectors_client provides direct access to your vector index for mathematical similarity searches, the bedrock_runtime_client handles embedding generation to transform text queries into searchable vectors, and the bedrock_agent_runtime_client orchestrates the complete RAG workflow by coordinating retrieval and generation operations. The model specifications ensure consistency between our query processing and the original document embeddings: EMBEDDING_MODEL_ID must match exactly what we used when populating our vector index, while GENERATION_MODEL_ID specifies the large language model that will synthesize retrieved information into coherent responses.

To perform direct vector similarity search, we first need to transform our text query into the same vector representation used for our stored documents:

# Test 1: Direct vector search using S3 Vectors
test_query = "What is Nimbus Assist?"

# Generate embedding for the test query
query_embedding_request = {
    "inputText": test_query,             # The query text to embed
    "dimensions": EMBEDDING_DIMENSIONS,  # Number of dimensions for the embedding vector
    "normalize": True,                   # Normalize the embedding vector (recommended for similarity search)
}

# Get embedding from Bedrock for the query
query_response = bedrock_runtime_client.invoke_model(
    modelId=EMBEDDING_MODEL_ID,
    body=json.dumps(query_embedding_request)
)

# Load the response body
query_response_body = json.loads(query_response["body"].read())

# Extract the embedding vector from the response
query_embedding = [float(x) for x in query_response_body["embedding"]]

The embedding generation process mirrors exactly what we performed when populating our vector index in lesson two, ensuring mathematical consistency between our query vectors and stored document vectors. The "normalize": True parameter is particularly important for similarity search operations, as it standardizes vector magnitudes to focus purely on directional relationships rather than magnitude differences. The response processing involves JSON parsing to extract the raw numerical vector, which we convert to a list of Python floats for compatibility with the S3 Vectors API. This query embedding becomes the mathematical representation that will be compared against all stored document vectors to identify the most semantically similar content.

With our query vector prepared, we can now execute the direct similarity search against our S3 Vectors index:

# Search for similar vectors in S3 Vectors index
search_response = s3_vectors_client.query_vectors(
    vectorBucketName=VECTOR_BUCKET_NAME,       # S3 Vectors bucket name
    indexName=VECTOR_INDEX_NAME,               # S3 Vectors index name
    queryVector={"float32": query_embedding},  # Query vector to search for
    topK=3,                                    # Number of top results to return
    returnDistance=True,                       # Include distance in response
)

# Display direct search results
print(f"Direct search results for '{test_query}':")
for i, result in enumerate(search_response.get("vectors", []), 1):
    distance = result.get("distance")
    if distance is not None:
        print(f"  {i}. Key: {result['key']}, Distance: {distance:.4f}")
    else:
        print(f"  {i}. Key: {result['key']}")

The query_vectors operation performs the core mathematical computation that drives semantic search, calculating similarity scores between your query embedding and every vector in your index. The topK parameter controls result quantity, returning only the three most similar documents rather than exhaustive results, while returnDistance=True provides the numerical similarity scores that quantify how closely each result matches your query. The results processing demonstrates how to safely extract information from the S3 Vectors response structure, using the get() method to handle potential missing fields gracefully.

When executed with our test query, the direct search produces output that reveals the mathematical relationships within your vector space:

Direct search results for 'What is Nimbus Assist?':
  1. Key: brd-nimbus-assist, Distance: 0.4789
  2. Key: prd-nimbus-assist-v1, Distance: 0.5374
  3. Key: design-rag-nimbus-assist, Distance: 0.5871

These distance scores provide valuable insights into your vector index's behavior: smaller distances indicate higher semantic similarity, with the business requirements document showing the closest match at 0.4789 distance, followed by the product requirements at 0.5374, and the design document at 0.5871. The consistent ranking demonstrates that our vector embeddings successfully capture semantic relationships, correctly identifying documents about Nimbus Assist as most relevant to queries about this topic.

Moving beyond direct search, we now implement the complete RAG workflow that transforms document retrieval into intelligent question-answering:

# Test 2: Full RAG workflow using Bedrock Knowledge Base
rag_queries = [
    "What is Nimbus Assist?",
    "In Nimbus Assist's RAG pipeline, what are the default retrieval parameters and which metadata filters are applied by default?",
    "What is Tech Company Inc.'s paid time off (PTO) policy?",
]

# Process each test query through the full RAG pipeline
for query in rag_queries:
    print(f"\nQuery: {query}")
    
    # Use Bedrock Knowledge Base for retrieval and generation
    response = bedrock_agent_runtime_client.retrieve_and_generate(
        input={"text": query},
        retrieveAndGenerateConfiguration={
            "type": "KNOWLEDGE_BASE",
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": KNOWLEDGE_BASE_ID,
                "modelArn": (
                    f"arn:aws:bedrock:{REGION_NAME}::foundation-model/{GENERATION_MODEL_ID}"
                ),
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "numberOfResults": 3
                    }
                },
                "generationConfiguration": {
                    "inferenceConfig": {
                        "textInferenceConfig": {
                            "maxTokens": 500,
                            "temperature": 0.1,
                        }
                    },
                },
            },
        },
    )

The retrieve_and_generate operation represents the sophisticated orchestration that makes RAG workflows so powerful. This single API call coordinates multiple AI operations: it embeds your query using the same model configured in your knowledge base, performs vector similarity search to identify relevant documents, ranks and selects the most appropriate content, and feeds this context into a large language model for response generation.

The input query is represented as a dict containing the text key. The retrieveAndGenerateConfiguration parameter provides a hierarchical structure that organizes all workflow settings:

• Top-level configuration: Specifies "type": "KNOWLEDGE_BASE" to indicate we're using a Bedrock Knowledge Base as our source.
• Knowledge base settings: Links to your specific knowledge base via knowledgeBaseId and defines which generation model to use via modelArn.
• Retrieval configuration: Controls the document search phase through vectorSearchConfiguration, where "numberOfResults": 3 specifies how many relevant documents should be retrieved as context.
• Generation configuration: Manages the response synthesis phase through nested inference settings, where "maxTokens": 500 limits response length for concise answers and "temperature": 0.1 keeps the model focused on factual accuracy by reducing creative variations.

This structured approach ensures that each stage of the RAG pipeline operates with precisely defined parameters while maintaining clean separation between retrieval and generation concerns.

After the RAG workflow completes, we need to extract and display both the generated answer and the source document citations:

    # Extract and display the generated answer
    answer = response["output"]["text"]
    print(f"Answer: {answer}")

    # Display source documents used for generation
    citations = response.get("citations", [])
    if citations:
        print("Sources:")
        for citation in citations:
            for ref in citation.get("retrievedReferences", []):
                metadata = ref.get("metadata", {})
                source_key = metadata.get("x-amz-bedrock-kb-source-uri", metadata.get("filename", "Unknown"))
                print(f"  - {source_key}")
    else:
        print("Sources: (none returned)")

This response processing code demonstrates proper extraction of both the generated text and the critical citation information that makes RAG responses trustworthy and verifiable. The citation processing implements a fallback hierarchy for source identification: it first looks for x-amz-bedrock-kb-source-uri, which provides the full source URI, then falls back to filename from the metadata, and finally defaults to "Unknown" if neither is available, ensuring robust handling of various document metadata configurations.

When executed, our RAG workflow produces comprehensive results that showcase the system's ability to handle diverse query types and provide detailed, cited responses:

Query: What is Nimbus Assist?
Answer: Nimbus Assist is an AI-powered support agent designed to reduce human ticket volume by providing fast, accurate, and citeable answers to customer questions using curated internal knowledge. It is intended to deliver answers with top 1-3 citations, ensuring compliance and traceability. The system uses retrieval-augmented generation (RAG) with Bedrock Knowledge Bases, supports multilingual capabilities (primarily English with fallbacks for Spanish and German), and includes guardrails for handling personally identifiable information (PII) and unsupported topics. Nimbus Assist is designed to integrate with various channels such as a web widget and Zendesk for agent assistance, aiming to deflect a significant percentage of Tier-0/Tier-1 topics and maintain high answer helpfulness ratings.
Sources:
  - brd-nimbus-assist.md
  - brd-nimbus-assist.md
  - brd-nimbus-assist.md
  - brd-nimbus-assist.md
  - prd-nimbus-assist-v1.md

Query: In Nimbus Assist's RAG pipeline, what are the default retrieval parameters and which metadata filters are applied by default?
Answer: Based on the retrieved results, the default retrieval parameters and metadata filters in Nimbus Assist's RAG pipeline are as follows:

- **Default Retrieval Parameters:**
  - Chunk size: 1200–1800 tokens, with an overlap of 150 tokens.
  - Top-k: 8, using a hybrid rerank by score and recency (with a half-life of 30 days).

- **Default Metadata Filters:**
  - `product`
  - `version`
  - `complianceLevel` These parameters and filters are used to optimize the retrieval process for relevant and up-to-date information.
Sources:
  - design-rag-nimbus-assist.md
  - design-rag-nimbus-assist.md
  - design-rag-nimbus-assist.md

Query: What is Tech Company Inc.'s paid time off (PTO) policy?
Answer: The model cannot find sufficient information to answer the question about Tech Company Inc.'s paid time off (PTO) policy. The provided search results do not contain any details regarding the company's PTO policy.
Sources:

These results demonstrate the capabilities of well-configured RAG systems. The first query produces a comprehensive, well-structured answer that synthesizes information from multiple source documents, providing detailed context about Nimbus Assist's purpose, capabilities, and integration points. The second query showcases the system's ability to handle technical, detailed questions by extracting specific configuration parameters and organizing them clearly with proper formatting. Most importantly, the third query demonstrates responsible AI behavior: when the knowledge base lacks relevant information (in this specific case, the PTO policy), the system explicitly states this limitation rather than fabricating an answer, maintaining trust and reliability by acknowledging the boundaries of its knowledge.

Congratulations on completing the final lesson of Managing Data for GenAI with Bedrock Knowledge Bases! You've successfully mastered the complete spectrum of GenAI data management, from initial vector storage setup through RAG querying implementations. Throughout this course, you've built production-ready infrastructure for intelligent document retrieval, learned to balance direct mathematical search with AI-powered response generation, and developed the skills to create trustworthy, citation-backed question-answering systems. Your achievement in reaching this point represents a significant milestone in your GenAI expertise.

Your journey through this foundational course has equipped you with essential knowledge of AWS's GenAI ecosystem and practical experience implementing the core components that power modern AI applications, which you will apply in the upcoming practice section. Then, as you move forward to the next course, "Putting Bedrock Models to Action with Strands Agents", you'll build upon this solid foundation to create even more sophisticated conversational agents that can take real-world actions. You'll explore advanced capabilities, including tool integration, knowledge base connections, and cutting-edge Model Context Protocol features that will take your GenAI expertise to the next level of practical application and real-world impact. Keep learning!