Changeset 0.14.0 (#67)
This commit is contained in:
244
server/knowledge/chunker.go
Normal file
244
server/knowledge/chunker.go
Normal file
@@ -0,0 +1,244 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user