libredesk/internal/stringutil/htmlchunker.go

package stringutil

import (
	"fmt"
	"regexp"
	"slices"
	"strings"
	"unicode/utf8"

	"golang.org/x/net/html"
)

// Bump whenever ChunkHTMLContent changes the text it emits for unchanged input; it feeds reindex fingerprints.
const ChunkerVersion = 3

var (
	sentenceRegex     = regexp.MustCompile(`[.!?]+[\s]+`)
	headingInnerRegex = regexp.MustCompile(`(?is)<h[1-6][^>]*>(.*?)</h[1-6]>`)

	// nonContentElements hold markup, not prose. HTML2Text only skips head, script and style, so
	// svg, noscript and template have to be excluded here or their text reaches the index.
	nonContentElements = map[string]bool{
		"head": true, "script": true, "style": true, "noscript": true, "template": true, "svg": true,
	}

	blockElements = map[string]bool{
		"h1": true, "h2": true, "h3": true, "h4": true, "h5": true, "h6": true,
		"p": true, "div": true, "section": true, "article": true, "aside": true,
		"header": true, "footer": true, "main": true, "nav": true,
		"ul": true, "ol": true, "li": true, "dl": true, "dt": true, "dd": true,
		"table": true, "thead": true, "tbody": true, "tfoot": true, "tr": true, "td": true, "th": true,
		"form": true, "fieldset": true, "legend": true,
		"blockquote": true, "pre": true, "code": true, "figure": true, "figcaption": true,
		"address": true, "details": true, "summary": true, "hr": true,
	}
)

type ChunkConfig struct {
	MaxTokens      int
	MinTokens      int
	OverlapTokens  int
	TokenizerFunc  func(string) int
	PreserveBlocks []string
}

type htmlBoundary struct {
	Type     string
	Content  string
	Priority int
	Tokens   int
}

// DefaultChunkConfig returns a ChunkConfig with sensible defaults for HTML chunking.
func DefaultChunkConfig() ChunkConfig {
	return ChunkConfig{
		MaxTokens:      2000,
		MinTokens:      400,
		OverlapTokens:  75,
		TokenizerFunc:  defaultTokenizer,
		PreserveBlocks: []string{"pre", "code", "table"},
	}
}

func (c *ChunkConfig) validate() error {
	if c.MaxTokens <= c.MinTokens {
		return fmt.Errorf("MaxTokens must be greater than MinTokens")
	}
	if c.OverlapTokens >= c.MinTokens {
		return fmt.Errorf("OverlapTokens must be less than MinTokens")
	}
	if c.TokenizerFunc == nil {
		c.TokenizerFunc = defaultTokenizer
	}
	return nil
}

// ChunkHTMLContent splits HTML into structure-aware chunks for embedding, prepending the title and section heading to each chunk's text.
func ChunkHTMLContent(title, htmlContent string, config ...ChunkConfig) ([]string, error) {
	cfg := DefaultChunkConfig()
	if len(config) > 0 {
		cfg = config[0]
	}

	if err := cfg.validate(); err != nil {
		return nil, err
	}

	if strings.TrimSpace(htmlContent) == "" {
		return []string{buildEmbeddingText(title, "", "")}, nil
	}

	htmlContent = prepareHTMLForEmbedding(htmlContent)

	boundaries, err := parseHTMLBoundaries(htmlContent, cfg)
	if err != nil {
		return nil, fmt.Errorf("failed to parse HTML: %w", err)
	}

	chunks := createChunks(boundaries, cfg)
	result := make([]string, len(chunks))

	var lastHeading string
	for i, chunk := range chunks {
		if h := extractLeadingHeading(chunk.Content); h != "" {
			lastHeading = h
		}
		result[i] = buildEmbeddingText(title, lastHeading, HTML2Text(chunk.Content))
	}

	// Markup with no text at all still needs a chunk; embedding nothing is indistinguishable from success.
	if len(result) == 0 {
		return []string{buildEmbeddingText(title, "", "")}, nil
	}

	return result, nil
}

// defaultTokenizer estimates token count with a conservative rune-based ratio.
func defaultTokenizer(text string) int {
	return int(float64(utf8.RuneCountInString(text)) * 0.4)
}

func isBlockElement(tag string) bool {
	return blockElements[tag]
}

func parseHTMLBoundaries(htmlContent string, cfg ChunkConfig) ([]htmlBoundary, error) {
	doc, err := html.Parse(strings.NewReader(htmlContent))
	if err != nil {
		return nil, err
	}

	var boundaries []htmlBoundary

	var extract func(*html.Node)
	extract = func(n *html.Node) {
		switch n.Type {
		case html.TextNode:
			// Text outside any block element has no boundary of its own and would drop out of the index.
			text := strings.TrimSpace(n.Data)
			if text == "" {
				return
			}
			// Unwrapped, so mergeBoundaries reassembles a run broken up by inline tags into one sentence.
			boundaries = append(boundaries, htmlBoundary{
				Type:     "p",
				Content:  html.EscapeString(text) + " ",
				Priority: getPriority("p"),
				Tokens:   cfg.TokenizerFunc(text),
			})
			return
		case html.ElementNode:
			tag := strings.ToLower(n.Data)
			if nonContentElements[tag] {
				return
			}

			// Rendering a non-block wrapper would redo its whole subtree once per nesting level.
			if isBlockElement(tag) {
				var content strings.Builder
				_ = html.Render(&content, n)
				contentStr := content.String()

				cleanText := HTML2Text(contentStr)
				if strings.TrimSpace(cleanText) == "" {
					return
				}

				tokens := cfg.TokenizerFunc(cleanText)
				// An oversized container taken as one atomic boundary would be truncated at
				// MaxTokens, silently dropping the rest; descend into its block children instead.
				if tokens <= cfg.MaxTokens || isPreservedBlock(tag, cfg.PreserveBlocks) || !splittableIntoBlocks(n) {
					boundaries = append(boundaries, htmlBoundary{
						Type:     tag,
						Content:  contentStr,
						Priority: getPriority(tag),
						Tokens:   tokens,
					})
					return
				}
			}
		}

		for c := n.FirstChild; c != nil; c = c.NextSibling {
			extract(c)
		}
	}

	extract(doc)

	return mergeBoundaries(boundaries, cfg), nil
}

// getPriority ranks tags for chunking; lower is higher priority.
func getPriority(tag string) int {
	switch tag {
	case "h1", "h2":
		return 1
	case "h3", "h4", "h5", "h6", "pre", "code":
		return 2
	case "p", "table", "ul", "ol", "blockquote":
		return 3
	case "div", "section", "article", "figure":
		return 4
	default:
		return 5
	}
}

func isPreservedBlock(blockType string, preserveBlocks []string) bool {
	return slices.Contains(preserveBlocks, blockType)
}

// splittableIntoBlocks reports whether all of a node's visible content lives inside block-element
// children, so splitting the node into its children loses nothing.
func splittableIntoBlocks(n *html.Node) bool {
	hasBlock := false
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		switch c.Type {
		case html.TextNode:
			if strings.TrimSpace(c.Data) != "" {
				return false
			}
		case html.ElementNode:
			if isBlockElement(strings.ToLower(c.Data)) {
				hasBlock = true
				continue
			}
			var buf strings.Builder
			_ = html.Render(&buf, c)
			if strings.TrimSpace(HTML2Text(buf.String())) != "" {
				return false
			}
		}
	}
	return hasBlock
}

func mergeBoundaries(boundaries []htmlBoundary, cfg ChunkConfig) []htmlBoundary {
	if len(boundaries) == 0 {
		return boundaries
	}

	var merged []htmlBoundary
	current := boundaries[0]

	for i := 1; i < len(boundaries); i++ {
		next := boundaries[i]

		if next.Priority == 1 {
			merged = append(merged, current)
			current = next
			continue
		}

		if current.Priority == 1 && current.Tokens >= cfg.MinTokens {
			merged = append(merged, current)
			current = next
			continue
		}

		if isPreservedBlock(current.Type, cfg.PreserveBlocks) || isPreservedBlock(next.Type, cfg.PreserveBlocks) {
			merged = append(merged, current)
			current = next
			continue
		}

		combinedTokens := current.Tokens + next.Tokens
		shouldMerge := false

		if combinedTokens < cfg.MinTokens {
			shouldMerge = true
		} else if current.Priority >= 3 && next.Priority >= 3 && combinedTokens < cfg.MaxTokens {
			shouldMerge = true
		}

		if shouldMerge {
			current.Content += next.Content
			current.Tokens = combinedTokens
			current.Priority = min(current.Priority, next.Priority)
		} else {
			merged = append(merged, current)
			current = next
		}
	}

	merged = append(merged, current)
	return merged
}

// splitOversizedBoundary breaks an atomic boundary into MaxTokens-sized pieces on sentence boundaries.
func splitOversizedBoundary(boundary htmlBoundary, cfg ChunkConfig) []htmlBoundary {
	var out []htmlBoundary
	// Only Content is read downstream; setting Tokens would be a second tokenizer pass over the whole block.
	for _, piece := range chunkPlainText(HTML2Text(boundary.Content), boundary.Tokens, cfg) {
		piece = strings.TrimSpace(piece)
		if piece == "" {
			continue
		}
		out = append(out, htmlBoundary{Content: "<p>" + html.EscapeString(piece) + "</p>"})
	}
	return out
}

func createChunks(boundaries []htmlBoundary, cfg ChunkConfig) []htmlBoundary {
	if len(boundaries) == 0 {
		return boundaries
	}

	var chunks []htmlBoundary
	var currentChunk htmlBoundary
	currentChunk.Priority = 10

	for _, boundary := range boundaries {
		// Split here, not at the boundary source, so an overlap prefix can't push a piece back over the limit.
		if boundary.Tokens > cfg.MaxTokens {
			if currentChunk.Content != "" {
				chunks = append(chunks, currentChunk)
				currentChunk = htmlBoundary{Priority: 10}
			}
			chunks = append(chunks, splitOversizedBoundary(boundary, cfg)...)
			continue
		}

		shouldStartNewChunk := false

		if boundary.Priority == 1 && currentChunk.Tokens >= cfg.MinTokens {
			shouldStartNewChunk = true
		}

		if currentChunk.Tokens+boundary.Tokens > cfg.MaxTokens {
			if currentChunk.Content != "" {
				shouldStartNewChunk = true
			}
		}

		if shouldStartNewChunk && currentChunk.Content != "" {
			chunks = append(chunks, currentChunk)

			var overlapContent string
			if !isPreservedBlock(boundary.Type, cfg.PreserveBlocks) {
				overlapContent = extractOverlap(currentChunk.Content, cfg)
			}

			currentChunk = htmlBoundary{
				Content:  overlapContent,
				Tokens:   cfg.TokenizerFunc(HTML2Text(overlapContent)),
				Priority: 10,
			}
		}

		currentChunk.Content += boundary.Content
		currentChunk.Tokens += boundary.Tokens

		if boundary.Priority < currentChunk.Priority {
			currentChunk.Priority = boundary.Priority
		}
	}

	if currentChunk.Content != "" {
		chunks = append(chunks, currentChunk)
	}

	return chunks
}

// extractOverlap carries trailing whole sentences into the next chunk for context continuity.
func extractOverlap(content string, cfg ChunkConfig) string {
	cleanText := HTML2Text(content)
	sentences := splitSentences(cleanText)

	if len(sentences) <= 1 {
		return ""
	}

	var overlap []string
	tokens := 0
	for i := len(sentences) - 1; i >= 0 && tokens < cfg.OverlapTokens; i-- {
		sentence := strings.TrimSpace(sentences[i])
		if sentence == "" {
			continue
		}
		sentTokens := cfg.TokenizerFunc(sentence)
		if tokens+sentTokens <= cfg.OverlapTokens {
			overlap = append([]string{sentence}, overlap...)
			tokens += sentTokens
		} else {
			break
		}
	}

	if len(overlap) == 0 {
		return ""
	}

	return "<p>" + html.EscapeString(strings.Join(overlap, " ")) + "</p>\n"
}

// extractLeadingHeading returns the plain text of the first heading in the chunk HTML, if any.
func extractLeadingHeading(htmlContent string) string {
	m := headingInnerRegex.FindStringSubmatch(htmlContent)
	if len(m) < 2 {
		return ""
	}
	return strings.TrimSpace(HTML2Text(m[1]))
}

// buildEmbeddingText prepends the title and section heading (when present) to the chunk text.
func buildEmbeddingText(title, heading, cleanText string) string {
	title = strings.TrimSpace(title)
	heading = strings.TrimSpace(heading)
	cleanText = strings.TrimSpace(cleanText)

	if title == "" && heading == "" {
		return cleanText
	}
	if cleanText == "" {
		return title
	}

	var b strings.Builder
	if title != "" {
		fmt.Fprintf(&b, "Title: %s\n", title)
	}
	if heading != "" && !strings.EqualFold(heading, title) {
		fmt.Fprintf(&b, "Section: %s\n", heading)
	}
	fmt.Fprintf(&b, "Content: %s", cleanText)
	return b.String()
}

// chunkPlainText packs text into <=MaxTokens pieces on sentence boundaries, hard-splitting any single sentence over the cap; tokens is the caller's already-measured count for text.
func chunkPlainText(text string, tokens int, cfg ChunkConfig) []string {
	if tokens <= cfg.MaxTokens {
		return []string{text}
	}

	var chunks []string
	var cur strings.Builder
	curTokens := 0
	flush := func() {
		if cur.Len() > 0 {
			chunks = append(chunks, strings.TrimSpace(cur.String()))
			cur.Reset()
			curTokens = 0
		}
	}

	for _, s := range splitSentences(text) {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}
		st := cfg.TokenizerFunc(s)
		if st > cfg.MaxTokens {
			flush()
			chunks = append(chunks, hardSplit(s, cfg)...)
			continue
		}
		if curTokens+st > cfg.MaxTokens {
			flush()
		}
		if cur.Len() > 0 {
			cur.WriteString(" ")
		}
		cur.WriteString(s)
		curTokens += st
	}
	flush()
	return chunks
}

// hardSplit cuts s into consecutive pieces each within MaxTokens, taking the largest fitting rune prefix per piece.
func hardSplit(s string, cfg ChunkConfig) []string {
	runes := []rune(s)
	var out []string
	for len(runes) > 0 {
		// Probing upward first keeps tokenizer calls near the piece size instead of the whole remaining string.
		hi := 2
		for hi < len(runes) && cfg.TokenizerFunc(string(runes[:hi])) <= cfg.MaxTokens {
			hi *= 2
		}
		hi = min(hi, len(runes))
		lo, best := 1, 1
		for lo <= hi {
			mid := (lo + hi) / 2
			if cfg.TokenizerFunc(string(runes[:mid])) <= cfg.MaxTokens {
				best = mid
				lo = mid + 1
			} else {
				hi = mid - 1
			}
		}
		out = append(out, string(runes[:best]))
		runes = runes[best:]
	}
	return out
}

// splitSentences splits on sentence boundaries, keeping each sentence's terminating punctuation.
func splitSentences(text string) []string {
	seps := sentenceRegex.FindAllStringIndex(text, -1)
	if len(seps) == 0 {
		return []string{text}
	}
	out := make([]string, 0, len(seps)+1)
	prev := 0
	for _, sep := range seps {
		out = append(out, text[prev:sep[0]]+strings.TrimSpace(text[sep[0]:sep[1]]))
		prev = sep[1]
	}
	return append(out, text[prev:])
}