Chunking strategies

import tiktoken
from dataclasses import dataclass

enc = tiktoken.get_encoding("cl100k_base")

@dataclass
class Chunk:
    id: str
    text: str
    token_count: int
    source: str

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

def chunk_fixed(text: str, source: str, size: int = 512, overlap: int = 64) -> list[Chunk]:
    tokens = enc.encode(text)
    chunks: list[Chunk] = []
    start, i = 0, 0
    while start < len(tokens):
        end = min(start + size, len(tokens))
        piece = enc.decode(tokens[start:end])
        chunks.append(Chunk(
            id=f"{source}-{i}",
            text=piece,
            token_count=end - start,
            source=source,
        ))
        i += 1
        if end == len(tokens):
            break
        start = end - overlap  # slide window back by overlap
    return chunks

import nltk
import spacy
from chunking.tokens import token_len  # from fixed-size section

nltk.download("punkt", quiet=True)
nlp = spacy.blank("en")
nlp.add_pipe("sentencizer")

def sentences_nltk(text: str) -> list[str]:
    return nltk.sent_tokenize(text)

def sentences_spacy(text: str) -> list[str]:
    return [s.text.strip() for s in nlp(text).sents if s.text.strip()]

def chunk_by_sentences(text: str, source: str, max_tokens: int = 400) -> list[dict]:
    chunks, buf, i = [], [], 0
    for sent in sentences_spacy(text):
        trial = " ".join(buf + [sent])
        if token_len(trial) > max_tokens and buf:
            piece = " ".join(buf)
            chunks.append({"id": f"{source}-{i}", "text": piece, "source": source})
            i += 1
            buf = [sent]
        else:
            buf.append(sent)
    if buf:
        chunks.append({"id": f"{source}-{i}", "text": " ".join(buf), "source": source})
    return chunks

from langchain_text_splitters import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=64,
    length_function=token_len,
    separators=["\n\n", "\n", ". ", " ", ""],
)

docs = splitter.create_documents(
    texts=[open("policy.txt").read()],
    metadatas=[{"source": "policy.txt"}],
)
chunks = [{"id": f"policy-{i}", "text": d.page_content, **d.metadata}
          for i, d in enumerate(docs)]

import numpy as np
from openai import OpenAI

client = OpenAI()

def embed_batch(texts: list[str]) -> np.ndarray:
    resp = client.embeddings.create(model="text-embedding-3-small", input=texts)
    return np.array([d.embedding for d in resp.data])

def cosine_dist(a: np.ndarray, b: np.ndarray) -> float:
    return 1.0 - float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

def semantic_breakpoints(sents: list[str], percentile: float = 95) -> list[int]:
    vecs = embed_batch(sents)
    dists = [cosine_dist(vecs[i], vecs[i + 1]) for i in range(len(sents) - 1)]
    thresh = np.percentile(dists, percentile)
    return [i + 1 for i, d in enumerate(dists) if d >= thresh]

def chunk_semantic(text: str, source: str) -> list[dict]:
    sents = sentences_spacy(text)
    breaks = set(semantic_breakpoints(sents))
    groups, buf = [], []
    for i, s in enumerate(sents):
        buf.append(s)
        if (i + 1) in breaks:
            groups.append(" ".join(buf))
            buf = []
    if buf:
        groups.append(" ".join(buf))
    return [{"id": f"{source}-{i}", "text": g, "source": source}
            for i, g in enumerate(groups)]

from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
import ast

md_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")],
)
token_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=50, length_function=token_len)

def chunk_markdown(text: str, source: str) -> list[dict]:
    sections = md_splitter.split_text(text)
    out = []
    for i, doc in enumerate(sections):
        for j, sub in enumerate(token_splitter.split_documents([doc])):
            out.append({
                "id": f"{source}-{i}-{j}",
                "text": sub.page_content,
                "source": source,
                "section": sub.metadata.get("h2") or sub.metadata.get("h1", ""),
            })
    return out

def chunk_python_source(code: str, source: str) -> list[dict]:
    tree = ast.parse(code)
    chunks = []
    for i, node in enumerate(tree.body):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            segment = ast.get_source_segment(code, node)
            chunks.append({"id": f"{source}-{node.name}", "text": segment, "source": source, "symbol": node.name})
    return chunks

from pydantic import BaseModel
from openai import OpenAI

client = OpenAI()

class Segment(BaseModel):
    title: str
    start_char: int
    end_char: int

class ChunkPlan(BaseModel):
    segments: list[Segment]

def agentic_chunk(text: str, source: str) -> list[dict]:
    resp = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": f"Split this document into coherent retrieval chunks. Return char offsets.\n\n{text[:12000]}"
        }],
        response_format=ChunkPlan,
    )
    plan = resp.choices[0].message.parsed
    return [{
        "id": f"{source}-{i}",
        "text": text[s.start_char:s.end_char],
        "source": source,
        "section": s.title,
    } for i, s in enumerate(plan.segments)]

from datetime import datetime, timezone

def enrich_metadata(chunk: dict, doc_meta: dict) -> dict:
    return {
        "source": chunk["source"],
        "page": doc_meta.get("page"),
        "section": chunk.get("section", ""),
        "updated_at": doc_meta.get("updated_at", datetime.now(timezone.utc).isoformat()),
        "tenant_id": doc_meta["tenant_id"],
        "tokens": token_len(chunk["text"]),
        "chunk_policy": "recursive_v3_512",
    }

# Pinecone upsert
index.upsert(vectors=[{
    "id": c["id"],
    "values": embed(c["text"]),
    "metadata": enrich_metadata(c, doc_meta),
} for c in chunks])

# pgvector: metadata in JSONB column alongside embedding

def build_parent_child(sections: list[dict]) -> tuple[list[dict], list[dict]]:
    parents, children = [], []
    for sec in sections:
        pid = f"{sec['source']}#{sec['section']}"
        parents.append({"id": pid, "text": sec["text"], "source": sec["source"]})
        for i, child in enumerate(chunk_fixed(sec["text"], pid, size=256, overlap=32)):
            child["parent_id"] = pid
            children.append(child)
    return parents, children

def retrieve_with_parents(query: str, k: int = 5) -> list[str]:
    child_hits = vector_search(embed(query), k=k * 3)  # over-fetch
    parent_ids = list(dict.fromkeys(h["metadata"]["parent_id"] for h in child_hits))[:k]
    return [parent_store[pid]["text"] for pid in parent_ids]

def recall_at_k(gold_ids: set[str], retrieved_ids: list[str], k: int = 5) -> float:
    top = set(retrieved_ids[:k])
    return len(gold_ids & top) / len(gold_ids)

def eval_policy(chunks: list[dict], golden: list[dict], k: int = 5) -> float:
    scores = []
    for row in golden:
        hits = vector_search(embed(row["query"]), index=chunks, k=k)
        scores.append(recall_at_k(set(row["gold_chunk_ids"]), [h["id"] for h in hits], k))
    return sum(scores) / len(scores)

# Compare policies before prod promotion
for name, chunker in [("fixed512", chunk_fixed), ("sentence400", chunk_by_sentences)]:
    chunks = chunker(corpus_text, "golden.pdf")
    print(name, eval_policy(chunks, golden_queries))