Building RAG Pipelines

import os
import re
from dotenv import load_dotenv
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi
import chromadb

load_dotenv()

PROXY_URL = "https://litellm.6640.ucf.spencerlyon.com"


def get_model(model_name: str) -> OpenAIChatModel:
    """Create a model connection through our LiteLLM proxy."""
    return OpenAIChatModel(
        model_name,
        provider=OpenAIProvider(
            base_url=PROXY_URL,
            api_key=os.environ["CAP6640_API_KEY"],
        ),
    )


# Load models we'll use throughout
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

document = """
Natural language processing (NLP) is a subfield of linguistics, computer science, and artificial intelligence concerned with the interactions between computers and human language. The goal is a computer capable of understanding the contents of documents, including the contextual nuances of the language within them. NLP combines computational linguistics with statistical, machine learning, and deep learning models to process human language.

Tokenization is the process of breaking text into smaller units called tokens. These tokens can be words, subwords, or characters depending on the chosen strategy. Modern NLP systems typically use subword tokenization methods like Byte Pair Encoding (BPE) or WordPiece, which balance vocabulary size with the ability to handle rare and out-of-vocabulary words. The choice of tokenization strategy affects all downstream tasks.

Text can be represented as sparse vectors using methods like bag-of-words or TF-IDF, which count word occurrences. Alternatively, dense representations like word embeddings (Word2Vec, GloVe) capture semantic meaning in continuous vector spaces. The key advantage of dense representations is that semantically similar words have similar vectors, enabling operations like analogy solving and similarity search.

The attention mechanism, introduced in the landmark 2017 paper "Attention Is All You Need," allows each position in a sequence to attend to all other positions. This enables transformer models to capture long-range dependencies without the sequential processing bottleneck of recurrent neural networks. Multi-head attention runs multiple attention operations in parallel, letting the model focus on different aspects of the input simultaneously.

Retrieval-augmented generation (RAG) combines information retrieval with language model generation. In a RAG system, relevant documents are first retrieved from an external knowledge base using embedding similarity search. These retrieved passages are then provided as context to a large language model, which generates a response grounded in the specific documents rather than relying solely on its parametric knowledge. This approach reduces hallucination and enables citation of sources.

Large language models like GPT-4, Claude, and Gemini are trained on massive text corpora and can perform a wide variety of NLP tasks through prompting alone. However, they have important limitations: their knowledge is frozen at training time, they can hallucinate plausible-sounding but incorrect information, and they cannot access private or domain-specific data unless it is provided in the prompt or through techniques like RAG.
"""

print(f"Document length: {len(document.split())} words, {len(document)} characters")

def chunk_fixed(text, chunk_size=50, overlap=10):
    """Split text into fixed-size word chunks with overlap."""
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i : i + chunk_size])
        if chunk.strip():
            chunks.append(chunk.strip())
    return chunks


fixed_chunks = chunk_fixed(document, chunk_size=50, overlap=10)
print(f"Fixed-size chunking: {len(fixed_chunks)} chunks\n")
for i, chunk in enumerate(fixed_chunks[:3]):
    print(f"Chunk {i+1} ({len(chunk.split())} words): {chunk[:80]}...")
    print()

def chunk_sentences(text, max_words=60):
    """Split into sentence groups, respecting sentence boundaries."""
    sentences = re.split(r"(?<=[.!?])\s+", text.strip())
    chunks = []
    current_chunk = []
    current_size = 0

    for sent in sentences:
        sent_words = len(sent.split())
        if current_size + sent_words > max_words and current_chunk:
            chunks.append(" ".join(current_chunk))
            current_chunk = [sent]
            current_size = sent_words
        else:
            current_chunk.append(sent)
            current_size += sent_words

    if current_chunk:
        chunks.append(" ".join(current_chunk))
    return chunks


sentence_chunks = chunk_sentences(document, max_words=60)
print(f"Sentence-aware chunking: {len(sentence_chunks)} chunks\n")
for i, chunk in enumerate(sentence_chunks[:3]):
    print(f"Chunk {i+1} ({len(chunk.split())} words): {chunk[:80]}...")
    print()

def chunk_paragraphs(text, max_words=80):
    """Split on paragraph boundaries, with sentence fallback for long paragraphs."""
    paragraphs = [p.strip() for p in text.strip().split("\n\n") if p.strip()]
    chunks = []
    for para in paragraphs:
        if len(para.split()) <= max_words:
            chunks.append(para)
        else:
            # Fall back to sentence grouping for long paragraphs
            chunks.extend(chunk_sentences(para, max_words=max_words))
    return chunks


para_chunks = chunk_paragraphs(document, max_words=80)
print(f"Paragraph-based chunking: {len(para_chunks)} chunks\n")
for i, chunk in enumerate(para_chunks[:3]):
    print(f"Chunk {i+1} ({len(chunk.split())} words): {chunk[:80]}...")
    print()

query = "How does RAG reduce hallucination?"

def retrieve_top(chunks, query, embed_model, n=2):
    """Quick dense retrieval over a list of chunks."""
    q_emb = embed_model.encode([query])
    c_embs = embed_model.encode(chunks)
    from sentence_transformers.util import cos_sim
    sims = cos_sim(q_emb, c_embs)[0]
    top_indices = sims.argsort(descending=True)[:n]
    return [(chunks[i], sims[i].item()) for i in top_indices]

print("=" * 70)
print(f"Query: '{query}'\n")

for name, chunks in [("Fixed-size", fixed_chunks), ("Sentence-aware", sentence_chunks), ("Paragraph-based", para_chunks)]:
    results = retrieve_top(chunks, query, embed_model, n=1)
    chunk_text, score = results[0]
    print(f"--- {name} (top result, similarity: {score:.3f}) ---")
    print(f"{chunk_text[:120]}...")
    print()

# We'll use paragraph chunks for the rest of this lecture
chunks = para_chunks
print(f"Working with {len(chunks)} paragraph-based chunks\n")

# BM25 needs tokenized documents (just lowercase word splits for now)
tokenized_chunks = [chunk.lower().split() for chunk in chunks]
bm25 = BM25Okapi(tokenized_chunks)

# Try a keyword-heavy query
query_keyword = "BPE tokenization subword"
bm25_scores = bm25.get_scores(query_keyword.lower().split())

print(f"Query: '{query_keyword}'\n")
print("BM25 ranking:")
for rank, idx in enumerate(bm25_scores.argsort()[::-1][:3]):
    print(f"  [{rank+1}] (score: {bm25_scores[idx]:.2f}) {chunks[idx][:70]}...")

# Set up ChromaDB with our paragraph chunks
client = chromadb.Client()
collection = client.create_collection("advanced_rag", metadata={"hnsw:space": "cosine"})

chunk_embeddings = embed_model.encode(chunks).tolist()
collection.add(
    ids=[f"chunk_{i}" for i in range(len(chunks))],
    documents=chunks,
    embeddings=chunk_embeddings,
)

# Same query, dense retrieval
dense_results = collection.query(
    query_embeddings=embed_model.encode([query_keyword]).tolist(),
    n_results=3,
)

print(f"Query: '{query_keyword}'\n")
print("Dense ranking:")
for rank, (doc, dist) in enumerate(zip(dense_results["documents"][0], dense_results["distances"][0])):
    print(f"  [{rank+1}] (similarity: {1-dist:.3f}) {doc[:70]}...")

def reciprocal_rank_fusion(ranked_lists, k=60):
    """Merge multiple ranked lists using RRF. Each list is [(doc_id, ...), ...]."""
    rrf_scores = {}
    for ranked_list in ranked_lists:
        for rank, doc_id in enumerate(ranked_list):
            if doc_id not in rrf_scores:
                rrf_scores[doc_id] = 0
            rrf_scores[doc_id] += 1 / (k + rank + 1)  # rank is 0-indexed
    return sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)


def hybrid_search(query, chunks, bm25, collection, embed_model, n=5):
    """Run BM25 + dense search and fuse with RRF."""
    # BM25 ranking
    bm25_scores = bm25.get_scores(query.lower().split())
    bm25_ranked = bm25_scores.argsort()[::-1][:n].tolist()

    # Dense ranking
    dense_results = collection.query(
        query_embeddings=embed_model.encode([query]).tolist(),
        n_results=n,
    )
    dense_ranked = [int(id.split("_")[1]) for id in dense_results["ids"][0]]

    # Fuse
    fused = reciprocal_rank_fusion([bm25_ranked, dense_ranked])
    return fused


# Test hybrid search
query_test = "How does RAG use retrieved documents to reduce hallucination?"
hybrid_results = hybrid_search(query_test, chunks, bm25, collection, embed_model)

print(f"Query: '{query_test}'\n")
print("Hybrid (RRF) ranking:")
for rank, (idx, score) in enumerate(hybrid_results[:3]):
    print(f"  [{rank+1}] (RRF score: {score:.4f}) {chunks[idx][:70]}...")

# Load a cross-encoder reranking model
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

query_rerank = "What are the key limitations of using LLMs without retrieval?"

# Step 1: Retrieve wide with hybrid search (all chunks in our small example)
hybrid_results = hybrid_search(query_rerank, chunks, bm25, collection, embed_model, n=len(chunks))
candidate_indices = [idx for idx, _ in hybrid_results]

# Step 2: Rerank with cross-encoder
pairs = [[query_rerank, chunks[idx]] for idx in candidate_indices]
rerank_scores = reranker.predict(pairs)

# Sort by cross-encoder score
reranked = sorted(
    zip(candidate_indices, rerank_scores),
    key=lambda x: x[1],
    reverse=True,
)

print(f"Query: '{query_rerank}'\n")
print("Before reranking (hybrid order):")
for rank, (idx, rrf_score) in enumerate(hybrid_results[:3]):
    print(f"  [{rank+1}] {chunks[idx][:80]}...")

print("\nAfter cross-encoder reranking:")
for rank, (idx, score) in enumerate(reranked[:3]):
    print(f"  [{rank+1}] (score: {score:.3f}) {chunks[idx][:80]}...")

def advanced_rag(query, chunks, bm25, collection, embed_model, reranker, top_k=3):
    """Full Advanced RAG pipeline: hybrid search + reranking + generation."""
    # 1. Hybrid retrieve (wide)
    hybrid_results = hybrid_search(query, chunks, bm25, collection, embed_model, n=min(len(chunks), 20))
    candidate_indices = [idx for idx, _ in hybrid_results]

    # 2. Rerank (narrow)
    pairs = [[query, chunks[idx]] for idx in candidate_indices]
    scores = reranker.predict(pairs)
    reranked = sorted(zip(candidate_indices, scores), key=lambda x: x[1], reverse=True)
    top_chunks = [(chunks[idx], score) for idx, score in reranked[:top_k]]

    # 3. Build context
    context = "\n".join(f"[{i+1}] {chunk}" for i, (chunk, _) in enumerate(top_chunks))
    return context, top_chunks


query_final = "Why is RAG better than just using a language model alone?"

context, top_chunks = advanced_rag(
    query_final, chunks, bm25, collection, embed_model, reranker
)

print("Retrieved context (after hybrid search + reranking):\n")
for i, (chunk, score) in enumerate(top_chunks):
    print(f"[{i+1}] (rerank score: {score:.3f})")
    print(f"    {chunk[:100]}...")
    print()

# Generate a grounded answer
rag_agent = Agent(
    get_model("claude-sonnet-4-6"),
    instructions=(
        "Answer the user's question based ONLY on the provided context. "
        "Cite your sources using [1], [2], etc. "
        "If the context doesn't contain enough information, say so."
    ),
)

prompt = f"""Context:
{context}

Question: {query_final}"""

result = await rag_agent.run(prompt)
print(f"Query: {query_final}\n")
print(f"Answer:\n{result.output}")