Embedding in Practice

// chunker-3 builds on chunker-2 by adding two things:
// 1. A token budget: sections that exceed the target are split into smaller chunks
// 2. Overlap: the end of each chunk is repeated at the start of the next, so context
//    isn't lost at a boundary
import type { Chunk } from "../types/chunk";

// Record<string, any> is a TypeScript generic type meaning "an object whose keys are strings
// and whose values can be anything". We use it here because metadata fields vary by document.

// --- Chunking parameters ---
// targetTokens: the maximum number of tokens we want in a single chunk.
//   We use 400 rather than the 500 mentioned in the article as our hard ceiling.
//   The overlap (60 tokens) is added on top, so an emitted chunk can reach ~460 tokens
//   in practice. Keeping the content budget at 400 leaves headroom without wasting space.
//
// overlapTokens: how many tokens of the previous chunk we copy to the start of the next.
//   60 tokens (~240 characters) is roughly 2-3 sentences — enough to preserve context
//   at a boundary without significantly inflating chunk size.
//
// tokenLength: our character-to-token approximation (4 chars per token).
//   This is a rough rule of thumb for English text. In production you would use the
//   actual tokenizer for your embedding model (e.g. tiktoken for OpenAI models) to
//   get an exact count. For this example the approximation is close enough.
const def_targetTokens = 400;
const def_overlapTokens = 60;
const def_tokenLength = 4; // avg chars per token for English text

function estimateTokens(s: string): number {
  // In production, replace this with a real tokenizer for your embedding model.
  // tiktoken (OpenAI) and @anthropic-ai/tokenizer are both good options.
  return Math.ceil(s.length / def_tokenLength);
}

// Reads the document-level metadata out of the policy header block.
// Each field is marked up in bold in the markdown, e.g. **Policy ID:** LOYALTY-050
function extractDocMetadata(md: string) {

  // Match the top-level # heading to get the document title.
  // ?.[1] is optional chaining — it safely accesses index [1] of the match result,
  // returning undefined instead of throwing if match() found nothing.
  // ?? "" is nullish coalescing — it means "use this value, or "" if it's null/undefined".
  const docTitle = (md.match(/^#\s+(.+)$/m)?.[1] ?? "").trim();

  // A helper that extracts a single metadata field by its label.
  function pick(label: string): string {
    const re = new RegExp(`^\\*\\*${label}:\\*\\*\\s*(.+)\\s*$`, "mi");
    return (md.match(re)?.[1] ?? "").trim();
  }

  // A helper for fields that contain comma-separated lists, like Topics or Audience.
  // .map(s => s.trim()) strips whitespace from each item in the array.
  // .filter(Boolean) removes any empty strings that might result from trailing commas.
  function pickArray(label: string): string[] {
    const raw = pick(label);
    return raw ? raw.split(",").map(s => s.trim()).filter(Boolean) : [];
  }

  return {
    doc_title: docTitle,
    policy_id: pick("Policy ID"),
    version: pick("Version"),
    effective: pick("Effective"),
    region: pick("Region"),
    owner: pick("Owner"),
    classification: pick("Classification"),
    audience: pickArray("Audience"),
    sensitivity: pick("Sensitivity"),
    topics: pickArray("Topics"),
    applies_to: pickArray("Applies To"),
  };
}

// Splits a single section into one or more chunks, respecting the token budget.
// If the section fits within targetTokens it is returned as-is.
// If it is too large we walk through it block by block (splitting on blank lines),
// emitting a new chunk whenever the next block would push us over the limit.
// The overlap from the previous chunk is prepended to each new chunk so that
// a sentence or table row that straddles a boundary appears in both chunks.
// Note: overlap is calculated from the fresh content only — it does not compound
// across multiple splits, which would cause the chunks to grow unboundedly.
function* splitIntoTokenBudgetChunks(
  content: string,
  headingPath: string[],
  metadata: Chunk["metadata"],
  targetTokens: number,
  overlapTokens: number
): Generator<Chunk> {
  const tokens = estimateTokens(content);

  // Fast path: section is already within budget
  if (tokens <= targetTokens) {
    yield { chunkIndex: 0, headingPath, content, metadata };
    return;
  }

  // Split on blank lines rather than at a raw character position.
  // This keeps paragraphs, list items, and table rows intact — cutting
  // mid-sentence would produce incoherent chunks and hurt embedding quality.
  const blocks = content.split(/\n\n+/).filter(b => b.trim());

  let chunkCount = 0;
  let currentChunk: string[] = [];
  let currentTokens = 0;
  let overlapContent = "";

  for (const block of blocks) {
    const blockTokens = estimateTokens(block);

    // This block would push us over budget — emit what we have and start fresh
    if (currentTokens + blockTokens > targetTokens && currentChunk.length > 0) {
      const chunkContent = currentChunk.join("\n\n");
      yield {
        chunkIndex: chunkCount++,
        headingPath,
        content: overlapContent + chunkContent,
        metadata,
      };

      // Carry the tail of this chunk forward as overlap for the next one.
      // We take the last ~60 tokens worth of characters from the fresh content
      // (not including the overlap we prepended) so that overlap doesn't compound.
      const overlapChars = Math.floor(overlapTokens * def_tokenLength);
      overlapContent = chunkContent.slice(-overlapChars);
      if (overlapContent && !overlapContent.endsWith("\n\n")) {
        overlapContent += "\n\n";
      }

      // A table row without its column headers is meaningless — an LLM reading
      // "| $300 | $400 |" with no header has no idea what those values represent
      // and may hallucinate column names. Tables are self-contained units, so
      // if the overlap slice landed inside a table, discard it rather than
      // carry misleading context into the next chunk.
      if (overlapContent.trimStart().startsWith("|")) {
        overlapContent = "";
      }

      // Start the new chunk's token count from the overlap we're carrying forward
      currentChunk = [];
      currentTokens = estimateTokens(overlapContent);
    }

    currentChunk.push(block);
    currentTokens += blockTokens;
  }

  // Emit whatever is left in the final chunk
  if (currentChunk.length > 0) {
    yield {
      chunkIndex: chunkCount,
      headingPath,
      content: overlapContent + currentChunk.join("\n\n"),
      metadata,
    };
  }
}

// chunkPolicyMarkdown is our main export.
// It splits the document on H2 headings (##), extracts metadata from the header block,
// and then passes each section through splitIntoTokenBudgetChunks.
// Sections that are already under the token budget come through unchanged.
// Sections that are over budget get split into multiple chunks with overlap.
export function* chunkPolicyMarkdown(
  md: string,
  targetTokens = def_targetTokens,
  overlapTokens = def_overlapTokens
): Generator<Chunk> {
  const meta = extractDocMetadata(md);
  const docTitle = meta.doc_title || "Untitled";

  // Split on H2 headings and drop anything before the first one (the metadata header block).
  // That content has already been captured in `meta` above.
  const parts = md.split(/\n(?=##\s+)/g).filter(s => s.trim().startsWith("## "));

  let globalIndex = 0;

  for (let i = 0; i < parts.length; i++) {
    const section = parts[i].trimEnd() + "\n";
    const h2 = section.match(/^##\s+(.+)$/m)?.[1]?.trim() ?? `Section ${i}`;
    const headingPath = [docTitle, h2];

    // yield* delegates to the inner generator, forwarding each chunk as it is produced.
    // We re-index here to keep chunkIndex sequential across the whole document.
    for (const chunk of splitIntoTokenBudgetChunks(section, headingPath, meta, targetTokens, overlapTokens)) {
      yield { ...chunk, chunkIndex: globalIndex++ };
    }
  }
}