llmgw/internal/kb/chunk/chunk.go

// Package chunk splits extracted text into the pieces a knowledge base keeps
// and searches (BACKLOG F2).
//
// A chunk is a few paragraphs under the headings they sit beneath: about 500
// tokens, which an embedding model reads whole and a prompt can carry several
// of. Paragraphs are not cut while they fit; a section starts a new chunk
// when the current one is already a fair size, so a chunk rarely mixes two
// topics; and the last paragraph of a chunk is repeated at the start of the
// next, so an answer that spans the boundary is found whole in one of them.
package chunk

import (
	"strings"
	"unicode"
	"unicode/utf8"
)

// Options size the chunks. Zero values take the defaults.
type Options struct {
	// MaxChars is the most characters in a chunk's body: 1800, about 500
	// tokens of Turkish or English.
	MaxChars int
	// Overlap is the most characters of the previous chunk repeated at the
	// start of the next: 300.
	Overlap int
}

func (o Options) max() int {
	if o.MaxChars > 0 {
		return o.MaxChars
	}
	return 1800
}

func (o Options) overlap() int {
	if o.Overlap >= 0 && o.MaxChars > 0 {
		return min(o.Overlap, o.max()/2)
	}
	return 300
}

// Chunk is one piece: the headings above it and its text.
type Chunk struct {
	// Headings is the path of headings the chunk starts under, widest first.
	Headings []string
	Body     string
}

// Text is what is stored and embedded: the heading path, then the body. The
// path is part of it so that "how many days?" under "Annual leave" finds the
// paragraph that says "twenty days" without naming leave.
func (c Chunk) Text() string {
	if len(c.Headings) == 0 {
		return c.Body
	}
	return strings.Join(c.Headings, " › ") + "\n\n" + c.Body
}

// Split cuts Markdown-shaped text into chunks.
func Split(text string, o Options) []Chunk {
	max, overlap := o.max(), o.overlap()
	var out []Chunk
	type head struct {
		level int
		title string
	}
	var stack []head // the headings in force, widest first
	titles := func() []string {
		out := make([]string, len(stack))
		for i, h := range stack {
			out[i] = h.title
		}
		return out
	}
	var cur []string // blocks of the chunk being built
	var curHeads []string
	size := 0
	// carried is the overlap at the start of cur: it does not count as the
	// chunk having content of its own.
	carried := 0

	emit := func(keepOverlap bool) {
		if size-carried <= 0 {
			cur, size, carried = nil, 0, 0
			return
		}
		out = append(out, Chunk{Headings: curHeads, Body: strings.Join(cur, "\n\n")})
		last := cur[len(cur)-1]
		cur, size, carried = nil, 0, 0
		if keepOverlap && overlap > 0 && runes(last) <= overlap {
			cur, size, carried = []string{last}, runes(last), runes(last)
		}
	}
	add := func(b string) {
		if len(cur) == 0 || size == carried {
			curHeads = titles()
		}
		if size > 0 && size+2+runes(b) > max {
			emit(true)
			curHeads = titles()
		}
		cur = append(cur, b)
		size += runes(b) + 2
	}

	for _, b := range blocks(text) {
		if lvl, title := heading(b); lvl > 0 {
			// A new section starts a new chunk once this one is worth
			// keeping by itself; otherwise the heading goes into it.
			if size-carried >= max/3 {
				emit(false)
			}
			for len(stack) > 0 && stack[len(stack)-1].level >= lvl {
				stack = stack[:len(stack)-1]
			}
			stack = append(stack, head{lvl, title})
			if size-carried > 0 {
				add(b)
			} else {
				// Nothing of its own yet: the chunk starts under this heading.
				cur, size, carried = nil, 0, 0
				curHeads = titles()
			}
			continue
		}
		for _, piece := range fit(b, max) {
			add(piece)
		}
	}
	emit(false)
	return out
}

func runes(s string) int { return utf8.RuneCountInString(s) }

// heading reads a Markdown ATX heading line: its level and title.
func heading(b string) (int, string) {
	if strings.Contains(b, "\n") {
		return 0, ""
	}
	n := 0
	for n < len(b) && b[n] == '#' {
		n++
	}
	if n == 0 || n > 6 || n >= len(b) || b[n] != ' ' {
		return 0, ""
	}
	return n, strings.TrimSpace(strings.TrimRight(strings.TrimSpace(b[n:]), "#"))
}

// blocks splits text on blank lines, keeping fenced code blocks whole.
func blocks(text string) []string {
	var out []string
	var b []string
	fenced := false
	flush := func() {
		s := strings.TrimSpace(strings.Join(b, "\n"))
		if s != "" {
			out = append(out, s)
		}
		b = nil
	}
	for _, line := range strings.Split(text, "\n") {
		t := strings.TrimSpace(line)
		if strings.HasPrefix(t, "```") || strings.HasPrefix(t, "~~~") {
			fenced = !fenced
		}
		if !fenced && t == "" {
			flush()
			continue
		}
		// A heading line is a block of its own even without blank lines
		// around it.
		if !fenced && len(b) > 0 {
			if lvl, _ := heading(t); lvl > 0 {
				flush()
			}
		}
		b = append(b, strings.TrimRight(line, " \t"))
		if !fenced {
			if lvl, _ := heading(t); lvl > 0 {
				flush()
			}
		}
	}
	flush()
	return out
}

// fit cuts a block longer than max: at line ends, then sentence ends, then
// spaces, and only as a last resort mid-word.
func fit(b string, max int) []string {
	if runes(b) <= max {
		return []string{b}
	}
	for _, split := range []func(string) []string{lines, sentences, words} {
		parts := split(b)
		if len(parts) < 2 {
			continue
		}
		var out []string
		var cur strings.Builder
		for _, p := range parts {
			if cur.Len() > 0 && runes(cur.String())+runes(p) > max {
				out = append(out, strings.TrimSpace(cur.String()))
				cur.Reset()
			}
			if runes(p) > max {
				out = append(out, fit(strings.TrimSpace(p), max)...)
				continue
			}
			cur.WriteString(p)
		}
		if strings.TrimSpace(cur.String()) != "" {
			out = append(out, strings.TrimSpace(cur.String()))
		}
		return out
	}
	// One word longer than a chunk: a hash, a line of base64.
	r := []rune(b)
	var out []string
	for len(r) > max {
		out, r = append(out, string(r[:max])), r[max:]
	}
	return append(out, string(r))
}

// lines splits keeping each line's newline.
func lines(s string) []string { return strings.SplitAfter(s, "\n") }

// sentences splits after ., ? or ! followed by a space, keeping both.
func sentences(s string) []string {
	var out []string
	start := 0
	r := []rune(s)
	for i := 0; i < len(r)-1; i++ {
		if (r[i] == '.' || r[i] == '?' || r[i] == '!' || r[i] == '…') && unicode.IsSpace(r[i+1]) {
			out = append(out, string(r[start:i+2]))
			start = i + 2
		}
	}
	if start < len(r) {
		out = append(out, string(r[start:]))
	}
	return out
}

func words(s string) []string { return strings.SplitAfter(s, " ") }