Code sync, security fixes

2026-08-11 10:40:10 +01:00
parent 4e00deedf8
commit 53dbc28383
141 changed files with 36194 additions and 0 deletions
@@ -0,0 +1,390 @@
+// Package chunking splits a parsed document into retrievable chunks.
+//
+// Chunking follows document structure rather than counting characters blindly:
+// headings, paragraphs, list groups, tables and JSON object boundaries are
+// preferred split points, and a character ceiling is applied only when a
+// structural block is itself too large.
+package chunking
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"strings"
+	"unicode/utf8"
+
+	"github.com/andydixon/relata/internal/parsing"
+)
+
+// Options configure the chunker. Zero values are replaced by defaults.
+type Options struct {
+	// MaxChars is the hard ceiling on one chunk's core span.
+	MaxChars int
+	// TargetChars is the size the chunker aims for when merging small blocks.
+	TargetChars int
+	// MinChars is the size below which a chunk is merged into its neighbour
+	// rather than emitted on its own.
+	MinChars int
+	// OverlapChars extends each chunk backwards into its predecessor, so a
+	// sentence spanning a boundary is still retrievable from both sides.
+	OverlapChars int
+	// MaxChunks bounds output; excess content is reported via Truncated.
+	MaxChunks int
+	// MinSplitChars is the document size below which the whole document becomes
+	// a single chunk. Splitting a short note only dilutes its embedding.
+	MinSplitChars int
+}
+
+func (o Options) withDefaults() Options {
+	if o.MaxChars <= 0 {
+		o.MaxChars = 1800
+	}
+	if o.TargetChars <= 0 || o.TargetChars > o.MaxChars {
+		o.TargetChars = o.MaxChars * 2 / 3
+	}
+	if o.MinChars <= 0 {
+		o.MinChars = o.MaxChars / 9
+	}
+	if o.OverlapChars < 0 || o.OverlapChars >= o.MaxChars {
+		o.OverlapChars = 0
+	}
+	if o.MaxChunks <= 0 {
+		o.MaxChunks = 400
+	}
+	if o.MinSplitChars <= 0 {
+		o.MinSplitChars = o.MaxChars / 2
+	}
+	return o
+}
+
+// Chunk is one produced chunk. Start and End are byte offsets into the source
+// text and always satisfy Text == text[Start:End], including any overlap.
+type Chunk struct {
+	Sequence int
+	Text     string
+	Start    int
+	End      int
+	Hash     string
+	Page     int
+	Section  string
+	Heading  string
+	// OverlapChars is how many leading bytes were carried over from the
+	// preceding chunk, so callers can tell borrowed context from new content.
+	OverlapChars int
+}
+
+// Result is the chunker's output.
+type Result struct {
+	Chunks []Chunk
+	// Truncated is true when MaxChunks stopped the chunker before the end of
+	// the document.
+	Truncated bool
+}
+
+// Split divides doc into chunks.
+//
+// The returned chunks cover the document in order. Very short documents are
+// returned as a single chunk: splitting a two-sentence note produces fragments
+// that embed worse than the whole.
+func Split(doc *parsing.Document, opts Options) Result {
+	o := opts.withDefaults()
+	text := doc.Text
+
+	if strings.TrimSpace(text) == "" {
+		return Result{}
+	}
+
+	// A short document is one chunk, whatever its internal structure.
+	if len(text) <= o.MinSplitChars {
+		return Result{Chunks: []Chunk{finish(0, text, 0, len(text), 0, firstPage(doc), firstSection(doc), firstHeading(doc))}}
+	}
+
+	blocks := doc.Blocks
+	if len(blocks) == 0 {
+		blocks = []parsing.Block{{Kind: parsing.BlockParagraph, Text: text, Start: 0, End: len(text)}}
+	}
+
+	// Phase 1: group blocks into core spans that respect structure.
+	spans := groupBlocks(blocks, text, o)
+
+	// Phase 2: apply overlap and materialise chunks.
+	res := Result{}
+	for i, sp := range spans {
+		if len(res.Chunks) >= o.MaxChunks {
+			res.Truncated = true
+			break
+		}
+		start := sp.start
+		overlap := 0
+		if o.OverlapChars > 0 && i > 0 {
+			lowerBound := spans[i-1].start
+			cand := start - o.OverlapChars
+			if cand < lowerBound {
+				cand = lowerBound
+			}
+			cand = alignRuneStart(text, cand)
+			// Prefer to begin the overlap at a word boundary so the borrowed
+			// context does not start mid-word.
+			if b := lastBoundary(text, cand, start); b > cand {
+				cand = b
+			}
+			if cand < start {
+				overlap = start - cand
+				start = cand
+			}
+		}
+		res.Chunks = append(res.Chunks, finish(
+			len(res.Chunks), text[start:sp.end], start, sp.end, overlap, sp.page, sp.section, sp.heading))
+	}
+	return res
+}
+
+type coreSpan struct {
+	start, end int
+	page       int
+	section    string
+	heading    string
+}
+
+// groupBlocks merges consecutive blocks into spans no larger than MaxChars,
+// starting a new span at structural boundaries.
+func groupBlocks(blocks []parsing.Block, text string, o Options) []coreSpan {
+	var spans []coreSpan
+	var cur *coreSpan
+	var curHeading string
+
+	flush := func() {
+		if cur != nil {
+			spans = append(spans, *cur)
+			cur = nil
+		}
+	}
+
+	for _, b := range blocks {
+		if strings.TrimSpace(b.Text) == "" {
+			continue
+		}
+		if b.Kind == parsing.BlockHeading {
+			curHeading = strings.TrimLeft(strings.TrimSpace(b.Text), "# ")
+		}
+
+		size := b.End - b.Start
+
+		// An oversized single block must be split internally.
+		if size > o.MaxChars {
+			flush()
+			for _, sub := range splitLarge(text, b.Start, b.End, o) {
+				spans = append(spans, coreSpan{
+					start: sub.start, end: sub.end,
+					page: b.Page, section: b.Section, heading: curHeading,
+				})
+			}
+			continue
+		}
+
+		switch {
+		case cur == nil:
+			cur = &coreSpan{start: b.Start, end: b.End, page: b.Page, section: b.Section, heading: curHeading}
+
+		// A heading opens a new section: start a fresh chunk once the current
+		// one carries enough content to stand alone.
+		case b.Kind == parsing.BlockHeading && cur.end-cur.start >= o.MinChars:
+			flush()
+			cur = &coreSpan{start: b.Start, end: b.End, page: b.Page, section: b.Section, heading: curHeading}
+
+		// A page change in a paginated document is a real boundary.
+		case b.Page > 0 && cur.page > 0 && b.Page != cur.page && cur.end-cur.start >= o.MinChars:
+			flush()
+			cur = &coreSpan{start: b.Start, end: b.End, page: b.Page, section: b.Section, heading: curHeading}
+
+		// Adding this block would exceed the ceiling.
+		case b.End-cur.start > o.MaxChars:
+			flush()
+			cur = &coreSpan{start: b.Start, end: b.End, page: b.Page, section: b.Section, heading: curHeading}
+
+		// A JSON object boundary is worth respecting once the chunk is big
+		// enough, so an object's fields stay together.
+		case b.Kind == parsing.BlockJSONField && b.Section != cur.section &&
+			cur.end-cur.start >= o.TargetChars:
+			flush()
+			cur = &coreSpan{start: b.Start, end: b.End, page: b.Page, section: b.Section, heading: curHeading}
+
+		default:
+			cur.end = b.End
+			if cur.section == "" {
+				cur.section = b.Section
+			}
+			if cur.page == 0 {
+				cur.page = b.Page
+			}
+			// Reaching the target size is a good place to stop, provided the
+			// next block is not mid-thought (handled by the loop restarting).
+			if cur.end-cur.start >= o.TargetChars {
+				flush()
+			}
+		}
+	}
+	flush()
+
+	return mergeTiny(spans, o)
+}
+
+// mergeTiny folds spans below MinChars into an adjacent span, so the output
+// does not contain fragments too small to embed usefully.
+func mergeTiny(spans []coreSpan, o Options) []coreSpan {
+	if len(spans) < 2 {
+		return spans
+	}
+	out := make([]coreSpan, 0, len(spans))
+	for _, sp := range spans {
+		if len(out) > 0 {
+			prev := &out[len(out)-1]
+			tiny := sp.end-sp.start < o.MinChars
+			fits := sp.end-prev.start <= o.MaxChars
+			contiguous := sp.start >= prev.end
+			if tiny && fits && contiguous {
+				prev.end = sp.end
+				if prev.section == "" {
+					prev.section = sp.section
+				}
+				continue
+			}
+		}
+		out = append(out, sp)
+	}
+	return out
+}
+
+// splitLarge divides one oversized block at the best available boundary:
+// blank lines, then line breaks, then sentence ends, then a hard character cut.
+func splitLarge(text string, start, end int, o Options) []span {
+	var out []span
+	for start < end {
+		limit := start + o.MaxChars
+		if limit >= end {
+			out = append(out, span{start, end})
+			break
+		}
+		limit = alignRuneStart(text, limit)
+		// Look backwards from the ceiling for a boundary, but not so far back
+		// that chunks become tiny.
+		floor := start + o.MinChars
+		if floor > limit {
+			floor = start + (limit-start)/2
+		}
+		cut := findBoundary(text, floor, limit)
+		if cut <= start {
+			cut = limit
+		}
+		out = append(out, span{start, cut})
+		start = cut
+	}
+	return out
+}
+
+type span struct{ start, end int }
+
+// findBoundary searches [floor,limit) for the latest good split point.
+func findBoundary(text string, floor, limit int) int {
+	if floor < 0 {
+		floor = 0
+	}
+	if limit > len(text) {
+		limit = len(text)
+	}
+	// Blank line.
+	for i := limit - 1; i > floor; i-- {
+		if text[i] == '\n' && i > 0 && text[i-1] == '\n' {
+			return i + 1
+		}
+	}
+	// Sentence terminator followed by whitespace.
+	for i := limit - 1; i > floor; i-- {
+		if (text[i] == '.' || text[i] == '!' || text[i] == '?' || text[i] == ';') &&
+			i+1 < len(text) && isSpace(text[i+1]) {
+			return i + 1
+		}
+	}
+	// Line break.
+	for i := limit - 1; i > floor; i-- {
+		if text[i] == '\n' {
+			return i + 1
+		}
+	}
+	// Word break.
+	for i := limit - 1; i > floor; i-- {
+		if isSpace(text[i]) {
+			return i + 1
+		}
+	}
+	return limit
+}
+
+// lastBoundary finds a word start within [from,upto) to begin overlap text at.
+func lastBoundary(text string, from, upto int) int {
+	for i := from; i < upto && i < len(text); i++ {
+		if isSpace(text[i]) {
+			return i + 1
+		}
+	}
+	return from
+}
+
+func isSpace(b byte) bool { return b == ' ' || b == '\n' || b == '\t' || b == '\r' }
+
+// alignRuneStart moves i back to the nearest UTF-8 rune boundary.
+func alignRuneStart(text string, i int) int {
+	if i <= 0 {
+		return 0
+	}
+	if i >= len(text) {
+		return len(text)
+	}
+	for i > 0 && !utf8.RuneStart(text[i]) {
+		i--
+	}
+	return i
+}
+
+func finish(seq int, body string, start, end, overlap, page int, section, heading string) Chunk {
+	return Chunk{
+		Sequence:     seq,
+		Text:         body,
+		Start:        start,
+		End:          end,
+		Hash:         HashText(body),
+		Page:         page,
+		Section:      section,
+		Heading:      heading,
+		OverlapChars: overlap,
+	}
+}
+
+// HashText returns the hex SHA-256 of s. Chunk hashes make re-chunking
+// idempotent: an unchanged chunk keeps its identity across reprocessing.
+func HashText(s string) string {
+	sum := sha256.Sum256([]byte(s))
+	return hex.EncodeToString(sum[:])
+}
+
+func firstPage(doc *parsing.Document) int {
+	if len(doc.Blocks) > 0 {
+		return doc.Blocks[0].Page
+	}
+	return 0
+}
+
+func firstSection(doc *parsing.Document) string {
+	if len(doc.Blocks) > 0 {
+		return doc.Blocks[0].Section
+	}
+	return ""
+}
+
+func firstHeading(doc *parsing.Document) string {
+	for _, b := range doc.Blocks {
+		if b.Kind == parsing.BlockHeading {
+			return strings.TrimLeft(strings.TrimSpace(b.Text), "# ")
+		}
+	}
+	return ""
+}