Chapter 3: Custom RAG Implementations

When to Build Custom

✅ Build custom when:

• Need fine-grained control
• Specific performance requirements
• Unique retrieval logic
• Minimal dependencies

❌ Use frameworks when:

• Rapid prototyping
• Standard use cases
• Team familiarity matters

Minimal RAG Implementation

import openai
from chromadb import Client
from chromadb.config import Settings

class SimpleRAG:
    def __init__(self, collection_name="docs"):
        self.client = Client(Settings(anonymized_telemetry=False))
        self.collection = self.client.create_collection(collection_name)
        self.openai = openai.OpenAI()
    
    def add_documents(self, texts, ids=None):
        """Add documents to the knowledge base"""
        if ids is None:
            ids = [f"doc_{i}" for i in range(len(texts))]
        
        embeddings = self._embed(texts)
        self.collection.add(
            embeddings=embeddings,
            documents=texts,
            ids=ids
        )
    
    def _embed(self, texts):
        """Generate embeddings"""
        response = self.openai.embeddings.create(
            input=texts,
            model="text-embedding-3-small"
        )
        return [e.embedding for e in response.data]
    
    def query(self, question, top_k=3):
        """Query the RAG system"""
        # Retrieve
        query_embedding = self._embed([question])[0]
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k
        )
        
        # Generate
        context = "\n".join(results["documents"][0])
        prompt = f"""Context: {context}

Question: {question}

Answer based on the context:"""
        
        response = self.openai.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}]
        )
        
        return response.choices[0].message.content

# Usage
rag = SimpleRAG()
rag.add_documents(["Doc 1 content...", "Doc 2 content..."])
answer = rag.query("What is discussed?")

Adding Features

Streaming

def query_stream(self, question, top_k=3):
    context = self._retrieve(question, top_k)
    prompt = self._build_prompt(context, question)
    
    stream = self.openai.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        stream=True
    )
    
    for chunk in stream:
        if chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content

Caching

from functools import lru_cache

class CachedRAG(SimpleRAG):
    @lru_cache(maxsize=1000)
    def _embed_cached(self, text):
        return tuple(self._embed([text])[0])

Comparison Summary

Approach	Pros	Cons
LangChain	Feature-rich, composable	Complex, large
LlamaIndex	Data-focused, easy	Less flexible
Custom	Full control, minimal	More work

---

Next Steps

Test your knowledge with the Module 6 Quiz!