Lab — RAG Builder

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 sentence_transformers.util import cos_sim
from rank_bm25 import BM25Okapi
import chromadb
import wikipedia

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"],
        ),
    )


embed_model = SentenceTransformer("all-MiniLM-L6-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

topics = [
    "Natural language processing",
    "Tokenization (lexical analysis)",
    "Word embedding",
    "Tf–idf",
    "Named-entity recognition",
    "Recurrent neural network",
    "Transformer (deep learning architecture)",
    "BERT (language model)",
    "GPT-4",
    "Retrieval-augmented generation",
]

articles = {}
for topic in topics:
    try:
        page = wikipedia.page(topic, auto_suggest=False)
        articles[topic] = page.content
        print(f"  {topic}: {len(page.content.split())} words")
    except Exception as e:
        print(f"  {topic}: FAILED ({e})")

print(f"\nLoaded {len(articles)} articles")

def chunk_paragraphs(text, max_words=100):
    """Split on paragraph boundaries, grouping short paragraphs."""
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip() and len(p.split()) > 10]
    chunks = []
    current = []
    current_size = 0
    for para in paragraphs:
        words = len(para.split())
        if current_size + words > max_words and current:
            chunks.append(" ".join(current))
            current = [para]
            current_size = words
        else:
            current.append(para)
            current_size += words
    if current:
        chunks.append(" ".join(current))
    return chunks


# Chunk all articles, tracking source
all_chunks = []
chunk_sources = []

for topic, content in articles.items():
    topic_chunks = chunk_paragraphs(content, max_words=100)
    for chunk in topic_chunks:
        all_chunks.append(chunk)
        chunk_sources.append(topic)

print(f"Total chunks: {len(all_chunks)}")
print(f"Average chunk size: {sum(len(c.split()) for c in all_chunks) / len(all_chunks):.0f} words")
print(f"\nChunks per article:")
for topic in articles:
    count = chunk_sources.count(topic)
    print(f"  {topic}: {count} chunks")

# Dense index (ChromaDB)
client = chromadb.Client()

# Delete collection if it exists from a previous run
try:
    client.delete_collection("wiki_rag")
except Exception:
    pass

collection = client.create_collection("wiki_rag", metadata={"hnsw:space": "cosine"})

# Embed and store all chunks
print("Embedding chunks... ", end="")
chunk_embeddings = embed_model.encode(all_chunks, show_progress_bar=True).tolist()

collection.add(
    ids=[f"chunk_{i}" for i in range(len(all_chunks))],
    documents=all_chunks,
    embeddings=chunk_embeddings,
    metadatas=[{"source": src} for src in chunk_sources],
)
print(f"Indexed {collection.count()} chunks in ChromaDB")

# Sparse index (BM25)
tokenized_chunks = [chunk.lower().split() for chunk in all_chunks]
bm25 = BM25Okapi(tokenized_chunks)
print(f"Built BM25 index over {len(tokenized_chunks)} chunks")

def reciprocal_rank_fusion(ranked_lists, k=60):
    """Merge multiple ranked lists using RRF."""
    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)
    return sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)


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

    # Dense
    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
    return reciprocal_rank_fusion([bm25_ranked, dense_ranked])


def rerank(query, candidate_indices, top_k=5):
    """Rerank candidates with cross-encoder."""
    pairs = [[query, all_chunks[idx]] for idx in candidate_indices]
    scores = reranker.predict(pairs)
    reranked = sorted(zip(candidate_indices, scores), key=lambda x: x[1], reverse=True)
    return reranked[:top_k]


def rag_pipeline(query, top_k=5):
    """Full pipeline: hybrid search → rerank → format context."""
    # Retrieve wide
    hybrid_results = hybrid_search(query, n=30)
    candidate_indices = [idx for idx, _ in hybrid_results]

    # Rerank narrow
    top_chunks = rerank(query, candidate_indices, top_k=top_k)

    # Format context with source attribution
    context_parts = []
    sources = []
    for i, (idx, score) in enumerate(top_chunks):
        source = chunk_sources[idx]
        context_parts.append(f"[{i+1}] (Source: {source})\n{all_chunks[idx]}")
        sources.append(source)

    context = "\n\n".join(context_parts)
    return context, top_chunks, sources

query = "How do transformer models handle long-range dependencies in text?"
context, top_chunks, sources = rag_pipeline(query)

print(f"Query: '{query}'\n")
print(f"Retrieved {len(top_chunks)} chunks from: {', '.join(dict.fromkeys(sources))}\n")
for i, (idx, score) in enumerate(top_chunks):
    print(f"[{i+1}] (rerank: {score:.2f}, source: {chunk_sources[idx]})")
    print(f"    {all_chunks[idx][: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."
    ),
)

result = await rag_agent.run(f"Context:\n{context}\n\nQuestion: {query}")
print(result.output)

test_queries = [
    "What is the difference between BPE and WordPiece tokenization?",
    "How does BERT differ from GPT in architecture?",
    "What are the limitations of TF-IDF for text retrieval?",
]

for q in test_queries:
    ctx, chunks, srcs = rag_pipeline(q, top_k=3)
    print(f"Q: {q}")
    print(f"   Sources: {', '.join(dict.fromkeys(srcs))}")
    print(f"   Top chunk: {all_chunks[chunks[0][0]][:80]}...")
    print()

async def hyde_search(query, n=20):
    """Generate a hypothetical answer, embed it, and search."""
    # Step 1: Generate hypothetical answer
    hyde_agent = Agent(
        get_model("claude-sonnet-4-6"),
        instructions=(
            "Write a short, factual paragraph (3-4 sentences) that would answer "
            "the following question. Write as if you are a textbook. "
            "Do not say 'I don't know' — just write your best answer."
        ),
    )
    hypothetical = await hyde_agent.run(query)
    hypo_text = hypothetical.output

    # Step 2: Embed the hypothetical answer (not the original query)
    hypo_embedding = embed_model.encode([hypo_text]).tolist()

    # Step 3: Search with the hypothetical embedding
    results = collection.query(
        query_embeddings=hypo_embedding,
        n_results=n,
    )

    return results, hypo_text

# Compare standard vs. HyDE retrieval on a short query
short_query = "attention"

# Standard dense search
standard_results = collection.query(
    query_embeddings=embed_model.encode([short_query]).tolist(),
    n_results=3,
)

print(f"Query: '{short_query}'\n")
print("--- Standard Dense Search ---")
for i, (doc, dist) in enumerate(zip(standard_results["documents"][0], standard_results["distances"][0])):
    print(f"  [{i+1}] (sim: {1-dist:.3f}) {doc[:80]}...")

# HyDE search
hyde_results, hypo_text = await hyde_search(short_query)

print(f"\n--- HyDE Search ---")
print(f"Hypothetical answer: {hypo_text[:120]}...\n")
for i, (doc, dist) in enumerate(zip(hyde_results["documents"][0][:3], hyde_results["distances"][0][:3])):
    print(f"  [{i+1}] (sim: {1-dist:.3f}) {doc[:80]}...")

# Test suite: queries with expected content
test_suite = [
    {
        "query": "What is tokenization and why is it important for NLP?",
        "expected_keywords": ["token", "subword", "BPE"],
    },
    {
        "query": "How do word embeddings capture semantic meaning?",
        "expected_keywords": ["vector", "Word2Vec", "semantic"],
    },
    {
        "query": "What problem does the attention mechanism solve?",
        "expected_keywords": ["attention", "long-range", "parallel"],
    },
    {
        "query": "How does RAG reduce hallucination in language models?",
        "expected_keywords": ["retriev", "hallucin", "knowledge"],
    },
    {
        "query": "What are the main differences between BERT and GPT?",
        "expected_keywords": ["BERT", "GPT", "encoder"],
    },
]

print(f"Test suite: {len(test_suite)} queries")
for i, test in enumerate(test_suite):
    print(f"  {i+1}. {test['query']}")

print("Running pipeline on test suite...\n")

results = []
for test in test_suite:
    query = test["query"]
    context, top_chunks, sources = rag_pipeline(query, top_k=3)

    result = await rag_agent.run(f"Context:\n{context}\n\nQuestion: {query}")
    answer = result.output

    # Check how many expected keywords appear in the context
    context_lower = context.lower()
    keywords_found = sum(
        1 for kw in test["expected_keywords"]
        if kw.lower() in context_lower
    )

    results.append({
        "query": query,
        "answer": answer,
        "sources": sources,
        "keywords_found": keywords_found,
        "keywords_total": len(test["expected_keywords"]),
    })

    print(f"Q: {query}")
    print(f"   Sources: {', '.join(dict.fromkeys(sources))}")
    print(f"   Keywords in context: {keywords_found}/{len(test['expected_keywords'])}")
    print(f"   Answer: {answer[:120]}...")
    print()

# Summary
total_found = sum(r["keywords_found"] for r in results)
total_expected = sum(r["keywords_total"] for r in results)
print(f"Context coverage: {total_found}/{total_expected} expected keywords found ({total_found/total_expected*100:.0f}%)")
print()
print("Per-query breakdown:")
for r in results:
    score = r["keywords_found"] / r["keywords_total"]
    status = "GOOD" if score >= 0.67 else "WEAK" if score >= 0.33 else "POOR"
    print(f"  [{status}] {r['query'][:50]}... ({r['keywords_found']}/{r['keywords_total']})")