spb/note-spboucher-ai
Public
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