SPB Git
3.1 KB · 35 lines typescript
1import { createHash } from "node:crypto";2import { docBlocks, type PMNode } from "@/lib/editor/text";3export interface Chunk { ordinal: number; heading: string | null; blockId: string | null; content: string; tokenCount: number; contentHash: string }4export const approxTokens = (s: string) => Math.ceil(s.length / 3.6);5const sha = (s: string) => createHash("sha256").update(s).digest("hex").slice(0, 32);67/** Document-aware chunking: follows headings/blocks, targets ~350 tokens, 60-token overlap for long runs. */8export function chunkDocument(title: string, doc: PMNode | null | undefined, opts = { target: 350, max: 500, overlap: 60 }): Chunk[] {9  const blocks = docBlocks(doc); const chunks: Chunk[] = [];10  let buf: string[] = []; let bufTok = 0; let heading: string | null = null; let blockId: string | null = null;11  const flush = () => { if (!buf.length) return; const content = buf.join("\n").trim(); if (content) chunks.push({ ordinal: chunks.length, heading, blockId, content, tokenCount: approxTokens(content), contentHash: sha(content) }); buf = []; bufTok = 0; };12  for (const b of blocks) {13    if (b.type === "heading") { flush(); heading = b.text; blockId = b.id; continue; }14    const t = approxTokens(b.text);15    if (t > opts.max) { flush(); chunkLongText(b.text, opts).forEach((piece) => { chunks.push({ ordinal: chunks.length, heading, blockId: b.id, content: piece, tokenCount: approxTokens(piece), contentHash: sha(piece) }); }); continue; }16    if (bufTok + t > opts.target && buf.length) flush();17    if (!buf.length) blockId = b.id;18    buf.push(b.text); bufTok += t;19  }20  flush();21  if (!chunks.length && title) chunks.push({ ordinal: 0, heading: null, blockId: null, content: title, tokenCount: approxTokens(title), contentHash: sha(title) });22  return chunks;23}24export function chunkLongText(text: string, opts = { target: 350, max: 500, overlap: 60 }): string[] {25  const sentences = text.split(/(?<=[.!?。])\s+|\n+/).filter(Boolean); const out: string[] = []; let cur: string[] = []; let tok = 0;26  for (const s of sentences) { const t = approxTokens(s); if (tok + t > opts.target && cur.length) { out.push(cur.join(" ")); const keep: string[] = []; let kt = 0; for (let i = cur.length - 1; i >= 0 && kt < opts.overlap; i--) { keep.unshift(cur[i]); kt += approxTokens(cur[i]); } cur = keep; tok = kt; } cur.push(s); tok += t; }27  if (cur.length) out.push(cur.join(" ")); return out;28}29export function chunkPlainText(text: string): Chunk[] {30  const paras = text.split(/\n{2,}/).map((p) => p.trim()).filter(Boolean); const chunks: Chunk[] = []; let buf: string[] = []; let tok = 0;31  const flush = () => { if (!buf.length) return; const c = buf.join("\n\n"); chunks.push({ ordinal: chunks.length, heading: null, blockId: null, content: c, tokenCount: approxTokens(c), contentHash: sha(c) }); buf = []; tok = 0; };32  for (const p of paras) { const t = approxTokens(p); if (t > 500) { flush(); chunkLongText(p).forEach((x) => chunks.push({ ordinal: chunks.length, heading: null, blockId: null, content: x, tokenCount: approxTokens(x), contentHash: sha(x) })); continue; } if (tok + t > 350 && buf.length) flush(); buf.push(p); tok += t; }33  flush(); return chunks;34}35
Files · main
Menu

Theme