RAG That Doesn't Suck: Build a Search That Beats 'Chat With Your Docs'

# Clone
git clone https://github.com/yourusername/rag-that-works.git
cd rag-that-works
 
# Install
pip install -r requirements.txt
 
# Add your API key
export ANTHROPIC_API_KEY="sk-ant-..."
 
# Index a document
[Python](https://www.python.org/) cli.py index ./docs/employee_handbook.md
 
# Ask a question
python cli.py search "What's the PTO policy for contractors?"

Answer:
Contractors receive 15 days of paid time off annually [1], prorated based on 
their start date [2]. Unused PTO does not roll over to the following year [1].
 
Sources:
  [1]: employee_handbook.md (section: "Contractor Benefits", chars 2847-3201)
  [2]: employee_handbook.md (section: "Onboarding", chars 892-1104)

[1] "...membership fees for professional organizations are reimbursable 
    up to $500 annually. Submit receipts through Expensify within 30..."
    
[2] "...workspace ergonomics are important. The company provides a $200 
    home office stipend for remote employees to purchase..."
    
[3] "...co-working is encouraged for team collaboration. Regional teams 
    should coordinate quarterly meetups at designated..."

[1] "Co-working Space Reimbursement: Full-time remote employees may 
    expense co-working space memberships up to $300/month. Part-time 
    remote employees are eligible for up to $150/month. Requires 
    manager approval and monthly receipt submission."
    
[2] "Expense Eligibility: To qualify for workspace reimbursements, 
    employees must be designated as 'remote' or 'hybrid' in Workday. 
    Office-based employees are not eligible for co-working benefits."

┌─────────────────────────────────────────────────────────────────┐
│                        USER QUERY                                │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                    QUERY EMBEDDING                               │
│              (same model as document chunks)                     │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                 VECTOR RETRIEVAL (Top 20)                        │
│              Coarse filtering—cast a wide net                    │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                   CROSS-ENCODER RERANKING                        │
│              Scores each chunk against query                     │
│              Returns top 5—the actually relevant ones            │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                     LLM GENERATION                               │
│        Answer with inline citations [1], [2], etc.              │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                    CITATION MAPPING                              │
│          Link each [n] to exact chunk + character offset        │
└─────────────────────────────────────────────────────────────────┘

# chunker.py
import re
from dataclasses import dataclass
from typing import List
 
@dataclass
class Chunk:
    text: str
    source: str
    chunk_id: int
    start_char: int
    end_char: int
    section: str = ""  # Optional: section header for context
 
class SemanticChunker:
    def __init__(
        self,
        target_size: int = 512,      # target tokens (roughly)
        max_size: int = 1024,         # hard ceiling
        overlap_sentences: int = 1    # context carry-over
    ):
        self.target_size = target_size
        self.max_size = max_size
        self.overlap_sentences = overlap_sentences
        self.chars_per_token = 4  # Rough estimate for English
    
    def chunk(self, text: str, source: str) -> List[Chunk]:
        """
        Split text into semantic chunks.
        
        Strategy:
        1. Split on major boundaries (double newlines, headers)
        2. Split long sections on sentence boundaries
        3. Merge tiny chunks with neighbors
        4. Track section headers for citation context
        """
        text = re.sub(r'\n{3,}', '\n\n', text.strip())
        
        # Extract section structure
        sections = self._split_by_headers(text)
        
        all_chunks = []
        for section_title, section_text in sections:
            section_chunks = self._chunk_section(section_text, source, section_title)
            all_chunks.extend(section_chunks)
        
        # Assign sequential IDs and character offsets
        char_offset = 0
        for i, chunk in enumerate(all_chunks):
            chunk.chunk_id = i
            start = text.find(chunk.text[:50], char_offset)
            if start == -1:
                start = char_offset
            chunk.start_char = start
            chunk.end_char = start + len(chunk.text)
            char_offset = start + len(chunk.text) // 2
        
        return all_chunks
    
    def _split_by_headers(self, text: str) -> List[tuple]:
        """Split text by headers, returning (header, content) tuples."""
        # Match markdown headers or ALL CAPS lines
        header_pattern = r'^(#{1,6}\s+.+|[A-Z][A-Z\s]{10,})$'
        
        sections = []
        current_header = ""
        current_content = []
        
        for line in text.split('\n'):
            if re.match(header_pattern, line.strip()):
                if current_content:
                    sections.append((current_header, '\n'.join(current_content)))
                current_header = line.strip().lstrip('#').strip()
                current_content = []
            else:
                current_content.append(line)
        
        if current_content:
            sections.append((current_header, '\n'.join(current_content)))
        
        return sections if sections else [("", text)]
    
    def _chunk_section(self, text: str, source: str, section: str) -> List[Chunk]:
        """Chunk a single section."""
        paragraphs = re.split(r'\n\n+', text)
        
        # Split large paragraphs on sentences
        pieces = []
        for para in paragraphs:
            if self._estimate_tokens(para) > self.max_size:
                pieces.extend(self._split_sentences(para))
            else:
                pieces.append(para)
        
        # Merge small pieces
        merged = self._merge_pieces(pieces)
        
        return [
            Chunk(text=t.strip(), source=source, chunk_id=0, 
                  start_char=0, end_char=0, section=section)
            for t in merged if t.strip()
        ]
    
    def _estimate_tokens(self, text: str) -> int:
        return len(text) // self.chars_per_token
    
    def _split_sentences(self, text: str) -> List[str]:
        pattern = r'(?<=[.!?])\s+(?=[A-Z])'
        return [s.strip() for s in re.split(pattern, text) if s.strip()]
    
    def _merge_pieces(self, pieces: List[str]) -> List[str]:
        if not pieces:
            return []
        
        merged = []
        current = pieces[0]
        
        for piece in pieces[1:]:
            combined_size = self._estimate_tokens(current + " " + piece)
            
            if combined_size <= self.target_size:
                current = current + "\n\n" + piece
            elif self._estimate_tokens(current) < self.target_size // 3:
                current = current + "\n\n" + piece
            else:
                merged.append(current)
                current = piece
        
        merged.append(current)
        return merged

# embedder.py
import numpy as np
from typing import List
from sentence_transformers import SentenceTransformer
 
class Embedder:
    def __init__(self, model_name: str = "BAAI/bge-base-en-v1.5"):
        """
        BGE-base: trained on retrieval tasks, handles query/doc asymmetry,
        768 dimensions, instruction-following for query enhancement.
        """
        self.model = SentenceTransformer(model_name)
        self.dimension = 768
    
    def embed_documents(self, texts: List[str]) -> np.ndarray:
        """Embed document chunks. No prefix needed."""
        return self.model.encode(
            texts,
            normalize_embeddings=True,
            show_progress_bar=True
        )
    
    def embed_query(self, query: str) -> np.ndarray:
        """
        Embed a search query with instruction prefix.
        
        This is critical: BGE models gain 5-10% recall with the prefix.
        """
        instruction = "Represent this sentence for searching relevant passages: "
        return self.model.encode(
            instruction + query,
            normalize_embeddings=True
        )

# vector_store.py
import numpy as np
import faiss
import pickle
from pathlib import Path
from typing import List, Tuple, Optional
from dataclasses import asdict
 
from chunker import Chunk
from embedder import Embedder
 
class VectorStore:
    def __init__(self, embedder: Embedder, index_path: Optional[str] = None):
        self.embedder = embedder
        self.index: Optional[faiss.IndexFlatIP] = None
        self.chunks: List[Chunk] = []
        self.index_path = index_path
        
        if index_path and Path(index_path).exists():
            self.load(index_path)
    
    def add_documents(self, chunks: List[Chunk]):
        if not chunks:
            return
        
        texts = [c.text for c in chunks]
        embeddings = self.embedder.embed_documents(texts)
        
        if self.index is None:
            self.index = faiss.IndexFlatIP(embeddings.shape[1])
        
        self.index.add(embeddings.astype(np.float32))
        self.chunks.extend(chunks)
    
    def search(self, query: str, k: int = 20) -> List[Tuple[Chunk, float]]:
        if self.index is None or not self.chunks:
            return []
        
        query_emb = self.embedder.embed_query(query).reshape(1, -1).astype(np.float32)
        scores, indices = self.index.search(query_emb, min(k, len(self.chunks)))
        
        return [(self.chunks[idx], float(score)) 
                for score, idx in zip(scores[0], indices[0]) 
                if idx < len(self.chunks)]
    
    def save(self, path: str):
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        if self.index:
            faiss.write_index(self.index, str(path / "index.faiss"))
        with open(path / "chunks.pkl", "wb") as f:
            pickle.dump([asdict(c) for c in self.chunks], f)
    
    def load(self, path: str):
        path = Path(path)
        if (path / "index.faiss").exists():
            self.index = faiss.read_index(str(path / "index.faiss"))
        if (path / "chunks.pkl").exists():
            with open(path / "chunks.pkl", "rb") as f:
                self.chunks = [Chunk(**d) for d in pickle.load(f)]

# reranker.py
from typing import List, Tuple
from sentence_transformers import CrossEncoder
 
from chunker import Chunk
 
class Reranker:
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        """
        Cross-encoder: processes query+doc together, catches interactions
        bi-encoders miss. Slower but far more accurate.
        
        For critical accuracy: use BAAI/bge-reranker-base (10x larger)
        """
        self.model = CrossEncoder(model_name)
    
    def rerank(
        self, 
        query: str, 
        chunks: List[Tuple[Chunk, float]], 
        top_k: int = 5
    ) -> List[Tuple[Chunk, float]]:
        if not chunks:
            return []
        
        pairs = [(query, chunk.text) for chunk, _ in chunks]
        scores = self.model.predict(pairs)
        
        reranked = [(chunks[i][0], float(scores[i])) for i in range(len(chunks))]
        reranked.sort(key=lambda x: x[1], reverse=True)
        
        return reranked[:top_k]

# generator.py
import re
from typing import List, Tuple, Dict, Any
from dataclasses import dataclass
import anthropic
 
from chunker import Chunk
 
@dataclass
class CitedAnswer:
    answer: str
    citations: List[Dict[str, Any]]
    chunks_used: List[Chunk]
 
class Generator:
    def __init__(self, api_key: str):
        self.client = anthropic.Anthropic(api_key=api_key)
    
    def generate(
        self, 
        query: str, 
        chunks: List[Tuple[Chunk, float]],
        model: str = "claude-sonnet-4-20250514"
    ) -> CitedAnswer:
        # Build numbered context
        context_parts = []
        chunk_map = {}
        
        for i, (chunk, score) in enumerate(chunks, 1):
            section_info = f" (section: \"{chunk.section}\")" if chunk.section else ""
            context_parts.append(
                f"[{i}] Source: {chunk.source}{section_info}\n{chunk.text}"
            )
            chunk_map[i] = chunk
        
        context = "\n\n---\n\n".join(context_parts)
        
        system_prompt = """You are a precise research assistant. Answer using ONLY the provided excerpts.
 
Rules:
1. Use ONLY information from the excerpts. No outside knowledge.
2. Cite with [1], [2], etc. Every factual claim needs a citation.
3. If excerpts don't answer the question, say so explicitly.
4. Be concise. Don't repeat information.
5. If sources conflict, note the discrepancy."""
 
        response = self.client.messages.create(
            model=model,
            max_tokens=1024,
            system=system_prompt,
            messages=[{"role": "user", "content": f"Excerpts:\n\n{context}\n\n---\n\nQuestion: {query}"}]
        )
        
        answer_text = response.content[0].text
        citations = self._extract_citations(answer_text, chunk_map)
        
        return CitedAnswer(
            answer=answer_text,
            citations=citations,
            chunks_used=[chunk_map[i] for i in chunk_map]
        )
    
    def _extract_citations(self, text: str, chunk_map: Dict[int, Chunk]) -> List[Dict]:
        refs = set(int(m) for m in re.findall(r'\[(\d+)\]', text))
        
        citations = []
        for ref in sorted(refs):
            if ref in chunk_map:
                chunk = chunk_map[ref]
                citations.append({
                    "reference": f"[{ref}]",
                    "source": chunk.source,
                    "section": chunk.section,
                    "start_char": chunk.start_char,
                    "end_char": chunk.end_char,
                    "text_preview": chunk.text[:200] + "..."
                })
            else:
                citations.append({
                    "reference": f"[{ref}]",
                    "error": "Hallucinated citation—refers to non-existent source"
                })
        
        return citations