Building a RAG Pipeline

from langchain_community.document_loaders import (
    PyPDFLoader,
    WebBaseLoader,
    TextLoader,
    DirectoryLoader,
    UnstructuredMarkdownLoader
)

# --- Load a PDF ---
pdf_loader = PyPDFLoader("docs/architecture_guide.pdf")
pdf_pages = pdf_loader.load()  # One Document per page
print(f"PDF: {len(pdf_pages)} pages")
print(f"Page 1 content: {pdf_pages[0].page_content[:200]}")
print(f"Page 1 metadata: {pdf_pages[0].metadata}")
# {'source': 'docs/architecture_guide.pdf', 'page': 0}

# --- Load a web page ---
web_loader = WebBaseLoader("https://docs.python.org/3/tutorial/classes.html")
web_docs = web_loader.load()
print(f"Web page: {len(web_docs[0].page_content)} chars")

# --- Load all markdown files from a directory ---
dir_loader = DirectoryLoader(
    "docs/",
    glob="**/*.md",
    loader_cls=UnstructuredMarkdownLoader,
    show_progress=True
)
all_docs = dir_loader.load()
print(f"Directory: {len(all_docs)} documents loaded")

# --- Custom loader for any source ---
class NotionLoader:
    """Example: load from Notion API."""
    def __init__(self, database_id: str, token: str):
        self.database_id = database_id
        self.token = token
    
    def load(self):
        # Fetch pages from Notion API
        # Convert blocks to text
        # Return list of Document objects
        pass

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
import chromadb

# --- Chunking ---
splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", ". ", " ", ""]
)

# Split all documents, preserving metadata
all_chunks = []
for doc in all_docs:
    chunks = splitter.split_text(doc.page_content)
    for i, chunk in enumerate(chunks):
        all_chunks.append({
            "content": chunk,
            "metadata": {
                **doc.metadata,
                "chunk_index": i,
                "char_count": len(chunk)
            }
        })

print(f"Total chunks: {len(all_chunks)}")

# --- Embedding + Storage ---
client = chromadb.PersistentClient(path="./rag_db")
collection = client.get_or_create_collection(
    name="documentation",
    metadata={"hnsw:space": "cosine"}
)

# Batch insert (ChromaDB handles embedding internally if configured)
# Or embed manually for more control:
from openai import OpenAI
openai_client = OpenAI()

BATCH_SIZE = 100
for i in range(0, len(all_chunks), BATCH_SIZE):
    batch = all_chunks[i:i + BATCH_SIZE]
    
    # Embed batch
    response = openai_client.embeddings.create(
        input=[c["content"] for c in batch],
        model="text-embedding-3-small"
    )
    embeddings = [d.embedding for d in response.data]
    
    # Insert into ChromaDB
    collection.add(
        ids=[f"chunk_{i+j}" for j in range(len(batch))],
        documents=[c["content"] for c in batch],
        embeddings=embeddings,
        metadatas=[c["metadata"] for c in batch]
    )
    
    print(f"Ingested {min(i + BATCH_SIZE, len(all_chunks))}/{len(all_chunks)}")

print(f"✅ Ingestion complete: {collection.count()} chunks indexed")

def retrieve(query: str, top_k: int = 5, filter_metadata: dict = None):
    """Retrieve relevant chunks for a query."""
    # Build query parameters
    query_params = {
        "query_texts": [query],
        "n_results": top_k,
    }
    if filter_metadata:
        query_params["where"] = filter_metadata
    
    results = collection.query(**query_params)
    
    # Format results
    retrieved = []
    for doc, meta, dist in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0]
    ):
        retrieved.append({
            "content": doc,
            "source": meta.get("source", "unknown"),
            "score": 1 - dist,  # Convert distance to similarity
        })
    
    return retrieved

# Test retrieval
results = retrieve("How do Python classes handle inheritance?")
for r in results:
    print(f"[{r['score']:.3f}] ({r['source']}) {r['content'][:100]}...")

from openai import OpenAI

client = OpenAI()

def generate_answer(query: str, retrieved_chunks: list) -> dict:
    """Generate a grounded answer with source citations."""
    
    # Format context with source labels
    context_parts = []
    for i, chunk in enumerate(retrieved_chunks):
        context_parts.append(
            f"[Source {i+1}: {chunk['source']}]\n{chunk['content']}"
        )
    context = "\n\n---\n\n".join(context_parts)
    
    # System prompt enforcing grounding
    system_prompt = """You are a helpful assistant that answers questions based ONLY on the provided context. 

Rules:
1. Only use information from the provided sources
2. Cite sources using [Source N] notation
3. If the context doesn't contain the answer, say "I don't have enough information to answer this"
4. Be concise but complete
5. If multiple sources agree, cite all of them"""

    user_prompt = f"""Context:
{context}

---

Question: {query}

Answer (with citations):"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        temperature=0.1,  # Low temp for factual answers
        max_tokens=500
    )
    
    answer = response.choices[0].message.content
    
    return {
        "answer": answer,
        "sources": [
            {"source": c["source"], "score": c["score"], "excerpt": c["content"][:100]}
            for c in retrieved_chunks
        ],
        "model": "gpt-4o-mini",
        "tokens_used": response.usage.total_tokens
    }

"""Complete RAG Pipeline - Production-Ready Structure"""
from dataclasses import dataclass
from typing import Optional
import chromadb
from openai import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter

@dataclass
class RAGResponse:
    answer: str
    sources: list
    tokens_used: int
    retrieval_scores: list[float]

class RAGPipeline:
    def __init__(
        self,
        db_path: str = "./rag_db",
        collection_name: str = "documents",
        embedding_model: str = "text-embedding-3-small",
        generation_model: str = "gpt-4o-mini",
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        top_k: int = 5
    ):
        self.openai = OpenAI()
        self.chroma = chromadb.PersistentClient(path=db_path)
        self.collection = self.chroma.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )
        self.embedding_model = embedding_model
        self.generation_model = generation_model
        self.top_k = top_k
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )
    
    def ingest(self, text: str, source: str, metadata: dict = None):
        """Ingest a document into the vector store."""
        chunks = self.splitter.split_text(text)
        
        # Embed
        response = self.openai.embeddings.create(
            input=chunks, model=self.embedding_model
        )
        embeddings = [d.embedding for d in response.data]
        
        # Store
        base_id = f"{source}_{hash(text) % 10000}"
        self.collection.add(
            ids=[f"{base_id}_chunk{i}" for i in range(len(chunks))],
            documents=chunks,
            embeddings=embeddings,
            metadatas=[
                {**(metadata or {}), "source": source, "chunk_idx": i}
                for i in range(len(chunks))
            ]
        )
        return len(chunks)
    
    def query(self, question: str, filter_metadata: dict = None) -> RAGResponse:
        """Full RAG: retrieve → generate with citations."""
        # Retrieve
        results = self.collection.query(
            query_texts=[question],
            n_results=self.top_k,
            where=filter_metadata
        )
        
        if not results["documents"][0]:
            return RAGResponse(
                answer="No relevant documents found.",
                sources=[], tokens_used=0, retrieval_scores=[]
            )
        
        # Build context
        context_parts = []
        sources = []
        scores = []
        for i, (doc, meta, dist) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]
        )):
            score = 1 - dist
            context_parts.append(f"[Source {i+1}: {meta['source']}]\n{doc}")
            sources.append(meta["source"])
            scores.append(score)
        
        context = "\n\n---\n\n".join(context_parts)
        
        # Generate
        response = self.openai.chat.completions.create(
            model=self.generation_model,
            messages=[
                {"role": "system", "content": (
                    "Answer based ONLY on the provided context. "
                    "Cite sources as [Source N]. If unsure, say so."
                )},
                {"role": "user", "content": (
                    f"Context:\n{context}\n\n"
                    f"Question: {question}\n\nAnswer:"
                )}
            ],
            temperature=0.1
        )
        
        return RAGResponse(
            answer=response.choices[0].message.content,
            sources=sources,
            tokens_used=response.usage.total_tokens,
            retrieval_scores=scores
        )

# --- Usage ---
rag = RAGPipeline()

# Ingest documents
rag.ingest(open("docs/api_guide.md").read(), source="api_guide.md")
rag.ingest(open("docs/architecture.md").read(), source="architecture.md")
rag.ingest(open("docs/deployment.md").read(), source="deployment.md")

print(f"Total chunks: {rag.collection.count()}")

# Query
result = rag.query("How do I deploy the application to production?")
print(f"\nAnswer: {result.answer}")
print(f"\nSources: {result.sources}")
print(f"Scores: {[f'{s:.3f}' for s in result.retrieval_scores]}")
print(f"Tokens: {result.tokens_used}")

import tiktoken

def fit_context(chunks: list[str], max_tokens: int = 3000) -> list[str]:
    """Select chunks that fit within token budget."""
    encoder = tiktoken.encoding_for_model("gpt-4o-mini")
    
    selected = []
    total_tokens = 0
    
    for chunk in chunks:
        chunk_tokens = len(encoder.encode(chunk))
        if total_tokens + chunk_tokens > max_tokens:
            break
        selected.append(chunk)
        total_tokens += chunk_tokens
    
    return selected

# Strategies for large retrieval sets:
# 1. Stuff: Include all chunks (simple, works for small sets)
# 2. Map-Reduce: Summarize each chunk, then combine summaries
# 3. Refine: Process chunks sequentially, refining the answer
# 4. Rerank + Top-K: Score all, take only the best N that fit