core/server/knowledge/chunker.go

package knowledge

import (
	"strings"
	"unicode/utf8"
)

// ── Configuration ────────────────────────────

// ChunkConfig controls how text is split into chunks.
type ChunkConfig struct {
	ChunkSize    int      // target chars per chunk (default 1000)
	ChunkOverlap int      // overlap chars between adjacent chunks (default 200)
	Separators   []string // split hierarchy, tried in order
}

// DefaultChunkConfig returns sensible defaults for general documents.
// ~250 tokens per chunk at ~4 chars/token.
func DefaultChunkConfig() ChunkConfig {
	return ChunkConfig{
		ChunkSize:    1000,
		ChunkOverlap: 200,
		Separators:   []string{"\n\n", "\n", ". ", " "},
	}
}

// ── Output ───────────────────────────────────

// Chunk is a single text segment with position metadata.
type Chunk struct {
	Content    string         // chunk text
	Index      int            // ordinal within document (0-based)
	TokenCount int            // estimated token count (~chars/4)
	Metadata   map[string]any // extensible: page, heading, etc.
}

// ── Splitter ─────────────────────────────────

// SplitText splits text into overlapping chunks using recursive
// character splitting. It tries separators in order, splitting on
// the first one that produces segments smaller than ChunkSize.
// Segments that are still too large are recursively split with the
// next separator.
func SplitText(text string, cfg ChunkConfig) []Chunk {
	if cfg.ChunkSize <= 0 {
		cfg.ChunkSize = 1000
	}
	if cfg.ChunkOverlap < 0 {
		cfg.ChunkOverlap = 0
	}
	if cfg.ChunkOverlap >= cfg.ChunkSize {
		cfg.ChunkOverlap = cfg.ChunkSize / 5
	}
	if len(cfg.Separators) == 0 {
		cfg.Separators = DefaultChunkConfig().Separators
	}

	text = strings.TrimSpace(text)
	if text == "" {
		return nil
	}

	// If text fits in one chunk, return it directly.
	if utf8.RuneCountInString(text) <= cfg.ChunkSize {
		return []Chunk{{
			Content:    text,
			Index:      0,
			TokenCount: estimateTokens(text),
			Metadata:   map[string]any{},
		}}
	}

	// Recursively split, then merge into overlapping windows.
	segments := recursiveSplit(text, cfg.Separators, cfg.ChunkSize)
	return mergeSegments(segments, cfg.ChunkSize, cfg.ChunkOverlap)
}

// recursiveSplit tries each separator in order. For the first separator
// that actually splits the text, it checks if each piece fits within
// chunkSize. Pieces that are still too large recurse with the remaining
// separators. Pieces that fit are kept as-is.
func recursiveSplit(text string, separators []string, chunkSize int) []string {
	if utf8.RuneCountInString(text) <= chunkSize {
		return []string{text}
	}

	if len(separators) == 0 {
		// No separators left — hard split by character count.
		return hardSplit(text, chunkSize)
	}

	sep := separators[0]
	rest := separators[1:]

	parts := splitKeepSep(text, sep)
	if len(parts) <= 1 {
		// This separator didn't split anything — try next.
		return recursiveSplit(text, rest, chunkSize)
	}

	var result []string
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		if utf8.RuneCountInString(part) <= chunkSize {
			result = append(result, part)
		} else {
			// Still too large — recurse with remaining separators.
			result = append(result, recursiveSplit(part, rest, chunkSize)...)
		}
	}
	return result
}

// splitKeepSep splits text on sep. Unlike strings.Split, the separator
// is kept at the end of each segment (except the last) to preserve
// context like sentence-ending periods.
func splitKeepSep(text, sep string) []string {
	parts := strings.Split(text, sep)
	if len(parts) <= 1 {
		return parts
	}

	result := make([]string, 0, len(parts))
	for i, part := range parts {
		if i < len(parts)-1 {
			result = append(result, part+sep)
		} else if part != "" {
			result = append(result, part)
		}
	}
	return result
}

// hardSplit cuts text into pieces of at most chunkSize runes.
// Last resort when no separator works.
func hardSplit(text string, chunkSize int) []string {
	runes := []rune(text)
	var result []string
	for i := 0; i < len(runes); i += chunkSize {
		end := i + chunkSize
		if end > len(runes) {
			end = len(runes)
		}
		result = append(result, string(runes[i:end]))
	}
	return result
}

// mergeSegments combines small segments into chunks up to chunkSize,
// then applies overlap between adjacent chunks.
func mergeSegments(segments []string, chunkSize, overlap int) []Chunk {
	if len(segments) == 0 {
		return nil
	}

	// First pass: merge small adjacent segments into windows up to chunkSize.
	var merged []string
	var current strings.Builder

	for _, seg := range segments {
		segLen := utf8.RuneCountInString(seg)
		curLen := utf8.RuneCountInString(current.String())

		if curLen > 0 && curLen+segLen > chunkSize {
			// Flush current buffer.
			merged = append(merged, strings.TrimSpace(current.String()))
			current.Reset()
		}
		if current.Len() > 0 {
			current.WriteString(" ")
		}
		current.WriteString(seg)
	}
	if current.Len() > 0 {
		merged = append(merged, strings.TrimSpace(current.String()))
	}

	if overlap <= 0 || len(merged) <= 1 {
		// No overlap needed — convert directly to Chunks.
		chunks := make([]Chunk, len(merged))
		for i, text := range merged {
			chunks[i] = Chunk{
				Content:    text,
				Index:      i,
				TokenCount: estimateTokens(text),
				Metadata:   map[string]any{},
			}
		}
		return chunks
	}

	// Second pass: create overlapping chunks.
	// Each chunk includes up to `overlap` chars from the end of the
	// previous chunk as prefix context.
	chunks := make([]Chunk, 0, len(merged))
	for i, text := range merged {
		if i > 0 {
			prev := merged[i-1]
			suffix := overlapSuffix(prev, overlap)
			if suffix != "" {
				text = suffix + " " + text
			}
		}
		chunks = append(chunks, Chunk{
			Content:    text,
			Index:      i,
			TokenCount: estimateTokens(text),
			Metadata:   map[string]any{},
		})
	}
	return chunks
}

// overlapSuffix returns the last `n` characters of s, breaking at a
// word boundary to avoid splitting mid-word.
func overlapSuffix(s string, n int) string {
	runes := []rune(s)
	if len(runes) <= n {
		return s
	}

	start := len(runes) - n
	suffix := string(runes[start:])

	// Try to break at a space to avoid mid-word splits.
	if idx := strings.Index(suffix, " "); idx >= 0 && idx < len(suffix)/2 {
		suffix = suffix[idx+1:]
	}
	return suffix
}

// estimateTokens gives a rough token count. Most tokenizers average
// ~4 characters per token for English text.
func estimateTokens(s string) int {
	n := utf8.RuneCountInString(s)
	tokens := n / 4
	if tokens == 0 && n > 0 {
		tokens = 1
	}
	return tokens
}