Basic RAG Pipeline

import tiktoken

enc = tiktoken.get_encoding("cl100k_base")

def count_tokens(text: str) -> int:
    return len(enc.encode(text))

# ── Strategy 1: Fixed-size with overlap ──────────────────────
def chunk_fixed(text: str, size: int = 512, overlap: int = 50) -> list[str]:
    """Split into token-counted chunks with overlap for context continuity."""
    tokens = enc.encode(text)
    chunks = []
    start = 0
    while start < len(tokens):
        end = start + size
        chunk_tokens = tokens[start:end]
        chunks.append(enc.decode(chunk_tokens))
        start = end - overlap  # slide back by overlap
    return chunks

# ── Strategy 2: Sentence-based grouping ──────────────────────
import re
def chunk_sentences(text: str, max_tokens: int = 512) -> list[str]:
    """Group sentences until token budget is reached."""
    sentences = re.split(r'(?<=[.!?])\s+', text)
    chunks, current = [], []
    current_len = 0
    for sent in sentences:
        sent_len = count_tokens(sent)
        if current_len + sent_len > max_tokens and current:
            chunks.append(" ".join(current))
            current, current_len = [], 0
        current.append(sent)
        current_len += sent_len
    if current:
        chunks.append(" ".join(current))
    return chunks

# ── Strategy 3: Recursive splitting ──────────────────────────
def chunk_recursive(text: str, max_tokens: int = 512,
                    separators: list[str] = None) -> list[str]:
    """Try each separator in order; recurse on pieces that are too large."""
    if separators is None:
        separators = ["\n\n", "\n", ". ", " "]
    if count_tokens(text) <= max_tokens:
        return [text.strip()] if text.strip() else []
    sep = separators[0] if separators else " "
    parts = text.split(sep)
    chunks, current = [], ""
    for part in parts:
        candidate = current + sep + part if current else part
        if count_tokens(candidate) > max_tokens:
            if current:
                chunks.append(current.strip())
            # If single part exceeds limit, recurse with finer separator
            if count_tokens(part) > max_tokens and len(separators) > 1:
                chunks.extend(chunk_recursive(part, max_tokens, separators[1:]))
                current = ""
            else:
                current = part
        else:
            current = candidate
    if current.strip():
        chunks.append(current.strip())
    return chunks

# ── Strategy 4: Semantic chunking ────────────────────────────
import numpy as np
def chunk_semantic(text: str, model, threshold: float = 0.5,
                   max_tokens: int = 1024) -> list[str]:
    """Split where cosine similarity between consecutive sentences drops."""
    sentences = re.split(r'(?<=[.!?])\s+', text)
    if len(sentences) <= 1:
        return [text]
    embeddings = model.encode(sentences)
    # Compute cosine similarity between consecutive sentences
    similarities = [
        np.dot(embeddings[i], embeddings[i+1]) /
        (np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[i+1]))
        for i in range(len(embeddings) - 1)
    ]
    # Split where similarity drops below threshold
    chunks, current = [], [sentences[0]]
    for i, sim in enumerate(similarities):
        if sim < threshold or count_tokens(" ".join(current)) > max_tokens:
            chunks.append(" ".join(current))
            current = []
        current.append(sentences[i + 1])
    if current:
        chunks.append(" ".join(current))
    return chunks

"""
Minimal RAG pipeline — no frameworks, just numpy + openai.
pip install openai numpy tiktoken
"""
import numpy as np
import tiktoken
from openai import OpenAI

client = OpenAI()  # reads OPENAI_API_KEY from env
enc = tiktoken.get_encoding("cl100k_base")

# ── 1. CHUNK ─────────────────────────────────────────────────
def chunk_text(text: str, max_tokens: int = 512, overlap: int = 50) -> list[str]:
    """Fixed-size chunking with token-based overlap."""
    tokens = enc.encode(text)
    chunks = []
    start = 0
    while start < len(tokens):
        end = min(start + max_tokens, len(tokens))
        chunks.append(enc.decode(tokens[start:end]))
        if end == len(tokens):
            break
        start = end - overlap
    return chunks

# ── 2. EMBED ─────────────────────────────────────────────────
def embed(texts: list[str], model: str = "text-embedding-3-small") -> np.ndarray:
    """Embed a batch of texts. Returns (N, dim) array."""
    response = client.embeddings.create(input=texts, model=model)
    return np.array([e.embedding for e in response.data])

# ── 3. INDEX (in-memory for demo; use a vector DB in production) ─
class SimpleVectorIndex:
    def __init__(self):
        self.chunks: list[str] = []
        self.vectors: np.ndarray | None = None

    def add(self, texts: list[str]):
        """Add chunks to the index."""
        self.chunks.extend(texts)
        new_vecs = embed(texts)
        if self.vectors is None:
            self.vectors = new_vecs
        else:
            self.vectors = np.vstack([self.vectors, new_vecs])

    def search(self, query: str, k: int = 5, threshold: float = 0.3
               ) -> list[tuple[str, float]]:
        """Return top-k chunks above similarity threshold."""
        q_vec = embed([query])[0]
        # Cosine similarity (vectors are already normalized by OpenAI)
        scores = self.vectors @ q_vec
        # Filter by threshold, then take top-k
        indices = np.where(scores >= threshold)[0]
        indices = indices[np.argsort(scores[indices])[::-1]][:k]
        return [(self.chunks[i], float(scores[i])) for i in indices]

# ── 4. GENERATE ──────────────────────────────────────────────
def generate(query: str, context_chunks: list[tuple[str, float]],
             model: str = "gpt-4o-mini") -> str:
    """Generate answer grounded in retrieved context."""
    context = "\n\n".join(
        f"[{i+1}] (score: {score:.2f}) {chunk}"
        for i, (chunk, score) in enumerate(context_chunks)
    )
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": (
                "Answer the user's question using ONLY the provided context. "
                "Cite sources using [1], [2], etc. If the context doesn't contain "
                "the answer, say 'I don't have enough information to answer that.'"
            )},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
        ],
        temperature=0.1,  # Low temperature for factual accuracy
    )
    return response.choices[0].message.content

# ── USAGE ────────────────────────────────────────────────────
# Index some documents
index = SimpleVectorIndex()
documents = [
    "The Eiffel Tower was built in 1889 for the World's Fair in Paris...",
    "Python was created by Guido van Rossum and released in 1991...",
    "The mitochondria is the powerhouse of the cell, producing ATP...",
]
for doc in documents:
    chunks = chunk_text(doc)
    index.add(chunks)

# Query
query = "When was Python created?"
results = index.search(query, k=3)
answer = generate(query, results)
print(answer)
# → "Python was created by Guido van Rossum and released in 1991 [1]."

Production RAG Prompt Template