Retrieval-augmented generation is the workload we get asked about most, and the usual demo — a notebook, an in-memory FAISS index, and an API key — is not something you can hand to an application team. This tutorial builds the production shape of it on AWS: a SageMaker AI endpoint serving embeddings, an Amazon OpenSearch Serverless collection as the vector store, a chunk-and-ingest job, and a retrieval + generation path that returns citations. Everything stays inside your account and your VPC boundary.
You will need the SageMaker Python SDK (pip install "sagemaker>=2.200" opensearch-py), an execution role, and permission to create an OpenSearch Serverless collection.
The architecture in one paragraph
Documents land in S3. A SageMaker Processing job splits them into chunks, calls an embeddings endpoint, and writes vectors plus source metadata into an OpenSearch Serverless vector collection. At query time your application embeds the user question with the same endpoint, runs a k-NN search, and passes the top chunks to a generator — either a SageMaker AI endpoint hosting your own model or Amazon Bedrock. Two SageMaker endpoints at most, one serverless index, no vector database to patch.
If you are still deciding whether the generator should be Bedrock or a model you host, read Bedrock or SageMaker AI for Generative AI? first; the retrieval half of this tutorial is identical either way.
1. Deploy an embeddings endpoint
Embeddings are a small-model, high-QPS workload, which means CPU instances and a real autoscaling policy — not a GPU sitting at 4 percent utilization.
import sagemaker
from sagemaker.huggingface import HuggingFaceModel
role = sagemaker.get_execution_role()
model = HuggingFaceModel(
role=role,
env={
"HF_MODEL_ID": "BAAI/bge-base-en-v1.5",
"HF_TASK": "feature-extraction",
},
transformers_version="4.49",
pytorch_version="2.6",
py_version="py312",
)
embedder = model.deploy(
initial_instance_count=1,
instance_type="ml.c7i.2xlarge",
endpoint_name="rag-embeddings",
)
bge-base-en-v1.5 produces 768-dimensional vectors and is a sensible default for English text. Whatever you pick, write the dimension down — the index mapping in step 2 must match it exactly, and changing the embedding model later means re-indexing every document.
A detail that bites people: many sentence-embedding models expect L2-normalized output and a query prefix. For BGE models, prefix questions (not documents) with "Represent this sentence for searching relevant passages: ". Skipping this quietly costs a few points of recall.
For the full walkthrough of deploying, autoscaling, and tearing down a Hugging Face endpoint, see Deploy a Hugging Face Model to a SageMaker AI Real-Time Endpoint.
2. Create the OpenSearch Serverless vector collection
OpenSearch Serverless has a dedicated VECTORSEARCH collection type. It needs three policies before it will start: encryption, network, and data access.
import boto3, json
aoss = boto3.client("opensearchserverless")
NAME = "rag-kb"
aoss.create_security_policy(
name=f"{NAME}-enc", type="encryption",
policy=json.dumps({
"Rules": [{"ResourceType": "collection", "Resource": [f"collection/{NAME}"]}],
"AWSOwnedKey": True,
}),
)
aoss.create_security_policy(
name=f"{NAME}-net", type="network",
policy=json.dumps([{
"Rules": [{"ResourceType": "collection", "Resource": [f"collection/{NAME}"]}],
"AllowFromPublic": False,
"SourceVPCEs": ["vpce-0123456789abcdef0"],
}]),
)
collection = aoss.create_collection(name=NAME, type="VECTORSEARCH")
AllowFromPublic: False with a VPC endpoint keeps the collection off the internet; create the VPC endpoint first with create_vpc_endpoint if you do not already have one. Use a customer-managed KMS key instead of AWSOwnedKey if your compliance posture requires it. The same reasoning applies to the SageMaker side — see Private by Default.
Then a data access policy granting the SageMaker execution role and your application role index-level permissions:
aoss.create_access_policy(
name=f"{NAME}-access", type="data",
policy=json.dumps([{
"Rules": [
{"ResourceType": "index",
"Resource": [f"index/{NAME}/*"],
"Permission": ["aoss:CreateIndex", "aoss:WriteDocument",
"aoss:ReadDocument", "aoss:DescribeIndex"]},
{"ResourceType": "collection",
"Resource": [f"collection/{NAME}"],
"Permission": ["aoss:DescribeCollection"]},
],
"Principal": [
"arn:aws:iam::123456789012:role/SageMakerExecutionRole",
"arn:aws:iam::123456789012:role/RagAppRole",
],
}]),
)
Collections take a couple of minutes to become ACTIVE. Poll aoss.batch_get_collection(names=[NAME]) until the status flips, then grab the endpoint from the response.
3. Create the index with the right k-NN mapping
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth
region = boto3.Session().region_name
auth = AWSV4SignerAuth(boto3.Session().get_credentials(), region, "aoss")
host = collection_endpoint.replace("https://", "")
client = OpenSearch(
hosts=[{"host": host, "port": 443}],
http_auth=auth, use_ssl=True, verify_certs=True,
connection_class=RequestsHttpConnection, pool_maxsize=20,
)
client.indices.create(
index="kb-chunks",
body={
"settings": {"index": {"knn": True}},
"mappings": {
"properties": {
"embedding": {
"type": "knn_vector",
"dimension": 768,
"method": {
"name": "hnsw",
"engine": "faiss",
"space_type": "innerproduct",
"parameters": {"ef_construction": 256, "m": 16},
},
},
"text": {"type": "text"},
"source_uri": {"type": "keyword"},
"page": {"type": "integer"},
"tenant": {"type": "keyword"},
"updated_at": {"type": "date"},
}
},
},
)
Notes that matter more than they look:
space_type: innerproductis correct for normalized embeddings; usel2orcosinesimilif your model does not normalize.dimensionmust equal your model's output width. A mismatch fails at ingest, not at create time.- Store
source_uriandpagenow. Citations you cannot produce later are the number one complaint from RAG pilot users. tenant(oracl_group) as akeywordlets you filter results per customer or per permission group in the same index. Retrofitting multi-tenancy is painful; adding an unused field is free.
4. Chunk and ingest
Run ingestion as a SageMaker Processing job so it is repeatable, logged, and runs in the same VPC. The core loop:
import itertools
def chunk(text, size=800, overlap=120):
words, out, i = text.split(), [], 0
while i < len(words):
out.append(" ".join(words[i:i + size]))
i += size - overlap
return out
def embed(batch):
resp = embedder.predict({"inputs": batch})
import numpy as np
v = np.array(resp)
if v.ndim == 3: # token-level output -> mean pool
v = v.mean(axis=1)
return (v / np.linalg.norm(v, axis=1, keepdims=True)).tolist()
def batched(it, n):
it = iter(it)
while (b := list(itertools.islice(it, n))):
yield b
actions = []
for doc in documents: # {"uri":..., "page":..., "text":...}
chunks = chunk(doc["text"])
for group in batched(chunks, 16):
for text, vec in zip(group, embed(group)):
actions.append({"index": {"_index": "kb-chunks"}})
actions.append({
"text": text, "embedding": vec,
"source_uri": doc["uri"], "page": doc["page"],
"tenant": doc.get("tenant", "public"),
"updated_at": doc["updated_at"],
})
if len(actions) >= 400:
client.bulk(body=actions); actions = []
if actions:
client.bulk(body=actions)
Chunking is the highest-leverage knob in the whole system, and 800 words with 120 of overlap is only a starting point. Split on structure first — headings, sections, table rows — and fall back to fixed windows inside a section. Prepend the document title and heading path to each chunk before embedding; it costs a few tokens and measurably improves retrieval on questions that use vocabulary from the heading rather than the body.
One OpenSearch Serverless quirk: the bulk API does not accept a client-supplied _id on VECTORSEARCH collections. To make ingestion idempotent, delete by query on source_uri before re-indexing a document rather than relying on upserts.
5. Retrieve, then generate
Hybrid search — vector plus BM25 — beats pure vector search on almost every corpus that contains product codes, error strings, or names. Run both and fuse:
def search(question, k=8, tenant="public"):
qvec = embed([f"Represent this sentence for searching relevant passages: {question}"])[0]
body = {
"size": k,
"query": {
"bool": {
"filter": [{"term": {"tenant": tenant}}],
"should": [
{"knn": {"embedding": {"vector": qvec, "k": k}}},
{"match": {"text": {"query": question, "boost": 0.3}}},
],
}
},
"_source": ["text", "source_uri", "page"],
}
return client.search(index="kb-chunks", body=body)["hits"]["hits"]
Then build a prompt that forces grounding and citation, and send it to your generator:
def answer(question, tenant="public"):
hits = search(question, tenant=tenant)
context = "\n\n".join(
f"[{i+1}] ({h['_source']['source_uri']} p.{h['_source']['page']})\n{h['_source']['text']}"
for i, h in enumerate(hits)
)
prompt = (
"Answer the question using only the numbered context below. "
"Cite sources as [1], [2]. If the context does not contain the answer, "
"say you do not know.\n\n"
f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
)
br = boto3.client("bedrock-runtime")
resp = br.converse(
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
messages=[{"role": "user", "content": [{"text": prompt}]}],
inferenceConfig={"maxTokens": 800, "temperature": 0.1},
)
return resp["output"]["message"]["content"][0]["text"], hits
Swapping in a self-hosted generator is a one-function change: call invoke_endpoint on your SageMaker AI endpoint instead of converse, and keep the prompt identical so your evaluation numbers stay comparable.
6. Measure retrieval before you tune generation
Nearly every stalled RAG project we are called into is tuning prompts when the retriever is the problem. Build a small labelled set — 50 to 100 real questions with the chunk or document that should answer them — and track two numbers on every change:
- Recall@k: fraction of questions where a correct chunk appears in the top k. If this is below about 0.9 at k=8, no prompt will save you.
- Answer faithfulness: fraction of generated answers whose claims are supported by the retrieved context, scored by an LLM judge or by hand.
Run these as a CI gate the same way you would gate a model release; our post on automated LLM evaluation gates with fmeval shows the mechanics. Changes worth measuring, roughly in order of payoff: chunk boundaries, heading prefixes, hybrid weighting, a reranker (a cross-encoder on a second small endpoint, applied to the top 50), and only then the prompt.
Costs and teardown
The recurring costs are OpenSearch Serverless OCUs (billed per OCU-hour, with a minimum for indexing and search), the embeddings endpoint per instance-hour, and generation tokens. Two habits keep this sane: put the embeddings endpoint behind an autoscaling policy with MinCapacity=1, and do not leave a development collection running — an idle vector collection still bills its minimum OCUs, which is the single most common surprise line item on a RAG proof of concept.
embedder.delete_endpoint(delete_endpoint_config=True)
embedder.delete_model()
aoss.delete_collection(id=collection_id)
When to use Bedrock Knowledge Bases instead
If your documents are plain files in S3, your access model is simple, and you are happy with the managed chunking options, Amazon Bedrock Knowledge Bases will do all of the above with far less code — it can even provision the OpenSearch Serverless collection for you. Build the pipeline in this tutorial when you need a specific embedding model, custom chunking driven by document structure, per-tenant filtering, a reranking stage, or retrieval that other applications query directly. The two are not exclusive: several of our clients run a managed knowledge base for general document search and this pipeline for the corpus that pays the bills.
Want a second pair of eyes on a RAG design, or a retrieval evaluation before you commit to an architecture? Talk to us.