The Semantic Chunking Engine (Advanced RAG)

1. Problem Statement & Scope

A junior engineer comes to you complaining that their Retrieval-Augmented Generation (RAG) application is hallucinating. They loaded 10,000 enterprise PDFs into a vector database, but when they ask the LLM a question, it returns garbled nonsense.

You look at their ingestion script and see they are using a naive text splitter: rigidly cutting the text every 500 tokens. They have sliced sentences in half, separated Python function definitions from their return statements, and completely obliterated the structural hierarchy of financial tables.

We need to build a Context-Aware Dynamic Chunking Engine.

Functional Requirements

• Format-Aware Routing: The system must dynamically choose the chunking strategy based on the MIME type (e.g., Markdown, Python, PDF, HTML).
• Semantic Preservation: Chunks must respect logical boundaries (paragraphs, AST nodes, sections) rather than arbitrary token counts.
• Metadata Injection: Every chunk must retain a pointer to its parent document, section header, and absolute position.

Non-Functional Requirements

• High Throughput: The pipeline must process and embed millions of documents per day asynchronously.
• Resilience: The pipeline must not crash when encountering corrupted PDFs or malformed code files.

2. High-Level Design (HLD)

We treat document ingestion as a highly decoupled Event-Driven Pipeline. We do not block the user’s upload request waiting for embeddings to generate.

3. The Viva: Pushing the Boundaries

4. Low-Level Design (LLD): The Strategy Pattern

Here is how we design the core chunking router in Python. We use the Strategy Design Pattern to ensure we can easily add new parsers (like a CSV chunker or an Audio Transcript chunker) without modifying the core routing logic.

from abc import ABC, abstractmethod
from typing import List, Dict, Any
import ast

# 1. The Strategy Interface
class BaseChunker(ABC):
    @abstractmethod
    def chunk(self, content: str, metadata: Dict[str, Any]) -> List[Dict[str, Any]]:
        pass

# 2. Concrete Strategy: Markdown Header Splitter
class MarkdownChunker(BaseChunker):
    def chunk(self, content: str, metadata: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Extremely simplified logic for brevity
        sections = content.split("\n## ")
        chunks = []
        for i, section in enumerate(sections):
            if not section.strip(): continue
            header = section.split("\n")[0]
            chunks.append({
                "text": f"## {section}" if i > 0 else section,
                "metadata": {**metadata, "header": header, "chunk_index": i}
            })
        return chunks

# 3. Concrete Strategy: Code AST Semantic Splitter
class PythonASTChunker(BaseChunker):
    def chunk(self, content: str, metadata: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Parses Python code and chunks strictly by Classes and Functions."""
        chunks = []
        try:
            tree = ast.parse(content)
            for node in tree.body:
                if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
                    # In a real system, you'd extract the exact string segment using node lineno
                    func_name = node.name
                    chunks.append({
                        "text": f"Code Block for {func_name}...", 
                        "metadata": {**metadata, "entity_type": type(node).__name__, "name": func_name}
                    })
        except SyntaxError:
            # Fallback to naive splitting if the code is malformed
            pass 
        return chunks

# 4. The Context/Router
class ChunkingEngine:
    def __init__(self):
        # Register strategies dynamically
        self.strategies = {
            "text/markdown": MarkdownChunker(),
            "text/x-python": PythonASTChunker(),
        }

    def process_document(self, mime_type: str, content: str, doc_id: str) -> List[Dict[str, Any]]:
        strategy = self.strategies.get(mime_type)
        if not strategy:
            raise ValueError(f"No chunking strategy configured for {mime_type}")
        
        base_metadata = {"doc_id": doc_id, "mime_type": mime_type}
        return strategy.chunk(content, base_metadata)

# Usage
engine = ChunkingEngine()
code_chunks = engine.process_document("text/x-python", "def hello():\n    print('world')", "doc_123")

5. Fault Tolerance & Edge Cases

• The Infinite Token Problem: What happens if the Markdown strategy extracts a single section under a ## header, but that section contains a massive 10,000-word table that exceeds the embedding model’s maximum token limit (e.g., 8191 tokens for OpenAI)? The Strategy architecture must implement a recursive fallback. If a semantic chunk exceeds the token limit, it must automatically recursively invoke a secondary, more aggressive strategy (like an NLP sentence splitter or a rigid recursive character splitter) to break the massive section down further before yielding it to the embedder.
• Poisoned Documents: PDF parsing libraries (like PyPDF2 or pdfplumber) are notoriously fragile. A corrupted font encoding in a PDF can cause the extraction thread to hang indefinitely, locking up the Kafka worker. All chunking operations must execute within strict timeout wrappers, and failed extractions must be routed to a Dead-Letter Queue (DLQ) for manual inspection or OCR fallback processing.