Building Your First
RAG System

from langchain.text_splitter import (
    RecursiveCharacterTextSplitter,
    CharacterTextSplitter,
    TokenTextSplitter,
    MarkdownHeaderTextSplitter
)

class SmartChunker:
    """Intelligent document chunking with multiple strategies"""
    
    def __init__(self):
        self.splitters = {
            "recursive": RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200,
                separators=["\n\n", "\n", ". ", ", ", " ", ""]
            ),
            "token": TokenTextSplitter(
                chunk_size=500,
                chunk_overlap=50
            ),
            "semantic": self._create_semantic_splitter()
        }
    
    def _create_semantic_splitter(self):
        """Create splitter that preserves semantic boundaries"""
        return RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=[
                "\n## ",     # Markdown headers
                "\n### ",    
                "\n\n",     # Paragraphs
                "\n",        # Lines
                ". ",         # Sentences
                "? ",         # Questions
                "! ",         # Exclamations
                "; ",         # Semicolons
                ", ",         # Commas
                " ",          # Words
                ""            # Characters
            ]
        )
    
    def chunk_by_document_type(self, document):
        """Choose chunking strategy based on document type"""
        file_extension = document.metadata.get("source", "").split(".")[-1]
        
        if file_extension == "md":
            # Use markdown-aware splitting
            headers_to_split_on = [
                ("#", "Header 1"),
                ("##", "Header 2"),
                ("###", "Header 3"),
            ]
            markdown_splitter = MarkdownHeaderTextSplitter(
                headers_to_split_on=headers_to_split_on
            )
            return markdown_splitter.split_text(document.page_content)
        
        elif file_extension in ["py", "js", "java"]:
            # Code files need different handling
            return self._chunk_code(document)
        
        else:
            # Default recursive splitting
            return self.splitters["recursive"].split_documents([document])
    
    def _chunk_code(self, document):
        """Special handling for code files"""
        # Split by functions/classes while preserving context
        code_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1500,
            chunk_overlap=300,
            separators=[
                "\nclass ",
                "\ndef ",
                "\n\n",
                "\n",
                " ",
                ""
            ]
        )
        return code_splitter.split_documents([document])
    
    def experimental_sliding_window(self, text, window_size=500, step=250):
        """Sliding window approach for maximum context preservation"""
        chunks = []
        for i in range(0, len(text), step):
            chunk = text[i:i + window_size]
            if len(chunk) > 100:  # Minimum chunk size
                chunks.append({
                    "content": chunk,
                    "start": i,
                    "end": min(i + window_size, len(text))
                })
        return chunks

# Demo: Testing Different Chunking Strategies
print("=== Smart Chunking Demo ===")

# Create sample documents of different types
sample_documents = [
    {
        "content": """# Introduction to RAG Systems

RAG (Retrieval Augmented Generation) represents a significant advancement in AI systems. Unlike traditional language models that rely solely on their training data, RAG systems can access external knowledge bases to provide more accurate and up-to-date information.

## How RAG Works

The RAG process involves several key steps:

1. **Document Processing**: Raw documents are processed and split into manageable chunks
2. **Embedding Generation**: Each chunk is converted into vector embeddings
3. **Vector Storage**: Embeddings are stored in specialized vector databases
4. **Query Processing**: User queries are converted to embeddings
5. **Retrieval**: Similar chunks are retrieved based on semantic similarity
6. **Generation**: Retrieved context is used to generate accurate responses

## Benefits

RAG systems offer several advantages over traditional approaches:
- Reduced hallucinations in AI responses
- Access to up-to-date information
- Source attribution for answers
- Domain-specific knowledge integration""",
        "type": "markdown",
        "source": "rag_guide.md"
    },
    {
        "content": """def create_vector_store(documents, embeddings):
    '''Create a vector store from processed documents'''
    chunks = []
    
    for doc in documents:
        # Split document into chunks
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        doc_chunks = text_splitter.split_documents([doc])
        chunks.extend(doc_chunks)
    
    # Create vector store
    vector_store = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory="./vector_db"
    )
    
    return vector_store

class RAGPipeline:
    def __init__(self, vector_store):
        self.vector_store = vector_store
        self.retriever = vector_store.as_retriever(search_kwargs={"k": 5})
    
    def query(self, question):
        retrieved_docs = self.retriever.get_relevant_documents(question)
        return retrieved_docs""",
        "type": "code",
        "source": "rag_pipeline.py"
    }
]

# Initialize chunker
chunker = SmartChunker()

for doc_info in sample_documents:
    print(f"\n--- Processing {doc_info['source']} ({doc_info['type']}) ---")
    
    # Create mock document object
    class MockDocument:
        def __init__(self, content, source):
            self.page_content = content
            self.metadata = {"source": source}
    
    doc = MockDocument(doc_info["content"], doc_info["source"])
    
    # Test different chunking strategies
    print(f"Original length: {len(doc.page_content)} characters")
    
    # Try recursive splitter
    recursive_chunks = chunker.splitters["recursive"].split_documents([doc])
    print(f"Recursive chunks: {len(recursive_chunks)}")
    for i, chunk in enumerate(recursive_chunks[:2]):  # Show first 2 chunks
        print(f"  Chunk {i+1}: {len(chunk.page_content)} chars - '{chunk.page_content[:80]}...'")
    
    # Try document-type specific chunking
    if doc_info["type"] == "markdown":
        print("\nMarkdown-aware chunking:")
        try:
            md_chunks = chunker.chunk_by_document_type(doc)
            print(f"Markdown chunks: {len(md_chunks)}")
            for i, chunk in enumerate(md_chunks[:2]):
                chunk_text = chunk.page_content if hasattr(chunk, 'page_content') else str(chunk)
                print(f"  MD Chunk {i+1}: {len(chunk_text)} chars")
        except Exception as e:
            print(f"  Markdown chunking fallback to recursive")
    
    # Test sliding window
    print("\nSliding window chunking:")
    sliding_chunks = chunker.experimental_sliding_window(doc.page_content, window_size=300, step=150)
    print(f"Sliding window chunks: {len(sliding_chunks)}")
    for i, chunk in enumerate(sliding_chunks[:3]):
        print(f"  Window {i+1}: chars {chunk['start']}-{chunk['end']} (length: {len(chunk['content'])})")

print("\n=== Chunking Strategy Comparison ===")
test_text = "This is a sample document. It contains multiple sentences. Each sentence provides different information. We want to test how different chunking strategies handle this content. The goal is to preserve semantic meaning while creating manageable chunks."

strategies = {
    "recursive": chunker.splitters["recursive"],
    "token": chunker.splitters["token"]
}

for strategy_name, splitter in strategies.items():
    chunks = splitter.split_text(test_text)
    print(f"\n{strategy_name.title()} strategy:")
    print(f"  Number of chunks: {len(chunks)}")
    for i, chunk in enumerate(chunks):
        print(f"  Chunk {i+1}: '{chunk}'")

from typing import List, Dict
import tiktoken

class ContextManager:
    """Learn how to manage context windows for RAG systems step-by-step"""
    
    def __init__(self, model_name="gpt-3.5-turbo", max_tokens=4000):
        self.encoding = tiktoken.encoding_for_model(model_name)
        self.max_tokens = max_tokens
        self.reserved_tokens = 500  # Reserve for response
        
    def count_tokens(self, text: str) -> int:
        """Count tokens in text"""
        return len(self.encoding.encode(text))
    
    def optimize_context(self, 
                        documents: List[Dict], 
                        query: str,
                        strategy: str = "relevance") -> List[Dict]:
        """Optimize document selection for context window"""
        
        query_tokens = self.count_tokens(query)
        available_tokens = self.max_tokens - self.reserved_tokens - query_tokens
        
        if strategy == "relevance":
            return self._relevance_based_selection(documents, available_tokens)
        elif strategy == "diversity":
            return self._relevance_based_selection(documents, available_tokens)  # Simplified for demo
        elif strategy == "recency":
            return self._relevance_based_selection(documents, available_tokens)  # Simplified for demo
        else:
            return self._relevance_based_selection(documents, available_tokens)  # Simplified for demo
    
    def _relevance_based_selection(self, documents, available_tokens):
        """Select most relevant documents within token limit"""
        selected = []
        current_tokens = 0
        
        # Sort by relevance score
        sorted_docs = sorted(
            documents, 
            key=lambda x: x.get("score", 0), 
            reverse=True
        )
        
        for doc in sorted_docs:
            doc_tokens = self.count_tokens(doc["content"])
            if current_tokens + doc_tokens <= available_tokens:
                selected.append(doc)
                current_tokens += doc_tokens
            else:
                # Try to fit partial document
                remaining_tokens = available_tokens - current_tokens
                if remaining_tokens > 100:  # Minimum useful chunk
                    truncated_content = self._truncate_to_tokens(
                        doc["content"], 
                        remaining_tokens
                    )
                    doc_copy = doc.copy()
                    doc_copy["content"] = truncated_content
                    doc_copy["truncated"] = True
                    selected.append(doc_copy)
                break
        
        return selected
    
    def _truncate_to_tokens(self, text, max_tokens):
        """Truncate text to fit within token limit"""
        tokens = self.encoding.encode(text)
        if len(tokens) <= max_tokens:
            return text
        
        truncated_tokens = tokens[:max_tokens]
        return self.encoding.decode(truncated_tokens)
    
    def create_context_prompt(self, documents, query):
        """Create optimized prompt with context"""
        context_parts = []
        
        for i, doc in enumerate(documents):
            source = doc.get("metadata", {}).get("source", "Unknown")
            truncated = " (truncated)" if doc.get("truncated") else ""
            
            context_parts.append(
                f"[Document {i+1} - {source}{truncated}]\n"
                f"{doc['content']}\n"
            )
        
        context = "\n---\n".join(context_parts)
        
        return f"""Based on the following context, answer the question.
        
Context:
{context}

Question: {query}

Answer:"""

# Usage example - Context Optimization Demo
context_manager = ContextManager(max_tokens=4000)

# Simulate retrieved documents with realistic content
documents = [
    {
        "content": "RAG (Retrieval Augmented Generation) is a technique that combines information retrieval with text generation. It helps AI models access external knowledge to provide more accurate and up-to-date responses by retrieving relevant documents before generating answers.",
        "score": 0.95,
        "metadata": {"source": "rag_guide.pdf"}
    },
    {
        "content": "Vector databases like ChromaDB, Pinecone, and Weaviate are essential components of RAG systems. They store document embeddings and enable fast similarity search to find relevant context for user queries.",
        "score": 0.87,
        "metadata": {"source": "vector_db_tutorial.txt"}
    },
    {
        "content": "Traditional search relies on keyword matching, while RAG uses semantic similarity. This allows the system to understand meaning and context, making it more effective for question answering applications.",
        "score": 0.82,
        "metadata": {"source": "semantic_search.md"}
    },
]

print("=== Context Optimization Demo ===")
print(f"Maximum tokens: {context_manager.max_tokens}")
print(f"Reserved for response: {context_manager.reserved_tokens}")

query = "What is RAG and how does it work?"
query_tokens = context_manager.count_tokens(query)
print(f"Query tokens: {query_tokens}")

# Test different strategies
for strategy in ["relevance", "diversity"]:
    print(f"\n--- Strategy: {strategy} ---")
    
    # Optimize context
    optimized_docs = context_manager.optimize_context(
        documents, 
        query,
        strategy=strategy
    )
    
    print(f"Selected documents: {len(optimized_docs)}")
    
    total_tokens = query_tokens + context_manager.reserved_tokens
    for i, doc in enumerate(optimized_docs):
        doc_tokens = context_manager.count_tokens(doc["content"])
        total_tokens += doc_tokens
        truncated = " (TRUNCATED)" if doc.get("truncated") else ""
        print(f"  Doc {i+1}: {doc_tokens} tokens{truncated}")
    
    print(f"Total tokens used: {total_tokens}/{context_manager.max_tokens}")
    
    # Create and show sample prompt
    if strategy == "relevance":  # Show prompt for one strategy
        prompt = context_manager.create_context_prompt(optimized_docs, query)
        print(f"\nSample prompt preview (first 200 chars):")
        print(f"{prompt[:200]}...")

print("\n=== Token Counting Examples ===")
sample_texts = [
    "Hello world",
    "This is a longer sentence with more tokens to demonstrate counting.",
    "RAG systems are amazing!"
]

for text in sample_texts:
    token_count = context_manager.count_tokens(text)
    print(f"Text: '{text}' → {token_count} tokens")