RAG Pipeline Deep Dive — Retrieval-Augmented Generation

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", ".", " ", ""],
    length_function=len,
)

with open("documentation.md", "r") as f:
    document = f.read()

chunks = text_splitter.split_text(document)
print(f"Split {len(document)} characters into {len(chunks)} chunks")

for i, chunk in enumerate(chunks[:3]):
    print(f"\nChunk {i+1} ({len(chunk)} chars):")
    print(chunk[:150] + "...")

from openai import OpenAI
import numpy as np

client = OpenAI()

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

chunk_texts = [chunk for chunk in chunks[:10]]
embeddings = get_embeddings(chunk_texts)

print(f"Generated {len(embeddings)} embeddings")
print(f"Each embedding dimension: {len(embeddings[0])}")

# Check similarity between first two chunks
similarity = np.dot(embeddings[0], embeddings[1])
print(f"Similarity between chunk 1 and 2: {similarity:.4f}")

import chromadb
from chromadb.config import Settings

# Initialize Chroma client
chroma_client = chromadb.Client(Settings(
    chroma_db_impl="duckdb+parquet",
    persist_directory="./chroma_db"
))

# Create or get collection
collection = chroma_client.get_or_create_collection(
    name="documentation",
    metadata={"hnsw:space": "cosine"}
)

# Add chunks with embeddings
collection.add(
    ids=[f"chunk_{i}" for i in range(len(chunk_texts))],
    documents=chunk_texts,
    embeddings=embeddings,
    metadatas=[{"source": "documentation.md", "chunk_index": i}
               for i in range(len(chunk_texts))]
)

# Semantic search
query = "What are the rate limits for API calls?"
query_embedding = get_embeddings([query])[0]

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    include=["documents", "distances", "metadatas"]
)

for i, (doc, dist, meta) in enumerate(zip(
    results["documents"][0],
    results["distances"][0],
    results["metadatas"][0]
)):
    print(f"\nResult {i+1} (distance: {dist:.4f})")
    print(f"Source: {meta['source']}, Chunk: {meta['chunk_index']}")
    print(doc[:200] + "...")

def rag_answer(query, collection, llm_client, n_results=3):
    # 1. Embed the query
    query_embedding = get_embeddings([query])[0]

    # 2. Retrieve relevant chunks
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results,
        include=["documents", "metadatas"]
    )

    # 3. Build context from retrieved chunks
    context_parts = []
    for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
        context_parts.append(
            f"[Source: {meta['source']}, Section: {meta['chunk_index']}]\n{doc}"
        )
    context = "\n\n---\n\n".join(context_parts)

    # 4. Generate answer with context
    response = llm_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": ""]
You are a helpful documentation assistant. Answer based ONLY on the
provided context. If the context does not contain the answer, say
"I cannot find this information in the documentation."

Always cite the source of your information.
"""},
            {"role": "user", "content": f"""
Context:
{context}

Question: {query}

Answer with citations:
"""}
        ]
    )

    return response.choices[0].message.content

# Test the RAG pipeline
answer = rag_answer(
    "What happens if I exceed the API rate limit?",
    collection,
    client
)
print(answer)

def hybrid_search(query, collection, alpha=0.5):
    """Combine semantic and keyword search scores."""
    from sklearn.feature_extraction.text import TfidfVectorizer

    semantic_embedding = get_embeddings([query])[0]
    semantic_results = collection.query(
        query_embeddings=[semantic_embedding],
        n_results=10,
        include=["documents", "distances"]
    )

    # TF-IDF keyword search
    all_docs = collection.get(include=["documents"])["documents"]
    vectorizer = TfidfVectorizer().fit_transform(all_docs + [query])
    query_vec = vectorizer[-1]
    doc_vecs = vectorizer[:-1]
    keyword_scores = (doc_vecs @ query_vec.T).toarray().flatten()

    # Combine scores
    doc_ids = collection.get()["ids"]
    combined = {}
    for i, (sem_score, kid) in enumerate(zip(
        [1 - d for d in semantic_results["distances"][0]],
        [doc_ids.index(id_) for id_ in semantic_results["ids"][0]]
    )):
        combined[i] = alpha * sem_score + (1 - alpha) * keyword_scores[kid]

    return sorted(combined.items(), key=lambda x: x[1], reverse=True)[:3]