Lab: Build a RAG Pipeline

pip install sentence-transformers numpy

import numpy as np
from sentence_transformers import SentenceTransformer

# Source documents -- each is a short passage on a different topic
documents = [
    {
        "title": "Photosynthesis",
        "text": (
            "Photosynthesis is the process by which green plants and certain "
            "other organisms convert light energy into chemical energy. During "
            "photosynthesis, plants capture sunlight using chlorophyll in their "
            "leaves and use it to transform water and carbon dioxide into oxygen "
            "and glucose. The glucose provides energy for the plant's metabolic "
            "processes. Photosynthesis is responsible for producing the oxygen "
            "in Earth's atmosphere."
        ),
    },
    {
        "title": "TCP/IP",
        "text": (
            "TCP/IP is the foundational protocol suite for the internet. TCP "
            "(Transmission Control Protocol) provides reliable, ordered delivery "
            "of data between applications. IP (Internet Protocol) handles "
            "addressing and routing packets across networks. Together they enable "
            "applications to communicate across diverse networks without needing "
            "to know the underlying hardware. TCP uses a three-way handshake "
            "(SYN, SYN-ACK, ACK) to establish connections and sequence numbers "
            "to ensure packets arrive in order."
        ),
    },
    {
        "title": "The French Revolution",
        "text": (
            "The French Revolution began in 1789 with the storming of the "
            "Bastille and ended in 1799 with Napoleon's rise to power. It was "
            "driven by widespread discontent with the monarchy, economic hardship, "
            "and Enlightenment ideals of liberty and equality. The revolution "
            "abolished feudalism, established the Declaration of the Rights of Man "
            "and Citizen, and fundamentally transformed French governance. The Reign "
            "of Terror under Robespierre saw thousands executed by guillotine."
        ),
    },
    {
        "title": "Hash Functions",
        "text": (
            "A cryptographic hash function takes an arbitrary-length input and "
            "produces a fixed-size output called a digest. Good hash functions "
            "have three properties: preimage resistance (cannot reverse the hash), "
            "second preimage resistance (cannot find a different input with the "
            "same hash), and collision resistance (cannot find any two inputs with "
            "the same hash). SHA-256 produces a 256-bit digest and is widely used "
            "in TLS certificates, git, and blockchain systems."
        ),
    },
    {
        "title": "Mitochondria",
        "text": (
            "Mitochondria are membrane-bound organelles found in the cytoplasm "
            "of eukaryotic cells. They generate most of the cell's supply of "
            "adenosine triphosphate (ATP), the molecule used as energy currency. "
            "Mitochondria have their own DNA, separate from nuclear DNA, which "
            "supports the endosymbiotic theory that they originated as free-living "
            "bacteria that were engulfed by ancestral eukaryotic cells. Each cell "
            "can contain hundreds to thousands of mitochondria depending on its "
            "energy demands."
        ),
    },
]

print(f"Loaded {len(documents)} documents")

def chunk_documents(docs, max_chunk_length=200):
    """Split documents into chunks by sentence boundaries."""
    chunks = []
    for doc in docs:
        sentences = doc["text"].replace(". ", ".\n").split("\n")
        current_chunk = ""
        for sentence in sentences:
            if len(current_chunk) + len(sentence) > max_chunk_length and current_chunk:
                chunks.append({
                    "title": doc["title"],
                    "text": current_chunk.strip(),
                })
                current_chunk = sentence
            else:
                current_chunk += (" " if current_chunk else "") + sentence
        if current_chunk.strip():
            chunks.append({
                "title": doc["title"],
                "text": current_chunk.strip(),
            })
    return chunks


chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks\n")
for i, chunk in enumerate(chunks):
    print(f"  Chunk {i}: [{chunk['title']}] {chunk['text'][:60]}...")

print("\nLoading embedding model...")
model = SentenceTransformer("all-MiniLM-L6-v2")

# Embed all chunks
chunk_texts = [c["text"] for c in chunks]
embeddings = model.encode(chunk_texts, show_progress_bar=False)

print(f"Embedding shape: {embeddings.shape}")
print(f"Each chunk is now a {embeddings.shape[1]}-dimensional vector")

def cosine_similarity(a, b):
    """Compute cosine similarity between two vectors."""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def retrieve(query, model, embeddings, chunks, top_k=3):
    """Embed a query and return the top-k most similar chunks."""
    query_embedding = model.encode([query], show_progress_bar=False)[0]

    # Compute similarity against all chunks
    similarities = []
    for i, emb in enumerate(embeddings):
        sim = cosine_similarity(query_embedding, emb)
        similarities.append((sim, i))

    # Sort by similarity (descending) and return top-k
    similarities.sort(reverse=True)
    results = []
    for sim, idx in similarities[:top_k]:
        results.append({
            "score": float(sim),
            "title": chunks[idx]["title"],
            "text": chunks[idx]["text"],
        })
    return results

queries = [
    "How do plants produce energy from sunlight?",
    "How does TCP establish a connection?",
    "What is SHA-256 used for?",
]

for query in queries:
    print(f"\nQuery: {query}")
    results = retrieve(query, model, embeddings, chunks, top_k=3)
    for rank, r in enumerate(results):
        print(f"  {rank + 1}. [{r['title']}] score={r['score']:.4f}")
        print(f"     {r['text'][:80]}...")

def build_rag_prompt(query, results):
    """Format a RAG prompt with retrieved context."""
    context_parts = []
    for i, r in enumerate(results):
        context_parts.append(f"[Source {i + 1}: {r['title']}]\n{r['text']}")
    context = "\n\n".join(context_parts)

    prompt = (
        "Answer the question based only on the provided context. "
        "If the context does not contain enough information, say so.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {query}\n\n"
        "Answer:"
    )
    return prompt


# Build and display a RAG prompt
query = "How do plants convert sunlight into energy?"
results = retrieve(query, model, embeddings, chunks, top_k=2)
prompt = build_rag_prompt(query, results)

print("\nFull RAG Prompt:")
print("=" * 60)
print(prompt)
print("=" * 60)
print(f"\nPrompt length: {len(prompt)} characters")

# Uncomment the following to call the OpenAI API:
#
# import os
# from openai import OpenAI
#
# client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
#
# response = client.chat.completions.create(
#     model="gpt-4o-mini",
#     messages=[{"role": "user", "content": prompt}],
#     max_tokens=200,
# )
#
# print("\nLLM Response:")
# print(response.choices[0].message.content)