Lab 4 — Vanilla RAG (from scratch)¶

⏱️ ~30 min

Chunk → embed → store in Chroma → retrieve → answer. No framework.

Requires an OpenAI API key — these labs use OpenAI embeddings / function-calling specifically. (Labs 1–3 work with OpenAI or Anthropic.)

In [ ]:

!pip install -q openai chromadb
import os, getpass
if not os.environ.get("OPENAI_API_KEY"):
    os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API key: ")
from openai import OpenAI
client = OpenAI()

!pip install -q openai chromadb import os, getpass if not os.environ.get("OPENAI_API_KEY"): os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API key: ") from openai import OpenAI client = OpenAI()

1. Our "documents" (normally you'd load a PDF; we inline text to stay simple)¶

In [ ]:

documents = [
    "The Helios API rate limit is 100 requests per minute per API key.",
    "Helios refunds are processed within 5 business days to the original payment method.",
    "The Helios SDK supports Python 3.9 and above; Python 3.8 reached end of life.",
    "To rotate a Helios API key, go to Settings > Security > Rotate Key. Old keys expire in 24h.",
    "Helios stores data in the EU (Frankfurt) region by default for all new accounts.",
]

documents = [ "The Helios API rate limit is 100 requests per minute per API key.", "Helios refunds are processed within 5 business days to the original payment method.", "The Helios SDK supports Python 3.9 and above; Python 3.8 reached end of life.", "To rotate a Helios API key, go to Settings > Security > Rotate Key. Old keys expire in 24h.", "Helios stores data in the EU (Frankfurt) region by default for all new accounts.", ]

2. Chunk (already small) and embed¶

In [ ]:

def embed(texts):
    resp = client.embeddings.create(model="text-embedding-3-small", input=texts)
    return [d.embedding for d in resp.data]

doc_embeddings = embed(documents)
print(f"Embedded {len(doc_embeddings)} chunks, dim={len(doc_embeddings[0])}")

def embed(texts): resp = client.embeddings.create(model="text-embedding-3-small", input=texts) return [d.embedding for d in resp.data] doc_embeddings = embed(documents) print(f"Embedded {len(doc_embeddings)} chunks, dim={len(doc_embeddings[0])}")

3. Store in Chroma¶

In [ ]:

import chromadb
chroma = chromadb.Client()
collection = chroma.create_collection("helios-docs")
collection.add(
    ids=[f"doc-{i}" for i in range(len(documents))],
    documents=documents,
    embeddings=doc_embeddings,
)
print("Stored in Chroma.")

import chromadb chroma = chromadb.Client() collection = chroma.create_collection("helios-docs") collection.add( ids=[f"doc-{i}" for i in range(len(documents))], documents=documents, embeddings=doc_embeddings, ) print("Stored in Chroma.")

4. Retrieve the most relevant chunks for a question¶

In [ ]:

def retrieve(question, k=2):
    q_emb = embed([question])[0]
    res = collection.query(query_embeddings=[q_emb], n_results=k)
    return res["documents"][0]

question = "How long until my old API key stops working after I rotate it?"
chunks = retrieve(question)
print("Retrieved chunks:")
for c in chunks:
    print(" -", c)

def retrieve(question, k=2): q_emb = embed([question])[0] res = collection.query(query_embeddings=[q_emb], n_results=k) return res["documents"][0] question = "How long until my old API key stops working after I rotate it?" chunks = retrieve(question) print("Retrieved chunks:") for c in chunks: print(" -", c)

5. Augment the prompt and generate a grounded answer¶

In [ ]:

def answer(question):
    chunks = retrieve(question)
    context = "\n".join(f"- {c}" for c in chunks)
    messages = [
        {"role": "system", "content":
         "Answer ONLY from the context. If it's not in the context, say you don't know."},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
    ]
    resp = client.chat.completions.create(
        model="gpt-4o-mini", messages=messages, temperature=0)
    return resp.choices[0].message.content

print(answer(question))

def answer(question): chunks = retrieve(question) context = "\n".join(f"- {c}" for c in chunks) messages = [ {"role": "system", "content": "Answer ONLY from the context. If it's not in the context, say you don't know."}, {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}, ] resp = client.chat.completions.create( model="gpt-4o-mini", messages=messages, temperature=0) return resp.choices[0].message.content print(answer(question))

6. Watch grounding work: ask something NOT in the docs¶

In [ ]:

print(answer("What is the Helios phone support number?"))

print(answer("What is the Helios phone support number?"))

Takeaway: the model answered the first question from your data and correctly refused the second. That's RAG: retrieval + grounding beats relying on model memory.

✅ Did it work?¶

• The grounded answer cited the Helios docs — it said old keys expire in 24h, drawn from the retrieved chunk.
• The off-topic phone-support question was declined ("I don't know"), because that fact isn't in the documents.