WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,268 @@
|
||||
// Package chunker - header_tracker.go implements context-preserving header
|
||||
// tracking for document chunking, ported from docreader/splitter/header_hook.py.
|
||||
//
|
||||
// When a large Markdown table is split across multiple chunks, each chunk after
|
||||
// the first would lose the table header context. The headerTracker detects table
|
||||
// headers and signals the merge logic to prepend them to subsequent chunks.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// headerTrackerHook defines a pattern pair for detecting contextual headers.
|
||||
// When startPattern matches a unit's text, that text becomes an "active header".
|
||||
// The header stays active until endPattern matches a subsequent unit.
|
||||
type headerTrackerHook struct {
|
||||
startPattern *regexp.Regexp
|
||||
endPattern *regexp.Regexp
|
||||
priority int
|
||||
}
|
||||
|
||||
// defaultHeaderHooks returns header tracking hooks matching the Python defaults
|
||||
// in docreader/splitter/header_hook.py.
|
||||
var defaultHeaderHooks = []headerTrackerHook{
|
||||
{
|
||||
// Markdown table: header row + separator row (e.g. "| A | B |\n| --- | --- |\n")
|
||||
startPattern: regexp.MustCompile(`(?si)^\s*(?:\|[^|\n]*)+[\r\n]+\s*(?:\|\s*:?-{3,}:?\s*)+\|?[\r\n]+$`),
|
||||
// Empty/whitespace line or a line that doesn't start with | or whitespace
|
||||
endPattern: regexp.MustCompile(`(?si)^\s*$|^\s*[^|\s].*$`),
|
||||
priority: 15,
|
||||
},
|
||||
}
|
||||
|
||||
// tableRowPattern matches a single Markdown table row: "| cell | cell | ... |\n"
|
||||
var tableRowPattern = regexp.MustCompile(`(?m)^\s*(?:\|[^|\n]*)+\|\s*$`)
|
||||
|
||||
// markdownTableHookPriority matches DEFAULT_CONFIGS / defaultHeaderHooks table hook.
|
||||
const markdownTableHookPriority = 15
|
||||
|
||||
// headerTracker maintains the state of active headers across split units.
|
||||
type headerTracker struct {
|
||||
hooks []headerTrackerHook
|
||||
activeHeaders map[int]string // priority -> header text
|
||||
endedHeaders map[int]bool // priorities that have been ended
|
||||
pendingExtend map[int]bool // headers with empty column names awaiting first data row
|
||||
// pendingTableBreak is set when a table row unit ends with a paragraph break
|
||||
// (the blank line between tables is consumed by \n\n splitting). The header
|
||||
// stays active until the next unit is seen so we can detect a new table.
|
||||
pendingTableBreak bool
|
||||
// headerEndedThisUnit tells mergeUnits to flush before the current unit when a
|
||||
// new table starts (column mismatch or pendingTableBreak + table row).
|
||||
headerEndedThisUnit bool
|
||||
}
|
||||
|
||||
func newHeaderTracker() *headerTracker {
|
||||
return &headerTracker{
|
||||
hooks: defaultHeaderHooks,
|
||||
activeHeaders: make(map[int]string),
|
||||
endedHeaders: make(map[int]bool),
|
||||
pendingExtend: make(map[int]bool),
|
||||
}
|
||||
}
|
||||
|
||||
// update checks split text for header start/end markers and updates internal state.
|
||||
func (ht *headerTracker) update(split string) {
|
||||
ht.headerEndedThisUnit = false
|
||||
|
||||
if ht.pendingTableBreak {
|
||||
ht.pendingTableBreak = false
|
||||
if _, active := ht.activeHeaders[markdownTableHookPriority]; active {
|
||||
if firstTableRowColumnCount(split) > 0 {
|
||||
ht.clearTableHeader()
|
||||
ht.headerEndedThisUnit = true
|
||||
} else {
|
||||
ht.clearTableHeader()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 1. Check for header-end markers among currently active headers
|
||||
for _, hook := range ht.hooks {
|
||||
if _, active := ht.activeHeaders[hook.priority]; active {
|
||||
if hook.endPattern.MatchString(split) {
|
||||
ht.endedHeaders[hook.priority] = true
|
||||
delete(ht.activeHeaders, hook.priority)
|
||||
delete(ht.pendingExtend, hook.priority)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 1b. Paragraph splits consume the blank line between tables. Mark a break
|
||||
// after "| last row |\n\n" and resolve on the next unit; also end when a new
|
||||
// table row has a different column count than the active header.
|
||||
if _, active := ht.activeHeaders[markdownTableHookPriority]; active {
|
||||
if !ht.pendingExtend[markdownTableHookPriority] {
|
||||
if splitEndsWithParagraphBreak(split) {
|
||||
ht.pendingTableBreak = true
|
||||
} else {
|
||||
ht.endTableHeaderOnColumnMismatch(split)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. If a header has an empty column-name row (e.g. "||"), replace it with
|
||||
// a proper Markdown table header using the first data row as column names.
|
||||
//
|
||||
// Before: "||" + "| --- | --- |\n"
|
||||
// After: "| col1 | col2 |\n" + "| --- | --- |\n"
|
||||
for p := range ht.pendingExtend {
|
||||
if _, active := ht.activeHeaders[p]; active && tableRowPattern.MatchString(split) {
|
||||
sep := extractSeparatorLine(ht.activeHeaders[p])
|
||||
ht.activeHeaders[p] = split + sep
|
||||
}
|
||||
delete(ht.pendingExtend, p)
|
||||
}
|
||||
|
||||
// 3. Check for new header-start markers (only for hooks that are neither active nor ended)
|
||||
for _, hook := range ht.hooks {
|
||||
if _, active := ht.activeHeaders[hook.priority]; active {
|
||||
continue
|
||||
}
|
||||
if ht.endedHeaders[hook.priority] {
|
||||
continue
|
||||
}
|
||||
if loc := hook.startPattern.FindString(split); loc != "" {
|
||||
ht.activeHeaders[hook.priority] = loc
|
||||
if isEmptyTableHeaderRow(loc) {
|
||||
ht.pendingExtend[hook.priority] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. If all headers ended, clear the ended set so future tables can be tracked
|
||||
if len(ht.activeHeaders) == 0 {
|
||||
for k := range ht.endedHeaders {
|
||||
delete(ht.endedHeaders, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getHeaders returns all active headers concatenated, sorted by priority descending.
|
||||
func (ht *headerTracker) getHeaders() string {
|
||||
if len(ht.activeHeaders) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
type entry struct {
|
||||
priority int
|
||||
text string
|
||||
}
|
||||
entries := make([]entry, 0, len(ht.activeHeaders))
|
||||
for p, t := range ht.activeHeaders {
|
||||
entries = append(entries, entry{p, t})
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool {
|
||||
return entries[i].priority > entries[j].priority
|
||||
})
|
||||
|
||||
parts := make([]string, len(entries))
|
||||
for i, e := range entries {
|
||||
parts[i] = e.text
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
// isEmptyTableHeaderRow checks if the header row (the line before the separator)
|
||||
// contains only pipes and whitespace — meaning the column names are empty.
|
||||
// This is common with MarkItDown and similar converters that produce tables like:
|
||||
//
|
||||
// ||
|
||||
// | --- | --- |
|
||||
// | real column A | real column B |
|
||||
func isEmptyTableHeaderRow(header string) bool {
|
||||
idx := strings.IndexByte(header, '\n')
|
||||
if idx < 0 {
|
||||
return false
|
||||
}
|
||||
row := strings.TrimSpace(header[:idx])
|
||||
for _, r := range row {
|
||||
if r != '|' && r != ' ' && r != '\t' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// extractSeparatorLine returns the separator line (e.g. "| --- | --- |\n") from
|
||||
// a table header string. It looks for the line containing "---".
|
||||
func extractSeparatorLine(header string) string {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
if strings.Contains(line, "---") {
|
||||
return line + "\n"
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (ht *headerTracker) clearTableHeader() {
|
||||
ht.endedHeaders[markdownTableHookPriority] = true
|
||||
delete(ht.activeHeaders, markdownTableHookPriority)
|
||||
delete(ht.pendingExtend, markdownTableHookPriority)
|
||||
}
|
||||
|
||||
func (ht *headerTracker) endTableHeaderOnColumnMismatch(split string) {
|
||||
header, ok := ht.activeHeaders[markdownTableHookPriority]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
rowCols := firstTableRowColumnCount(split)
|
||||
headerCols := headerTableColumnCount(header)
|
||||
if rowCols > 0 && headerCols > 0 && rowCols != headerCols {
|
||||
ht.clearTableHeader()
|
||||
ht.headerEndedThisUnit = true
|
||||
}
|
||||
}
|
||||
|
||||
func splitEndsWithParagraphBreak(split string) bool {
|
||||
trimmed := strings.TrimRight(split, " \t\r")
|
||||
return strings.HasSuffix(trimmed, "\n\n") || strings.HasSuffix(trimmed, "\r\n\r\n")
|
||||
}
|
||||
|
||||
func tableRowColumnCount(line string) int {
|
||||
line = strings.TrimSpace(line)
|
||||
if !strings.HasPrefix(line, "|") {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(line, "|")
|
||||
if len(parts) > 0 && strings.TrimSpace(parts[0]) == "" {
|
||||
parts = parts[1:]
|
||||
}
|
||||
if len(parts) > 0 && strings.TrimSpace(parts[len(parts)-1]) == "" {
|
||||
parts = parts[:len(parts)-1]
|
||||
}
|
||||
return len(parts)
|
||||
}
|
||||
|
||||
func firstTableRowColumnCount(text string) int {
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line != "" && tableRowPattern.MatchString(line) {
|
||||
return tableRowColumnCount(line)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func headerTableColumnCount(header string) int {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.Contains(line, "---") {
|
||||
continue
|
||||
}
|
||||
if n := tableRowColumnCount(line); n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// headerColumnMismatch reports whether the next split unit starts a new table
|
||||
// whose width differs from the active markdown table header.
|
||||
func headerColumnMismatch(headers, nextUnit string) bool {
|
||||
headerCols := headerTableColumnCount(headers)
|
||||
rowCols := firstTableRowColumnCount(nextUnit)
|
||||
return headerCols > 0 && rowCols > 0 && headerCols != rowCols
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
// Package chunker - heading_hierarchy.go tracks Markdown heading nesting so
|
||||
// that the heading-aware splitter (heading_splitter.go) can prepend a
|
||||
// breadcrumb like "# Top > ## Section > ### Subsection" to each chunk.
|
||||
//
|
||||
// Conceptually similar to header_tracker.go (which tracks table headers via
|
||||
// start/end hooks) but a Markdown heading does not have an explicit end —
|
||||
// it is ended by the next heading of equal or shallower depth. The hook
|
||||
// abstraction in header_tracker.go does not fit that pattern, so we model
|
||||
// it as an explicit level-stack.
|
||||
package chunker
|
||||
|
||||
import "strings"
|
||||
|
||||
// HeadingHierarchy maintains a stack of active Markdown headings indexed by
|
||||
// level (1..6). Pushing a level-N heading pops every entry of level >= N
|
||||
// because the previous siblings/descendants are no longer in scope.
|
||||
type HeadingHierarchy struct {
|
||||
// stack[i] holds the heading text for level i+1 (so stack[0] = H1).
|
||||
// Entries beyond the deepest active level are empty strings.
|
||||
stack [6]string
|
||||
depth int // current deepest active level (0 if no active heading)
|
||||
}
|
||||
|
||||
// NewHeadingHierarchy returns an empty hierarchy.
|
||||
func NewHeadingHierarchy() *HeadingHierarchy {
|
||||
return &HeadingHierarchy{}
|
||||
}
|
||||
|
||||
// Observe parses line and updates the hierarchy if line is a Markdown
|
||||
// heading. Returns the (level, headingText) when a heading was recognized,
|
||||
// or (0, "") otherwise. Lines that look like headings inside fenced code
|
||||
// blocks are NOT detected here — callers must avoid feeding code-block
|
||||
// content to Observe (the heading splitter does so).
|
||||
func (h *HeadingHierarchy) Observe(line string) (int, string) {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
return 0, ""
|
||||
}
|
||||
level := len(m[1])
|
||||
if level < 1 || level > 6 {
|
||||
return 0, ""
|
||||
}
|
||||
heading := strings.TrimSpace(m[2])
|
||||
// Replace this level and clear deeper ones — siblings/descendants of
|
||||
// the previous heading at this level are no longer in scope.
|
||||
h.stack[level-1] = heading
|
||||
for i := level; i < 6; i++ {
|
||||
h.stack[i] = ""
|
||||
}
|
||||
if level > h.depth {
|
||||
h.depth = level
|
||||
} else {
|
||||
// Recompute depth: it might shrink if we just pushed a shallower heading.
|
||||
h.depth = 0
|
||||
for i := 0; i < 6; i++ {
|
||||
if h.stack[i] != "" {
|
||||
h.depth = i + 1
|
||||
}
|
||||
}
|
||||
}
|
||||
return level, heading
|
||||
}
|
||||
|
||||
// Breadcrumb returns the current heading path joined by " > ", e.g.
|
||||
// "Chapter 1 > Section 2 > Subsection a". Returns "" when no headings
|
||||
// are active.
|
||||
func (h *HeadingHierarchy) Breadcrumb() string {
|
||||
if h.depth == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, h.depth)
|
||||
for i := 0; i < h.depth; i++ {
|
||||
if h.stack[i] != "" {
|
||||
parts = append(parts, h.stack[i])
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " > ")
|
||||
}
|
||||
|
||||
// BreadcrumbWithHashes returns the path with the original `#` prefixes,
|
||||
// suitable for embedding back into chunk content as a context header.
|
||||
// Example: "# Chapter 1\n## Section 2\n### Subsection a"
|
||||
func (h *HeadingHierarchy) BreadcrumbWithHashes() string {
|
||||
if h.depth == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
for i := 0; i < h.depth; i++ {
|
||||
if h.stack[i] == "" {
|
||||
continue
|
||||
}
|
||||
if sb.Len() > 0 {
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
sb.WriteString(strings.Repeat("#", i+1))
|
||||
sb.WriteByte(' ')
|
||||
sb.WriteString(h.stack[i])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// Depth returns the current deepest active heading level.
|
||||
func (h *HeadingHierarchy) Depth() int { return h.depth }
|
||||
|
||||
// Reset clears all state.
|
||||
func (h *HeadingHierarchy) Reset() {
|
||||
for i := range h.stack {
|
||||
h.stack[i] = ""
|
||||
}
|
||||
h.depth = 0
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHeadingHierarchy_LinearNesting(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section 1.1")
|
||||
h.Observe("### Subsection 1.1.1")
|
||||
if got := h.Breadcrumb(); got != "Chapter 1 > Section 1.1 > Subsection 1.1.1" {
|
||||
t.Errorf("breadcrumb mismatch: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_PopsDeeperOnSiblingHeading(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section 1.1")
|
||||
h.Observe("### Subsection 1.1.1")
|
||||
h.Observe("## Section 1.2") // pops the H3
|
||||
if got := h.Breadcrumb(); got != "Chapter 1 > Section 1.2" {
|
||||
t.Errorf("breadcrumb after sibling: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_PopsAllOnNewTopLevel(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section A")
|
||||
h.Observe("### Sub")
|
||||
h.Observe("# Chapter 2")
|
||||
if got := h.Breadcrumb(); got != "Chapter 2" {
|
||||
t.Errorf("breadcrumb after new H1: %q", got)
|
||||
}
|
||||
if h.Depth() != 1 {
|
||||
t.Errorf("depth should be 1, got %d", h.Depth())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_NonHeadingIgnored(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Title")
|
||||
level, txt := h.Observe("just a paragraph")
|
||||
if level != 0 || txt != "" {
|
||||
t.Errorf("non-heading should not register: %d %q", level, txt)
|
||||
}
|
||||
if got := h.Breadcrumb(); got != "Title" {
|
||||
t.Errorf("breadcrumb after non-heading: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_BreadcrumbWithHashes(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# A")
|
||||
h.Observe("## B")
|
||||
got := h.BreadcrumbWithHashes()
|
||||
want := "# A\n## B"
|
||||
if got != want {
|
||||
t.Errorf("hashes breadcrumb:\n got: %q\n want: %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_EmptyState(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
if h.Breadcrumb() != "" || h.BreadcrumbWithHashes() != "" || h.Depth() != 0 {
|
||||
t.Error("empty hierarchy should have no breadcrumb / depth=0")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_Reset(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# A")
|
||||
h.Observe("## B")
|
||||
h.Reset()
|
||||
if h.Breadcrumb() != "" || h.Depth() != 0 {
|
||||
t.Error("Reset did not clear state")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_SkipLevels(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
// Document jumps from H1 directly to H3, skipping H2.
|
||||
h.Observe("# Top")
|
||||
h.Observe("### Deep")
|
||||
if got := h.Breadcrumb(); got != "Top > Deep" {
|
||||
t.Errorf("skipped-level breadcrumb: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,342 @@
|
||||
// Package chunker - heading_splitter.go implements Tier 1: Markdown
|
||||
// heading-aware chunking. Documents with proper heading structure are split
|
||||
// at heading boundaries and each chunk is prefixed with a breadcrumb of
|
||||
// active heading context (e.g. "# Chapter 1\n## Section 1.2").
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// init wires this implementation into the strategy resolver.
|
||||
func init() {
|
||||
splitByHeadings = splitByHeadingsImpl
|
||||
}
|
||||
|
||||
// headingBoundary marks where a section starts. The first boundary is at
|
||||
// rune offset 0 (covers any preamble before the first heading), subsequent
|
||||
// boundaries sit at headings whose level is <= primaryLevel.
|
||||
type headingBoundary struct {
|
||||
runeStart int
|
||||
line string // raw heading line, "" when this is the leading boundary
|
||||
}
|
||||
|
||||
// splitByHeadingsImpl is the Tier-1 implementation. It falls through to the
|
||||
// legacy splitter when the document has no usable heading structure or when
|
||||
// the heading split would produce a single section anyway.
|
||||
//
|
||||
// profile may be nil; we compute one on demand. When the strategy resolver
|
||||
// already ran the profiler (auto strategy), the same profile is threaded
|
||||
// through here so we don't re-scan the entire document.
|
||||
func splitByHeadingsImpl(text string, cfg SplitterConfig, profile *DocProfile) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
if profile == nil {
|
||||
profile = ProfileDocument(text)
|
||||
}
|
||||
primaryLevel := profile.DominantHeadingLevel()
|
||||
if primaryLevel == 0 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
bounds := findHeadingBoundaries(text, primaryLevel)
|
||||
if len(bounds) <= 1 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
runes := []rune(text)
|
||||
hierarchy := NewHeadingHierarchy()
|
||||
|
||||
// Pre-walk every heading (not just primary-level) so the hierarchy
|
||||
// reflects the full nesting context for each section's start. We only
|
||||
// snapshot the breadcrumb at section boundaries; deeper sub-headings
|
||||
// inside a section update the hierarchy but do not change the chunk's
|
||||
// breadcrumb (chunks within a section share one breadcrumb).
|
||||
var out []Chunk
|
||||
seq := 0
|
||||
|
||||
for i, b := range bounds {
|
||||
endRune := len(runes)
|
||||
if i+1 < len(bounds) {
|
||||
endRune = bounds[i+1].runeStart
|
||||
}
|
||||
if b.line != "" {
|
||||
hierarchy.Observe(b.line)
|
||||
}
|
||||
// Catch sub-headings that occur between this primary boundary and
|
||||
// the next so the hierarchy stays in sync for subsequent sections.
|
||||
// We intentionally do this after observing the section header so
|
||||
// the breadcrumb reflects the section-leading heading.
|
||||
breadcrumb := hierarchy.BreadcrumbWithHashes()
|
||||
sectionStart := *hierarchy
|
||||
observeSubHeadings(runes[b.runeStart:endRune], primaryLevel, hierarchy)
|
||||
|
||||
sectionRunes := runes[b.runeStart:endRune]
|
||||
sectionContent := string(sectionRunes)
|
||||
secLen := len(sectionRunes)
|
||||
if secLen == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
bcLen := utf8.RuneCountInString(breadcrumb)
|
||||
// Single-chunk section: emit as-is, breadcrumb tracked separately.
|
||||
// The breadcrumb is delivered via Chunk.ContextHeader (not Content)
|
||||
// to preserve End-Start == len(Content) invariants relied on by
|
||||
// document reconstruction (knowledge.go:2278+).
|
||||
if bcLen+2+secLen <= cfg.ChunkSize {
|
||||
out = append(out, Chunk{
|
||||
Content: sectionContent,
|
||||
ContextHeader: breadcrumb,
|
||||
Seq: seq,
|
||||
Start: b.runeStart,
|
||||
End: endRune,
|
||||
})
|
||||
seq++
|
||||
continue
|
||||
}
|
||||
|
||||
// Section too large: defer to the legacy splitter for inner
|
||||
// segmentation. We do NOT shrink the inner ChunkSize budget here
|
||||
// because the breadcrumb no longer counts against Content size.
|
||||
// Each sub-chunk gets a breadcrumb reflecting the deepest heading
|
||||
// active at its start, so deep `###`/`####` sub-headings inside a
|
||||
// long section aren't collapsed to the section-level header.
|
||||
subBreadcrumbs := sectionBreadcrumbs(sectionRunes, primaryLevel, sectionStart)
|
||||
subChunks := SplitText(sectionContent, cfg)
|
||||
for _, sub := range subChunks {
|
||||
out = append(out, Chunk{
|
||||
Content: sub.Content,
|
||||
ContextHeader: breadcrumbAtOffset(subBreadcrumbs, sub.Start, breadcrumb),
|
||||
Seq: seq,
|
||||
Start: b.runeStart + sub.Start,
|
||||
End: b.runeStart + sub.End,
|
||||
})
|
||||
seq++
|
||||
}
|
||||
}
|
||||
|
||||
return coalesceTinyChunks(out, cfg.ChunkSize)
|
||||
}
|
||||
|
||||
// coalesceTinyChunks merges adjacent small chunks under their shared heading
|
||||
// context so that documents whose primary sections are mostly short (FAQs,
|
||||
// install logs, change-lists) don't trip the validator's "too many tiny
|
||||
// chunks" rule and fall through all the way to legacy. The merged breadcrumb
|
||||
// is the line-prefix shared by both inputs; the original sub-headings remain
|
||||
// visible because heading_splitter includes the heading line in each
|
||||
// section's Content.
|
||||
//
|
||||
// Safety:
|
||||
// - We only merge when cur.End == next.Start. That preserves the
|
||||
// End-Start == len([]rune(Content)) invariant that document
|
||||
// reconstruction relies on, and naturally skips legacy sub-chunks (which
|
||||
// may overlap due to ChunkOverlap).
|
||||
// - We stop accumulating once the running chunk reaches the merge target
|
||||
// (≈ ChunkSize/2) so we don't aggressively pack chunks beyond what the
|
||||
// validator considers comfortable.
|
||||
func coalesceTinyChunks(in []Chunk, chunkSize int) []Chunk {
|
||||
if len(in) <= 1 || chunkSize <= 0 {
|
||||
return in
|
||||
}
|
||||
target := chunkSize / 2
|
||||
if target < 200 {
|
||||
target = 200
|
||||
}
|
||||
|
||||
out := make([]Chunk, 0, len(in))
|
||||
cur := in[0]
|
||||
curLen := utf8.RuneCountInString(cur.Content)
|
||||
|
||||
for i := 1; i < len(in); i++ {
|
||||
next := in[i]
|
||||
nextLen := utf8.RuneCountInString(next.Content)
|
||||
sharedHeader := commonHeadingPrefix(cur.ContextHeader, next.ContextHeader)
|
||||
// Adjacent + still-small + would not blow the size budget → merge.
|
||||
if sharedHeader != "" && cur.End == next.Start && curLen < target && curLen+nextLen <= chunkSize {
|
||||
cur.Content += next.Content
|
||||
cur.ContextHeader = sharedHeader
|
||||
cur.End = next.End
|
||||
curLen += nextLen
|
||||
continue
|
||||
}
|
||||
out = append(out, cur)
|
||||
cur = next
|
||||
curLen = nextLen
|
||||
}
|
||||
out = append(out, cur)
|
||||
|
||||
// Re-sequence — downstream code (knowledge.go) expects Seq to be a dense
|
||||
// 0..N-1 range over the returned slice.
|
||||
for i := range out {
|
||||
out[i].Seq = i
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// commonHeadingPrefix returns the longest line-aligned prefix shared by two
|
||||
// breadcrumb strings. Heading hierarchies are emitted as
|
||||
// "# Top\n## Section\n### Sub", so a line-by-line comparison is sufficient
|
||||
// and avoids partial-line truncation that would corrupt the breadcrumb.
|
||||
func commonHeadingPrefix(a, b string) string {
|
||||
if a == b {
|
||||
return a
|
||||
}
|
||||
la := strings.Split(a, "\n")
|
||||
lb := strings.Split(b, "\n")
|
||||
n := len(la)
|
||||
if len(lb) < n {
|
||||
n = len(lb)
|
||||
}
|
||||
common := 0
|
||||
for i := 0; i < n; i++ {
|
||||
if la[i] != lb[i] {
|
||||
break
|
||||
}
|
||||
common = i + 1
|
||||
}
|
||||
if common == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(la[:common], "\n")
|
||||
}
|
||||
|
||||
// findHeadingBoundaries returns one boundary at offset 0 plus one per
|
||||
// Markdown heading at level <= primaryLevel that sits outside fenced code
|
||||
// blocks. Heading detection is line-oriented — a heading must occupy a
|
||||
// whole line to be recognized.
|
||||
func findHeadingBoundaries(text string, primaryLevel int) []headingBoundary {
|
||||
runes := []rune(text)
|
||||
bounds := []headingBoundary{{runeStart: 0}}
|
||||
if len(runes) == 0 {
|
||||
return bounds
|
||||
}
|
||||
|
||||
pos := 0
|
||||
inFence := false
|
||||
lines := strings.Split(text, "\n")
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // newline
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !inFence {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m != nil {
|
||||
level := len(m[1])
|
||||
if level >= 1 && level <= primaryLevel && pos > 0 {
|
||||
bounds = append(bounds, headingBoundary{
|
||||
runeStart: pos,
|
||||
line: line,
|
||||
})
|
||||
}
|
||||
if level >= 1 && level <= primaryLevel && pos == 0 {
|
||||
// First line is a heading — replace the leading boundary
|
||||
bounds[0].line = line
|
||||
}
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // account for the \n that strings.Split removed
|
||||
}
|
||||
}
|
||||
return bounds
|
||||
}
|
||||
|
||||
// observeSubHeadings walks the section's lines and feeds every Markdown
|
||||
// heading deeper than primaryLevel into the hierarchy. This keeps the
|
||||
// hierarchy state correct so the breadcrumb at the next primary section
|
||||
// reflects the truly active stack.
|
||||
func observeSubHeadings(runes []rune, primaryLevel int, h *HeadingHierarchy) {
|
||||
if len(runes) == 0 {
|
||||
return
|
||||
}
|
||||
text := string(runes)
|
||||
inFence := false
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
continue
|
||||
}
|
||||
if inFence {
|
||||
continue
|
||||
}
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
level := len(m[1])
|
||||
if level > primaryLevel {
|
||||
h.Observe(line)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sectionBreadcrumb pairs a rune offset within a section with the breadcrumb
|
||||
// in effect from that offset onward.
|
||||
type sectionBreadcrumb struct {
|
||||
runeStart int
|
||||
breadcrumb string
|
||||
}
|
||||
|
||||
// sectionBreadcrumbs walks a section's deeper sub-headings (level >
|
||||
// primaryLevel) and records, for each, the rune offset where it takes effect
|
||||
// and the resulting breadcrumb. seed is the hierarchy state at the section's
|
||||
// start (already including the section heading and its ancestors). The
|
||||
// returned slice is ordered by runeStart and always begins with the seed
|
||||
// breadcrumb at offset 0, so a sub-chunk sitting far below a deep heading
|
||||
// still resolves to that heading's path rather than the section header.
|
||||
func sectionBreadcrumbs(sectionRunes []rune, primaryLevel int, seed HeadingHierarchy) []sectionBreadcrumb {
|
||||
h := seed
|
||||
result := []sectionBreadcrumb{{runeStart: 0, breadcrumb: h.BreadcrumbWithHashes()}}
|
||||
pos := 0
|
||||
inFence := false
|
||||
lines := strings.Split(string(sectionRunes), "\n")
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !inFence {
|
||||
if m := MarkdownHeadingPattern.FindStringSubmatch(line); m != nil && len(m[1]) > primaryLevel {
|
||||
h.Observe(line)
|
||||
result = append(result, sectionBreadcrumb{
|
||||
runeStart: pos,
|
||||
breadcrumb: h.BreadcrumbWithHashes(),
|
||||
})
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// breadcrumbAtOffset returns the breadcrumb in effect at the given rune offset
|
||||
// — the last entry whose runeStart <= offset. fallback covers the (unreachable
|
||||
// in practice) empty-slice case.
|
||||
func breadcrumbAtOffset(bcs []sectionBreadcrumb, offset int, fallback string) string {
|
||||
bc := fallback
|
||||
for _, e := range bcs {
|
||||
if e.runeStart > offset {
|
||||
break
|
||||
}
|
||||
bc = e.breadcrumb
|
||||
}
|
||||
return bc
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitByHeadings_BasicSections(t *testing.T) {
|
||||
// Each section is intentionally larger than the merge-target (≈
|
||||
// ChunkSize/2) so the post-split coalesce pass leaves them as distinct
|
||||
// chunks. We're testing per-section emission + breadcrumb here, not
|
||||
// merging.
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4)
|
||||
doc := "# Top\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body + "\n\n## Section C\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected ≥3 chunks (one per section), got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Breadcrumb is delivered via ContextHeader, not Content.
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Top") {
|
||||
t.Errorf("chunk %d missing H1 in ContextHeader:\n%q", i, c.ContextHeader)
|
||||
}
|
||||
// EmbeddingContent merges header + content for the embedder.
|
||||
if !strings.Contains(c.EmbeddingContent(), "# Top") {
|
||||
t.Errorf("chunk %d EmbeddingContent missing H1", i)
|
||||
}
|
||||
}
|
||||
|
||||
found := false
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, "## Section B") && strings.Contains(c.Content, "Lorem ipsum") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("no chunk contains Section B with its body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_FallsThroughForUnstructuredDoc(t *testing.T) {
|
||||
doc := "Just a plain paragraph without any headings at all in this text."
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
// no headings → falls through to SplitText, which keeps the whole thing
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("expected fallthrough single chunk, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_LargeSectionRecursesIntoLegacy(t *testing.T) {
|
||||
body := strings.Repeat("This is a long sentence repeated many times. ", 50)
|
||||
doc := "# Top\n## Big\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{". "}}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("large section should be sub-split, got %d chunks", len(chunks))
|
||||
}
|
||||
// Every sub-chunk should carry the breadcrumb via ContextHeader.
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Top") {
|
||||
t.Errorf("sub-chunk %d missing H1 in ContextHeader", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_BreadcrumbReflectsLatestPath(t *testing.T) {
|
||||
// Sized so each section stays its own chunk after the tiny-section
|
||||
// coalesce pass — we're verifying breadcrumb assignment per section,
|
||||
// not the merge behavior.
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4)
|
||||
doc := "# Chapter 1\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected ≥3 chunks, got %d", len(chunks))
|
||||
}
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, "text B") {
|
||||
if strings.Contains(c.ContextHeader, "## Section A") {
|
||||
t.Errorf("Section B chunk should not include Section A in breadcrumb:\n%s", c.ContextHeader)
|
||||
}
|
||||
if !strings.Contains(c.ContextHeader, "## Section B") {
|
||||
t.Errorf("Section B chunk should include its own heading in breadcrumb:\n%s", c.ContextHeader)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_IgnoresHeadingsInsideCodeFence(t *testing.T) {
|
||||
doc := "# Real\n\n```\n# Fake heading inside code\n```\n\nbody"
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.ContextHeader, "# Real") || strings.Contains(c.Content, "# Real") {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("expected real H1 breadcrumb on some chunk")
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_PreservesPositionRelativeToOriginal(t *testing.T) {
|
||||
doc := "# Top\nintro\n\n## A\nbody A\n\n## B\nbody B"
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for i, c := range chunks {
|
||||
if c.Start < 0 {
|
||||
t.Errorf("chunk %d has negative Start", i)
|
||||
}
|
||||
if c.End < c.Start {
|
||||
t.Errorf("chunk %d End < Start", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_PositionInvariant ensures End-Start == len(Content)
|
||||
// and runes[Start:End] == Content for every emitted chunk. This invariant
|
||||
// is required by knowledge.go:2278+ document reconstruction logic.
|
||||
func TestSplitByHeadings_PositionInvariant(t *testing.T) {
|
||||
doc := `# Top
|
||||
intro paragraph here.
|
||||
|
||||
## Section A
|
||||
content of A here, several sentences.
|
||||
|
||||
## Section B
|
||||
content of B here.
|
||||
|
||||
## Section C
|
||||
content of C here.`
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
docRunes := []rune(doc)
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
span := c.End - c.Start
|
||||
if span != contentRuneLen {
|
||||
t.Errorf("chunk %d: span(%d) != content_runes(%d)\nContent:\n%q", i, span, contentRuneLen, c.Content)
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(docRunes) {
|
||||
if string(docRunes[c.Start:c.End]) != c.Content {
|
||||
t.Errorf("chunk %d: runes[Start:End] != Content", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalescesTinyAdjacentSections covers the FAQ /
|
||||
// install-log case where a parent heading hosts many short sub-sections.
|
||||
// Without merging, each `##` becomes its own <50-char chunk and the
|
||||
// validator rejects the tier with "too many tiny chunks". After merging,
|
||||
// they collapse into a small number of properly-sized chunks while still
|
||||
// surfacing the shared parent breadcrumb.
|
||||
func TestSplitByHeadings_CoalescesTinyAdjacentSections(t *testing.T) {
|
||||
doc := `# Install Log
|
||||
|
||||
## Docker镜像
|
||||
使用 daocloud 部署 v0.3.1。
|
||||
|
||||
## 前端老版本
|
||||
浏览器缓存了旧前端资源。
|
||||
|
||||
## 登录报错
|
||||
ERROR: column missing.
|
||||
|
||||
## 解析失败
|
||||
embedding 表缺列。`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected at least one chunk")
|
||||
}
|
||||
if len(chunks) >= 5 {
|
||||
t.Errorf("expected coalesce to produce <5 chunks, got %d", len(chunks))
|
||||
}
|
||||
// Every merged chunk must carry the shared parent in its breadcrumb so
|
||||
// retrieval can still answer "what document is this from".
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Install Log") {
|
||||
t.Errorf("chunk %d missing parent H1 in breadcrumb: %q", i, c.ContextHeader)
|
||||
}
|
||||
}
|
||||
// All four sub-section headings must remain visible somewhere in the
|
||||
// merged content (heading_splitter keeps the heading line as part of
|
||||
// each section's Content).
|
||||
for _, h := range []string{"## Docker镜像", "## 前端老版本", "## 登录报错", "## 解析失败"} {
|
||||
seen := false
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, h) {
|
||||
seen = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !seen {
|
||||
t.Errorf("merged chunks should still contain heading %q somewhere", h)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_DoesNotCoalesceDistinctTopLevelHeadings(t *testing.T) {
|
||||
doc := `# Intro
|
||||
short intro.
|
||||
|
||||
# Usage
|
||||
short usage.
|
||||
|
||||
# FAQ
|
||||
short faq.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("expected one chunk per top-level heading, got %d:\n%v", len(chunks), chunks)
|
||||
}
|
||||
for i, heading := range []string{"# Intro", "# Usage", "# FAQ"} {
|
||||
if !strings.Contains(chunks[i].Content, heading) {
|
||||
t.Errorf("chunk %d should contain heading %q, got:\n%s", i, heading, chunks[i].Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalescePreservesPositionInvariant guards the
|
||||
// End-Start == len([]rune(Content)) invariant after merging. Adjacent
|
||||
// chunks (cur.End == next.Start) must concatenate cleanly; the merge must
|
||||
// refuse to combine non-adjacent chunks (e.g. legacy sub-chunks from an
|
||||
// oversized section that overlap).
|
||||
func TestSplitByHeadings_CoalescePreservesPositionInvariant(t *testing.T) {
|
||||
doc := `# Top
|
||||
|
||||
## A
|
||||
short A.
|
||||
|
||||
## B
|
||||
short B.
|
||||
|
||||
## C
|
||||
short C.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
docRunes := []rune(doc)
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
if c.End-c.Start != contentRuneLen {
|
||||
t.Errorf("chunk %d: End-Start(%d) != content_runes(%d) after merge",
|
||||
i, c.End-c.Start, contentRuneLen)
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(docRunes) {
|
||||
if string(docRunes[c.Start:c.End]) != c.Content {
|
||||
t.Errorf("chunk %d: source[Start:End] != Content after merge", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalesceRespectsChunkSize ensures the merge target
|
||||
// stays within the ChunkSize budget — the validator caps oversize chunks
|
||||
// at 2x and we should never approach that line via merging.
|
||||
func TestSplitByHeadings_CoalesceRespectsChunkSize(t *testing.T) {
|
||||
const sections = 30
|
||||
var sb strings.Builder
|
||||
sb.WriteString("# Doc\n")
|
||||
for i := 0; i < sections; i++ {
|
||||
sb.WriteString("\n## Section ")
|
||||
sb.WriteString(strings.Repeat("X", 1)) // unique-ish heading
|
||||
sb.WriteString("\nshort body line.\n")
|
||||
}
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(sb.String(), cfg, nil)
|
||||
for i, c := range chunks {
|
||||
if l := len([]rune(c.Content)); l > cfg.ChunkSize {
|
||||
t.Errorf("chunk %d exceeds ChunkSize: %d > %d", i, l, cfg.ChunkSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCommonHeadingPrefix exercises the breadcrumb-prefix helper directly.
|
||||
func TestCommonHeadingPrefix(t *testing.T) {
|
||||
cases := []struct {
|
||||
a, b, want string
|
||||
}{
|
||||
{"# Top\n## A", "# Top\n## B", "# Top"},
|
||||
{"# Top", "# Top", "# Top"},
|
||||
{"# X", "# Y", ""},
|
||||
{"# Top\n## A\n### x", "# Top\n## A\n### y", "# Top\n## A"},
|
||||
{"", "# Top", ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
got := commonHeadingPrefix(tc.a, tc.b)
|
||||
if got != tc.want {
|
||||
t.Errorf("commonHeadingPrefix(%q, %q) = %q, want %q", tc.a, tc.b, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_NoBreadcrumbDuplication ensures the section's own
|
||||
// heading line does not appear twice in the chunk content (once as part of
|
||||
// the breadcrumb, once as the section's first line).
|
||||
func TestSplitByHeadings_NoBreadcrumbDuplication(t *testing.T) {
|
||||
doc := `# Chapter 1
|
||||
intro.
|
||||
|
||||
## Section A
|
||||
body A.
|
||||
|
||||
## Section B
|
||||
body B.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for i, c := range chunks {
|
||||
// Count occurrences of "## Section A" / "## Section B"
|
||||
for _, heading := range []string{"## Section A", "## Section B"} {
|
||||
n := strings.Count(c.Content, heading)
|
||||
if n > 1 {
|
||||
t.Errorf("chunk %d contains %q %d times — duplicated by breadcrumb prepend:\n%s",
|
||||
i, heading, n, c.Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_DeepSubHeadingInLargeSection covers issue #1674: when a
|
||||
// shallow heading level repeats and therefore becomes the dominant split
|
||||
// backbone, a long section split into sub-chunks must still carry the deep
|
||||
// `###`/`####` heading that defines a sub-chunk's meaning, not just the
|
||||
// section-level header. The marker line sits far below its deep heading, in a
|
||||
// later sub-chunk that has no heading of its own.
|
||||
func TestSplitByHeadings_DeepSubHeadingInLargeSection(t *testing.T) {
|
||||
filler := strings.Repeat("clause body sentence that pads the section out. ", 20)
|
||||
doc := "# Standard XYZ\n" +
|
||||
"## Preface\n" + filler + "\n\n" +
|
||||
"## 5 Classification\n" + filler + "\n\n" +
|
||||
"### 5.9 Grade Nine\n" + filler + "\n\n" +
|
||||
"#### 5.9.2 Clause Series\n" + filler + "\n\n" +
|
||||
"the item users search for is MARKER_ITEM_23 graded here.\n\n" +
|
||||
"## Appendix A\n" + filler + "\n\n" +
|
||||
"## Appendix B\n" + filler + "\n\n" +
|
||||
"## Appendix C\n" + filler
|
||||
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0, Separators: []string{". "}}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
|
||||
var marker *Chunk
|
||||
for i := range chunks {
|
||||
if strings.Contains(chunks[i].Content, "MARKER_ITEM_23") {
|
||||
marker = &chunks[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if marker == nil {
|
||||
t.Fatal("no chunk contains the marker line")
|
||||
}
|
||||
if !strings.Contains(marker.ContextHeader, "#### 5.9.2 Clause Series") {
|
||||
t.Errorf("marker chunk lost its deep sub-heading in the breadcrumb:\nContextHeader=%q", marker.ContextHeader)
|
||||
}
|
||||
if !strings.Contains(marker.ContextHeader, "## 5 Classification") {
|
||||
t.Errorf("marker chunk lost its section heading in the breadcrumb:\nContextHeader=%q", marker.ContextHeader)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
// Package chunker - heuristic_splitter.go implements Tier 2: boundary-driven
|
||||
// chunking for documents that lack proper Markdown headings but contain
|
||||
// recognizable structural cues (page breaks, numbered sections, multilingual
|
||||
// chapter markers, visual separators, all-caps section titles, page footers).
|
||||
//
|
||||
// The algorithm finds all candidate boundary positions, then performs greedy
|
||||
// bin-packing — accumulating blocks between boundaries into chunks until the
|
||||
// next block would exceed cfg.ChunkSize. Blocks larger than ChunkSize are
|
||||
// recursively delegated to the legacy splitter for inner segmentation.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func init() {
|
||||
splitByHeuristics = splitByHeuristicsImpl
|
||||
}
|
||||
|
||||
// boundary marks a candidate split point in the document.
|
||||
type boundary struct {
|
||||
runeStart int // rune offset where the next chunk should start
|
||||
priority int
|
||||
}
|
||||
|
||||
// splitByHeuristicsImpl is the Tier-2 implementation. Falls through to the
|
||||
// legacy splitter when no heuristic boundaries are found.
|
||||
//
|
||||
// profile is currently unused (this tier scans for boundaries directly) but
|
||||
// is accepted to keep the splitByHeadings / splitByHeuristics signatures
|
||||
// uniform — see strategy.runTier.
|
||||
func splitByHeuristicsImpl(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
runes := []rune(text)
|
||||
totalRunes := len(runes)
|
||||
if totalRunes <= cfg.ChunkSize {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
bounds := findHeuristicBoundaries(text, cfg.Languages)
|
||||
// Drop any boundary that falls strictly inside a protected region (table,
|
||||
// fenced code block, LaTeX block, etc.) — splitting there would cut
|
||||
// through atomic content. Boundaries on a span edge are kept since they
|
||||
// align with the protected region start/end.
|
||||
if prot := protectedSpansRune(text, protectedSpans(text)); len(prot) > 0 {
|
||||
bounds = dropBoundsInsideSpans(bounds, prot)
|
||||
}
|
||||
if len(bounds) == 0 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// Append a sentinel at end-of-document so the bin-packer can flush.
|
||||
bounds = append(bounds, boundary{runeStart: totalRunes})
|
||||
// Always start with a boundary at offset 0 if not already there.
|
||||
if bounds[0].runeStart != 0 {
|
||||
bounds = append([]boundary{{runeStart: 0}}, bounds...)
|
||||
}
|
||||
|
||||
// Greedy bin-packing.
|
||||
var out []Chunk
|
||||
seq := 0
|
||||
chunkStart := bounds[0].runeStart
|
||||
curEnd := chunkStart
|
||||
minChunkSize := cfg.ChunkSize / 4
|
||||
if minChunkSize < 50 {
|
||||
minChunkSize = 50
|
||||
}
|
||||
|
||||
for i := 1; i < len(bounds); i++ {
|
||||
nextEnd := bounds[i].runeStart
|
||||
blockLen := nextEnd - curEnd
|
||||
|
||||
if blockLen > cfg.ChunkSize {
|
||||
// The block between the previous and this boundary is itself too
|
||||
// large to fit in any chunk. Flush current accumulation, then
|
||||
// recursively chunk the oversize block via the legacy splitter.
|
||||
if curEnd-chunkStart > 0 {
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
chunkStart = curEnd
|
||||
}
|
||||
out = appendOversizeBlock(out, runes, curEnd, nextEnd, cfg, &seq)
|
||||
curEnd = nextEnd
|
||||
chunkStart = nextEnd
|
||||
continue
|
||||
}
|
||||
|
||||
// Would adding this block exceed the budget?
|
||||
accumulated := nextEnd - chunkStart
|
||||
if accumulated > cfg.ChunkSize && curEnd-chunkStart >= minChunkSize {
|
||||
// Flush accumulated content as a chunk, restart at curEnd.
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
// Snap overlap start to the nearest semantic boundary or line
|
||||
// break instead of slicing mid-line / mid-word.
|
||||
chunkStart = applyOverlapAligned(runes, curEnd, cfg.ChunkOverlap, bounds)
|
||||
}
|
||||
curEnd = nextEnd
|
||||
}
|
||||
|
||||
// Flush remaining content.
|
||||
if curEnd > chunkStart {
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// findHeuristicBoundaries scans text and returns boundary positions in
|
||||
// ascending order. Lower-priority duplicates at the same offset are dropped.
|
||||
func findHeuristicBoundaries(text string, langs []string) []boundary {
|
||||
var bounds []boundary
|
||||
|
||||
// Form feeds — strongest single-character boundary.
|
||||
for _, idx := range allRuneIndices(text, "\f") {
|
||||
bounds = append(bounds, boundary{runeStart: idx, priority: PrioFormFeed})
|
||||
}
|
||||
|
||||
// Per-line patterns walk the text once, line by line.
|
||||
lines := strings.Split(text, "\n")
|
||||
chapterPatterns := ChapterPatternsForLangs(langs)
|
||||
pos := 0
|
||||
inFence := false
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
} else if !inFence {
|
||||
runeStart := pos
|
||||
added := false
|
||||
for _, pat := range chapterPatterns {
|
||||
if pat.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioChapterMarker})
|
||||
added = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !added && NumberedSectionPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioNumberedHead})
|
||||
added = true
|
||||
}
|
||||
if !added && AllCapsHeadingPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioAllCapsHeading})
|
||||
added = true
|
||||
}
|
||||
if !added && VisualSeparatorPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioVisualSep})
|
||||
added = true
|
||||
}
|
||||
if !added && PageFooterPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioPageFooter})
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // \n
|
||||
}
|
||||
}
|
||||
|
||||
// Excessive blank blocks (\n{3,}). Match at the *start* of the run so we
|
||||
// drop into the next paragraph cleanly.
|
||||
for _, idx := range ExcessiveBlanksPattern.FindAllStringIndex(text, -1) {
|
||||
runeStart := utf8.RuneCountInString(text[:idx[1]])
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioBlankBlock})
|
||||
}
|
||||
|
||||
if len(bounds) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Sort by position; drop near-duplicate offsets keeping the highest priority.
|
||||
sort.Slice(bounds, func(i, j int) bool {
|
||||
if bounds[i].runeStart != bounds[j].runeStart {
|
||||
return bounds[i].runeStart < bounds[j].runeStart
|
||||
}
|
||||
return bounds[i].priority > bounds[j].priority
|
||||
})
|
||||
deduped := bounds[:0]
|
||||
prev := -1
|
||||
for _, b := range bounds {
|
||||
if b.runeStart != prev {
|
||||
deduped = append(deduped, b)
|
||||
prev = b.runeStart
|
||||
}
|
||||
}
|
||||
return deduped
|
||||
}
|
||||
|
||||
// dropBoundsInsideSpans returns bounds with entries that fall strictly
|
||||
// inside any of the (rune-offset) protected spans removed. Bounds at a
|
||||
// span's start or end are kept — they align with the span edge and don't
|
||||
// split protected content. spans must be sorted by start.
|
||||
func dropBoundsInsideSpans(bounds []boundary, spans []span) []boundary {
|
||||
if len(spans) == 0 {
|
||||
return bounds
|
||||
}
|
||||
out := bounds[:0]
|
||||
boundLoop:
|
||||
for _, b := range bounds {
|
||||
for _, s := range spans {
|
||||
if s.start >= b.runeStart {
|
||||
break // remaining spans start at or after b — can't contain b
|
||||
}
|
||||
if b.runeStart < s.end {
|
||||
continue boundLoop
|
||||
}
|
||||
}
|
||||
out = append(out, b)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// allRuneIndices returns every rune offset where needle starts in text.
|
||||
// Only used for single-rune needles like form-feed.
|
||||
func allRuneIndices(text, needle string) []int {
|
||||
var out []int
|
||||
if needle == "" {
|
||||
return out
|
||||
}
|
||||
pos := 0
|
||||
for _, r := range text {
|
||||
if string(r) == needle {
|
||||
out = append(out, pos)
|
||||
}
|
||||
pos++
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// appendChunk slices runes[start:end] into a Chunk and appends it to out.
|
||||
// Pure-whitespace slices are skipped (boundary clustering can occasionally
|
||||
// produce them). The Content stored is the raw slice — Start/End rune
|
||||
// offsets must match utf8.RuneCountInString(Content) for downstream
|
||||
// reconstruction code; whitespace stripping for embedding happens in
|
||||
// Chunk.EmbeddingContent.
|
||||
func appendChunk(out []Chunk, runes []rune, start, end int, seq *int) []Chunk {
|
||||
if end <= start {
|
||||
return out
|
||||
}
|
||||
raw := string(runes[start:end])
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
return out
|
||||
}
|
||||
c := Chunk{Content: raw, Seq: *seq, Start: start, End: end}
|
||||
*seq++
|
||||
return append(out, c)
|
||||
}
|
||||
|
||||
// appendOversizeBlock recursively chunks a region that is itself larger than
|
||||
// cfg.ChunkSize, using the legacy splitter so internal length budgets and
|
||||
// protected patterns are still respected.
|
||||
func appendOversizeBlock(out []Chunk, runes []rune, start, end int, cfg SplitterConfig, seq *int) []Chunk {
|
||||
if end <= start {
|
||||
return out
|
||||
}
|
||||
subText := string(runes[start:end])
|
||||
subs := SplitText(subText, cfg)
|
||||
for _, s := range subs {
|
||||
out = append(out, Chunk{
|
||||
Content: s.Content,
|
||||
Seq: *seq,
|
||||
Start: start + s.Start,
|
||||
End: start + s.End,
|
||||
})
|
||||
*seq++
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// applyOverlapAligned returns the rune offset where the next chunk should
|
||||
// start. The target is `curEnd - overlap`, but we snap to the nearest
|
||||
// preceding boundary (within 2x overlap) or, failing that, the previous
|
||||
// newline so chunks don't begin mid-line / mid-word. Falls back to the raw
|
||||
// target only if neither option is available.
|
||||
//
|
||||
// curEnd itself is always a boundary (the bin-packer flushes at boundary
|
||||
// positions), so we exclude it from the search — picking it would yield
|
||||
// zero overlap, defeating the purpose of this function.
|
||||
func applyOverlapAligned(runes []rune, curEnd, overlap int, bounds []boundary) int {
|
||||
if overlap <= 0 {
|
||||
return curEnd
|
||||
}
|
||||
target := curEnd - overlap
|
||||
if target < 0 {
|
||||
target = 0
|
||||
}
|
||||
// Allowed search window: [curEnd - 2*overlap, curEnd)
|
||||
windowStart := curEnd - 2*overlap
|
||||
if windowStart < 0 {
|
||||
windowStart = 0
|
||||
}
|
||||
|
||||
// Prefer a semantic boundary strictly inside the window.
|
||||
bestBound := -1
|
||||
for _, b := range bounds {
|
||||
if b.runeStart >= windowStart && b.runeStart < curEnd && b.runeStart > bestBound {
|
||||
bestBound = b.runeStart
|
||||
}
|
||||
}
|
||||
if bestBound >= 0 {
|
||||
return bestBound
|
||||
}
|
||||
|
||||
// Fallback: scan backwards from `target` to the previous newline, but
|
||||
// not past windowStart so we keep the overlap roughly the right size.
|
||||
for i := target; i > windowStart && i < len(runes); i-- {
|
||||
if runes[i] == '\n' {
|
||||
return i + 1
|
||||
}
|
||||
}
|
||||
return target
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitByHeuristics_FormFeedBoundary(t *testing.T) {
|
||||
doc := strings.Repeat("page one body text. ", 30) + "\f" + strings.Repeat("page two body. ", 30)
|
||||
cfg := SplitterConfig{ChunkSize: 400, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("form feed should produce ≥2 chunks, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_NumberedSections(t *testing.T) {
|
||||
body := strings.Repeat("body sentence. ", 8)
|
||||
doc := "1. Introduction\n" + body + "\n\n2. Methods\n" + body + "\n\n3. Results\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("numbered sections should split: got %d chunks", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_GermanChapterMarkers(t *testing.T) {
|
||||
body := strings.Repeat("Beispieltext. ", 10)
|
||||
doc := "Kapitel 1: Einführung\n" + body + "\n\nKapitel 2: Hauptteil\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("German chapter markers should split: got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_ChineseChapterMarkers(t *testing.T) {
|
||||
body := strings.Repeat("内容内容内容。", 60)
|
||||
doc := "第一章 引言\n" + body + "\n\n第二章 方法\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{"。"}, Languages: []string{LangChinese}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("Chinese chapter markers should split: got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_FallsThroughForUnstructuredDoc(t *testing.T) {
|
||||
doc := strings.Repeat("plain prose without structure. ", 5)
|
||||
cfg := SplitterConfig{ChunkSize: 1000, ChunkOverlap: 20}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("unstructured short doc should be one chunk, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_OversizeBlockRecursesIntoLegacy(t *testing.T) {
|
||||
huge := strings.Repeat("This is a long sentence. ", 200) // ~5000 chars
|
||||
doc := "1. Intro\n" + huge
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 50, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 5 {
|
||||
t.Errorf("oversize block should produce many sub-chunks, got %d", len(chunks))
|
||||
}
|
||||
// No single chunk should massively exceed the budget.
|
||||
for i, c := range chunks {
|
||||
if len([]rune(c.Content)) > 2*cfg.ChunkSize {
|
||||
t.Errorf("chunk %d exceeds 2x size: %d runes", i, len([]rune(c.Content)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_BoundariesAreOrdered(t *testing.T) {
|
||||
doc := "Kapitel 1: A\nbody\n\n---\n\n2. Section B\nbody\n\nPage 3 of 10\n\n第三章 C\nbody"
|
||||
bounds := findHeuristicBoundaries(doc, nil)
|
||||
if len(bounds) < 2 {
|
||||
t.Fatalf("expected multiple boundaries, got %d", len(bounds))
|
||||
}
|
||||
for i := 1; i < len(bounds); i++ {
|
||||
if bounds[i].runeStart < bounds[i-1].runeStart {
|
||||
t.Errorf("bounds not sorted: %d before %d", bounds[i].runeStart, bounds[i-1].runeStart)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_EmptyText(t *testing.T) {
|
||||
if got := splitByHeuristicsImpl("", DefaultConfig(), nil); got != nil {
|
||||
t.Errorf("empty doc should be nil, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: applyOverlapAligned previously included curEnd itself in its
|
||||
// boundary search, and curEnd is always one of the bounds (the bin-packer
|
||||
// flushes at boundary positions). The function therefore always returned
|
||||
// curEnd, producing zero overlap regardless of cfg.ChunkOverlap.
|
||||
func TestSplitByHeuristics_OverlapActuallyOverlaps(t *testing.T) {
|
||||
// Build many small numbered sections so the bin-packer flushes mid-doc
|
||||
// with at least one earlier boundary inside the overlap window.
|
||||
var sb strings.Builder
|
||||
for i := 1; i <= 12; i++ {
|
||||
sb.WriteString("\n\n")
|
||||
sb.WriteByte(byte('0' + i%10))
|
||||
sb.WriteString(". ")
|
||||
sb.WriteString(strings.Repeat("alpha beta gamma. ", 4)) // ~72 chars / section
|
||||
}
|
||||
doc := sb.String()
|
||||
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 80, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("need >=2 chunks to test overlap, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// At least one consecutive chunk pair must share a non-trivial suffix /
|
||||
// prefix. We don't require *every* pair to overlap (oversize blocks
|
||||
// short-circuit through legacy and reset chunkStart), but at least one
|
||||
// regular flush boundary should produce real overlap.
|
||||
saw := false
|
||||
for i := 1; i < len(chunks); i++ {
|
||||
prev := strings.TrimSpace(chunks[i-1].Content)
|
||||
cur := strings.TrimSpace(chunks[i].Content)
|
||||
// Walk back the longest suffix of prev that prefixes cur.
|
||||
match := 0
|
||||
maxScan := len(prev)
|
||||
if len(cur) < maxScan {
|
||||
maxScan = len(cur)
|
||||
}
|
||||
for n := 1; n <= maxScan; n++ {
|
||||
if strings.HasPrefix(cur, prev[len(prev)-n:]) {
|
||||
match = n
|
||||
}
|
||||
}
|
||||
if match >= 20 {
|
||||
saw = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Fatalf("expected at least one chunk pair to overlap by >=20 chars, none did. chunk sizes: %v",
|
||||
chunkLengths(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// Heuristic boundaries that fall inside protected regions (LaTeX block,
|
||||
// table, link, etc.) must be dropped so the bin-packer doesn't break
|
||||
// atomic content. Without the protected-span filter, a numbered-section
|
||||
// looking line inside a $$...$$ math block would be picked as a boundary.
|
||||
func TestSplitByHeuristics_DropsBoundariesInsideProtectedSpans(t *testing.T) {
|
||||
body := strings.Repeat("filler. ", 30)
|
||||
// LaTeX block whose middle line matches NumberedSectionPattern. The
|
||||
// filter should drop that boundary so the math block stays intact.
|
||||
doc := body + "\n\n$$\nx = 1\n1. equation step one\ny = 2\n$$\n\n" + body
|
||||
|
||||
bounds := findHeuristicBoundaries(doc, nil)
|
||||
prot := protectedSpansRune(doc, protectedSpans(doc))
|
||||
if len(prot) == 0 {
|
||||
t.Fatalf("expected protected spans for doc, got none")
|
||||
}
|
||||
filtered := dropBoundsInsideSpans(bounds, prot)
|
||||
for _, b := range filtered {
|
||||
for _, s := range prot {
|
||||
if b.runeStart > s.start && b.runeStart < s.end {
|
||||
t.Errorf("boundary %d still inside protected span [%d,%d)", b.runeStart, s.start, s.end)
|
||||
}
|
||||
}
|
||||
}
|
||||
// And it should actually have removed at least one boundary.
|
||||
if len(filtered) >= len(bounds) {
|
||||
t.Errorf("filter removed nothing: before=%d after=%d", len(bounds), len(filtered))
|
||||
}
|
||||
}
|
||||
|
||||
func chunkLengths(chunks []Chunk) []int {
|
||||
out := make([]int, len(chunks))
|
||||
for i, c := range chunks {
|
||||
out[i] = len([]rune(c.Content))
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package chunker - patterns.go is the source of truth for multilingual
|
||||
// regex patterns used by the heading-aware and heuristic splitters.
|
||||
//
|
||||
// Patterns are grouped by purpose (chapter markers, numbering, separators)
|
||||
// and tagged with a priority that the heuristic splitter uses to rank
|
||||
// candidate chunk boundaries.
|
||||
package chunker
|
||||
|
||||
import "regexp"
|
||||
|
||||
// BoundaryPriority levels for heuristic chunk boundaries. Higher = stronger.
|
||||
const (
|
||||
PrioFormFeed = 100
|
||||
PrioNumberedHead = 90
|
||||
PrioChapterMarker = 85
|
||||
PrioAllCapsHeading = 70
|
||||
PrioVisualSep = 60
|
||||
PrioPageFooter = 50
|
||||
PrioBlankBlock = 40
|
||||
)
|
||||
|
||||
// MarkdownHeadingPattern matches an ATX-style Markdown heading at line start.
|
||||
// Capture groups: (1) hashes, (2) heading text.
|
||||
var MarkdownHeadingPattern = regexp.MustCompile(`(?m)^(#{1,6})\s+(.+?)\s*#*\s*$`)
|
||||
|
||||
// FormFeedPattern matches the form-feed control character used by some PDF
|
||||
// converters as a page break marker.
|
||||
var FormFeedPattern = regexp.MustCompile(`\f`)
|
||||
|
||||
// NumberedSectionPattern matches lines starting with numeric or roman numbering
|
||||
// followed by a non-empty title, e.g. "1. Intro", "2.3 Methods", "IV. Results",
|
||||
// "2.2.1 用户与权限". The trailing dot after a multi-level numeral is optional
|
||||
// because many technical documents write "1.1 Foo" without a closing dot.
|
||||
var NumberedSectionPattern = regexp.MustCompile(`(?m)^[ \t]*(?:\d+(?:\.\d+){1,3}\.?|(?:\d+|[IVX]{1,5})\.)[ \t]+\S.{0,200}$`)
|
||||
|
||||
// AllCapsHeadingPattern matches short all-caps lines (likely section titles
|
||||
// rendered without Markdown headings). It requires at least 4 letters and
|
||||
// up to ~10 words. Trailing colons are tolerated.
|
||||
var AllCapsHeadingPattern = regexp.MustCompile(`(?m)^[ \t]*([A-ZÄÖÜ][A-ZÄÖÜ \-]{3,80}):?\s*$`)
|
||||
|
||||
// VisualSeparatorPattern matches horizontal rules / divider lines used as
|
||||
// section separators in plain text or pre-Markdown documents.
|
||||
var VisualSeparatorPattern = regexp.MustCompile(`(?m)^[ \t]*(?:-{3,}|={3,}|\*{3,}|_{3,})[ \t]*$`)
|
||||
|
||||
// ExcessiveBlanksPattern matches three or more consecutive newlines, which
|
||||
// usually denote a hard section break.
|
||||
var ExcessiveBlanksPattern = regexp.MustCompile(`\n{3,}`)
|
||||
|
||||
// PageFooterPattern matches typical "Seite X von Y" / "Page X of Y" lines.
|
||||
var PageFooterPattern = regexp.MustCompile(`(?mi)^[ \t]*(?:Seite|Page|页码?)\s+\d+(?:\s*(?:von|of|/)\s*\d+)?[ \t]*$`)
|
||||
|
||||
// GermanChapterPattern matches German chapter / section markers.
|
||||
var GermanChapterPattern = regexp.MustCompile(`(?m)^[ \t]*(?:Kapitel|Abschnitt|Teil)\s+(?:[0-9]+|[IVX]{1,5})[\.: ].{0,200}$`)
|
||||
|
||||
// EnglishChapterPattern matches English chapter / section markers.
|
||||
var EnglishChapterPattern = regexp.MustCompile(`(?m)^[ \t]*(?:Chapter|Section|Part)\s+(?:[0-9]+|[IVX]{1,5})[\.: ].{0,200}$`)
|
||||
|
||||
// ChineseChapterPattern matches CJK chapter / section markers like 第一章,
|
||||
// 第3节, 第 1 章 (whitespace between 第 / numeral / unit is tolerated).
|
||||
var ChineseChapterPattern = regexp.MustCompile(`(?m)^[ \t]*第[ \t]*[一二三四五六七八九十百千零〇0-9]+[ \t]*(?:章|节|節|部分|篇)[ \t]?.{0,200}$`)
|
||||
|
||||
// SentenceSeparators returns sentence-level separators tuned for the language.
|
||||
// Used for fine-grained sub-splitting when a section is still too large.
|
||||
func SentenceSeparators(lang string) []string {
|
||||
switch lang {
|
||||
case LangChinese:
|
||||
return []string{"。", "!", "?", ";", "\n"}
|
||||
case LangGerman, LangEnglish:
|
||||
return []string{". ", "! ", "? ", "; ", "\n"}
|
||||
default:
|
||||
return []string{"。", "!", "?", ";", ". ", "! ", "? ", "; ", "\n"}
|
||||
}
|
||||
}
|
||||
|
||||
// ChapterPatternsForLangs returns the chapter-marker regexes that apply for
|
||||
// the given language hints. An empty / unknown list returns all of them so
|
||||
// that auto-detected documents still match.
|
||||
func ChapterPatternsForLangs(langs []string) []*regexp.Regexp {
|
||||
if len(langs) == 0 {
|
||||
return []*regexp.Regexp{GermanChapterPattern, EnglishChapterPattern, ChineseChapterPattern}
|
||||
}
|
||||
var out []*regexp.Regexp
|
||||
for _, l := range langs {
|
||||
switch l {
|
||||
case LangGerman:
|
||||
out = append(out, GermanChapterPattern)
|
||||
case LangEnglish:
|
||||
out = append(out, EnglishChapterPattern)
|
||||
case LangChinese:
|
||||
out = append(out, ChineseChapterPattern)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
out = []*regexp.Regexp{GermanChapterPattern, EnglishChapterPattern, ChineseChapterPattern}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestMarkdownHeadingPattern_BasicLevels(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"# Heading 1", true},
|
||||
{"## Heading 2", true},
|
||||
{"###### Heading 6", true},
|
||||
{"####### Too many", false},
|
||||
{"#NoSpace", false},
|
||||
{" # Indented", false},
|
||||
{"plain text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := MarkdownHeadingPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("MarkdownHeadingPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNumberedSectionPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"1. Introduction", true},
|
||||
{"2.3 Methodology", true}, // multi-level numerals no longer require trailing dot
|
||||
{"2.3. Methodology", true},
|
||||
{"2.2.1 用户与权限", true}, // three-level numbering, Chinese title, no trailing dot
|
||||
{"3.2.1 单机 Docker Compose", true},
|
||||
{"IV. Results", true},
|
||||
{"1.Introduction", false}, // requires whitespace
|
||||
{"1.", false}, // requires title
|
||||
{"1.1", false}, // requires title after numeral
|
||||
{"1 NoDotSingleLevel", false}, // single-level numerals still need trailing dot
|
||||
{"1.2.3.4.5 TooDeep", false}, // more than 3 sub-levels not accepted
|
||||
{"plain text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := NumberedSectionPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("NumberedSectionPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGermanChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Kapitel 1: Einführung", true},
|
||||
{"Abschnitt 2.3 Methodik", true},
|
||||
{"Abschnitt 3 Methodik", true},
|
||||
{"Teil II Ergebnisse", true},
|
||||
{"chapter 1", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := GermanChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("GermanChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnglishChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Chapter 1: Intro", true},
|
||||
{"Section 5 Methods", true},
|
||||
{"Part IV Results", true},
|
||||
{"Kapitel 1", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := EnglishChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("EnglishChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChineseChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"第一章 引言", true},
|
||||
{"第3节 方法论", true},
|
||||
{"第二部分 结果", true},
|
||||
{"第 1 章 引言", true}, // space between 第 / numeral / unit
|
||||
{"第 一 章 引言", true},
|
||||
{"第1章 引言", true},
|
||||
{"Chapter 1", false},
|
||||
{"第 章 空数字", false}, // missing numeral
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := ChineseChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("ChineseChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVisualSeparatorPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"---", true},
|
||||
{"========", true},
|
||||
{"***", true},
|
||||
{"____", true},
|
||||
{"--", false}, // needs at least 3
|
||||
{"-- text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := VisualSeparatorPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("VisualSeparatorPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPageFooterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Seite 3 von 24", true},
|
||||
{"Page 5 of 12", true},
|
||||
{"page 7", true},
|
||||
{"Seite 9", true},
|
||||
{"页 3", true},
|
||||
{"页码 3 / 12", true},
|
||||
{"Some text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := PageFooterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("PageFooterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllCapsHeadingPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"INTRODUCTION", true},
|
||||
{"METHODS AND MATERIALS", true},
|
||||
{"Mixed Case Heading", false},
|
||||
{"ABC", false}, // <4 chars
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := AllCapsHeadingPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("AllCapsHeadingPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSentenceSeparators(t *testing.T) {
|
||||
if got := SentenceSeparators(LangChinese); got[0] != "。" {
|
||||
t.Errorf("Chinese should start with 。, got %v", got)
|
||||
}
|
||||
if got := SentenceSeparators(LangEnglish); got[0] != ". " {
|
||||
t.Errorf("English should start with '. ', got %v", got)
|
||||
}
|
||||
if got := SentenceSeparators("xx"); len(got) < 5 {
|
||||
t.Errorf("Unknown lang should return mixed (>=5 separators), got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChapterPatternsForLangs(t *testing.T) {
|
||||
if got := ChapterPatternsForLangs(nil); len(got) != 3 {
|
||||
t.Errorf("nil langs should return all 3, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{LangGerman}); len(got) != 1 {
|
||||
t.Errorf("only DE requested should return 1, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{LangGerman, LangChinese}); len(got) != 2 {
|
||||
t.Errorf("DE+ZH requested should return 2, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{"xx"}); len(got) != 3 {
|
||||
t.Errorf("unknown lang should fall back to all 3, got %d", len(got))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
// Package chunker - profiler.go scans a document once to gather structure
|
||||
// indicators that drive strategy selection (heading-aware vs. heuristic vs.
|
||||
// recursive). Profiling is cheap (a few regex passes plus rune counting)
|
||||
// and runs before any chunking decision is made.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// DocProfile holds the document-level signals used to choose a chunking tier.
|
||||
//
|
||||
// The JSON shape (snake_case via struct tags) is part of the public preview
|
||||
// endpoint API — keep field names stable. Internal callers should use the Go
|
||||
// field names; only the preview handler relies on the wire format.
|
||||
type DocProfile struct {
|
||||
TotalChars int `json:"total_chars"`
|
||||
TotalLines int `json:"total_lines"`
|
||||
AvgLineLen float64 `json:"avg_line_len"`
|
||||
StdLineLen float64 `json:"std_line_len"`
|
||||
|
||||
// Markdown structure
|
||||
MdHeadingCounts map[int]int `json:"md_heading_counts"` // level (1..6) → count
|
||||
MdHeadingTotal int `json:"md_heading_total"`
|
||||
|
||||
// Heuristic indicators
|
||||
NumberedSectionCount int `json:"numbered_section_count"`
|
||||
AllCapsShortLineCount int `json:"all_caps_short_line_count"`
|
||||
BlankParagraphBreaks int `json:"blank_paragraph_breaks"`
|
||||
FormFeedCount int `json:"form_feed_count"`
|
||||
VisualSepCount int `json:"visual_sep_count"`
|
||||
GermanChapterCount int `json:"german_chapter_count"`
|
||||
EnglishChapterCount int `json:"english_chapter_count"`
|
||||
ChineseChapterCount int `json:"chinese_chapter_count"`
|
||||
RepeatedFooterCount int `json:"repeated_footer_count"`
|
||||
|
||||
// Content characteristics
|
||||
HasTables bool `json:"has_tables"`
|
||||
HasCode bool `json:"has_code"`
|
||||
CodeRatio float64 `json:"code_ratio"`
|
||||
|
||||
// Detected language hints (best-effort)
|
||||
DetectedLangs []string `json:"detected_langs"`
|
||||
}
|
||||
|
||||
// HeadingDensity returns the share of lines that are Markdown headings.
|
||||
func (p *DocProfile) HeadingDensity() float64 {
|
||||
if p.TotalLines == 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(p.MdHeadingTotal) / float64(p.TotalLines)
|
||||
}
|
||||
|
||||
// DominantHeadingLevel returns the heading level (1..6) that should drive
|
||||
// section splitting. Preference order:
|
||||
// 1. The lowest level (closest to root) that has at least 3 occurrences —
|
||||
// a "real" structural backbone of the document.
|
||||
// 2. Otherwise the deepest level present at least once — gives finer-grained
|
||||
// boundaries for small documents that just have an H1 + a few H2s.
|
||||
//
|
||||
// Returns 0 when no Markdown headings exist.
|
||||
func (p *DocProfile) DominantHeadingLevel() int {
|
||||
if p.MdHeadingTotal == 0 {
|
||||
return 0
|
||||
}
|
||||
for level := 1; level <= 6; level++ {
|
||||
if p.MdHeadingCounts[level] >= 3 {
|
||||
return level
|
||||
}
|
||||
}
|
||||
for level := 6; level >= 1; level-- {
|
||||
if p.MdHeadingCounts[level] > 0 {
|
||||
return level
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// HeuristicMarkerTotal sums the non-Markdown structural markers.
|
||||
func (p *DocProfile) HeuristicMarkerTotal() int {
|
||||
return p.NumberedSectionCount +
|
||||
p.GermanChapterCount + p.EnglishChapterCount + p.ChineseChapterCount +
|
||||
p.AllCapsShortLineCount + p.VisualSepCount + p.FormFeedCount
|
||||
}
|
||||
|
||||
// ProfileDocument runs a single pass over text and returns its profile.
|
||||
func ProfileDocument(text string) *DocProfile {
|
||||
p := &DocProfile{
|
||||
MdHeadingCounts: make(map[int]int),
|
||||
}
|
||||
if text == "" {
|
||||
return p
|
||||
}
|
||||
|
||||
p.TotalChars = len([]rune(text))
|
||||
p.FormFeedCount = strings.Count(text, "\f")
|
||||
|
||||
lines := strings.Split(text, "\n")
|
||||
p.TotalLines = len(lines)
|
||||
|
||||
// First pass: per-line markers and length stats
|
||||
var lengths []float64
|
||||
inFence := false
|
||||
codeChars := 0
|
||||
for _, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
|
||||
// Toggle fenced-code state. We use a 3-backtick prefix detector here
|
||||
// rather than a full regex so we don't have to fight with the
|
||||
// protected-pattern logic later.
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
p.HasCode = true
|
||||
continue
|
||||
}
|
||||
if inFence {
|
||||
codeChars += len([]rune(line))
|
||||
continue
|
||||
}
|
||||
|
||||
runeLen := len([]rune(line))
|
||||
lengths = append(lengths, float64(runeLen))
|
||||
|
||||
if matchHeading(line, &p.MdHeadingCounts) {
|
||||
p.MdHeadingTotal++
|
||||
continue
|
||||
}
|
||||
if NumberedSectionPattern.MatchString(line) {
|
||||
p.NumberedSectionCount++
|
||||
}
|
||||
if GermanChapterPattern.MatchString(line) {
|
||||
p.GermanChapterCount++
|
||||
}
|
||||
if EnglishChapterPattern.MatchString(line) {
|
||||
p.EnglishChapterCount++
|
||||
}
|
||||
if ChineseChapterPattern.MatchString(line) {
|
||||
p.ChineseChapterCount++
|
||||
}
|
||||
if AllCapsHeadingPattern.MatchString(line) {
|
||||
p.AllCapsShortLineCount++
|
||||
}
|
||||
if VisualSeparatorPattern.MatchString(line) {
|
||||
p.VisualSepCount++
|
||||
}
|
||||
if PageFooterPattern.MatchString(line) {
|
||||
p.RepeatedFooterCount++
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "|") && strings.HasSuffix(trimmed, "|") {
|
||||
p.HasTables = true
|
||||
}
|
||||
}
|
||||
|
||||
if len(lengths) > 0 {
|
||||
var sum float64
|
||||
for _, l := range lengths {
|
||||
sum += l
|
||||
}
|
||||
p.AvgLineLen = sum / float64(len(lengths))
|
||||
var variance float64
|
||||
for _, l := range lengths {
|
||||
d := l - p.AvgLineLen
|
||||
variance += d * d
|
||||
}
|
||||
variance /= float64(len(lengths))
|
||||
p.StdLineLen = math.Sqrt(variance)
|
||||
}
|
||||
|
||||
if p.TotalChars > 0 {
|
||||
p.CodeRatio = float64(codeChars) / float64(p.TotalChars)
|
||||
}
|
||||
|
||||
p.BlankParagraphBreaks = strings.Count(text, "\n\n\n")
|
||||
|
||||
// Sample a slice of the document for language detection — avoids paying
|
||||
// O(N) scan cost on huge inputs while still giving a stable signal.
|
||||
sample := text
|
||||
if len(sample) > 4096 {
|
||||
sample = sample[:4096]
|
||||
}
|
||||
lang := DetectLanguage(sample)
|
||||
p.DetectedLangs = []string{lang}
|
||||
if lang == LangMixed {
|
||||
// Provide all three for downstream pattern selection.
|
||||
p.DetectedLangs = []string{LangEnglish, LangGerman, LangChinese}
|
||||
}
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
// matchHeading checks whether line is an ATX heading and increments the
|
||||
// appropriate level counter when so. Returns true on match.
|
||||
func matchHeading(line string, counts *map[int]int) bool {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
return false
|
||||
}
|
||||
level := len(m[1])
|
||||
if level < 1 || level > 6 {
|
||||
return false
|
||||
}
|
||||
(*counts)[level]++
|
||||
return true
|
||||
}
|
||||
|
||||
// StrategyTier identifies which chunking implementation should run.
|
||||
type StrategyTier string
|
||||
|
||||
const (
|
||||
TierHeading StrategyTier = "heading"
|
||||
TierHeuristic StrategyTier = "heuristic"
|
||||
TierLegacy StrategyTier = "legacy"
|
||||
)
|
||||
|
||||
// SelectStrategy returns the ordered tier chain to attempt for this document.
|
||||
// The first tier is the primary choice; subsequent tiers are fallbacks if
|
||||
// validation rejects the previous output. The "legacy" tier is appended as
|
||||
// a final safety net so callers always receive at least one chunk-set.
|
||||
func SelectStrategy(p *DocProfile) []StrategyTier {
|
||||
if p == nil {
|
||||
return []StrategyTier{TierLegacy}
|
||||
}
|
||||
var chain []StrategyTier
|
||||
|
||||
// Tier 1 candidate: Markdown heading-aware
|
||||
if p.MdHeadingTotal >= 3 && p.HeadingDensity() > 0.005 && p.DominantHeadingLevel() > 0 {
|
||||
chain = append(chain, TierHeading)
|
||||
}
|
||||
|
||||
// Tier 2 candidate: heuristic boundary detection
|
||||
if p.HeuristicMarkerTotal() >= 5 || p.FormFeedCount > 0 ||
|
||||
p.GermanChapterCount+p.EnglishChapterCount+p.ChineseChapterCount > 0 {
|
||||
chain = append(chain, TierHeuristic)
|
||||
}
|
||||
|
||||
// Legacy is the ultimate fallback: always returns chunks even when
|
||||
// validation fails, so callers never get an empty result.
|
||||
chain = append(chain, TierLegacy)
|
||||
return chain
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestProfileDocument_Empty(t *testing.T) {
|
||||
p := ProfileDocument("")
|
||||
if p.TotalChars != 0 || p.TotalLines != 0 {
|
||||
t.Errorf("empty doc should have zero stats, got %+v", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_MarkdownHeadings(t *testing.T) {
|
||||
doc := `# Title
|
||||
Some intro text here.
|
||||
|
||||
## Section 1
|
||||
Body of section 1.
|
||||
|
||||
## Section 2
|
||||
Body of section 2.
|
||||
|
||||
### Subsection 2.1
|
||||
Detail.
|
||||
|
||||
## Section 3
|
||||
More body.`
|
||||
p := ProfileDocument(doc)
|
||||
if p.MdHeadingCounts[1] != 1 {
|
||||
t.Errorf("expected 1 H1, got %d", p.MdHeadingCounts[1])
|
||||
}
|
||||
if p.MdHeadingCounts[2] != 3 {
|
||||
t.Errorf("expected 3 H2, got %d", p.MdHeadingCounts[2])
|
||||
}
|
||||
if p.MdHeadingCounts[3] != 1 {
|
||||
t.Errorf("expected 1 H3, got %d", p.MdHeadingCounts[3])
|
||||
}
|
||||
if p.MdHeadingTotal != 5 {
|
||||
t.Errorf("expected 5 headings total, got %d", p.MdHeadingTotal)
|
||||
}
|
||||
if p.DominantHeadingLevel() != 2 {
|
||||
t.Errorf("dominant level should be 2 (≥3 occurrences), got %d", p.DominantHeadingLevel())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DominantLevelFallback(t *testing.T) {
|
||||
// No level reaches 3 occurrences — should fall back to most frequent.
|
||||
doc := "# Single H1\n## H2 a\n## H2 b\n"
|
||||
p := ProfileDocument(doc)
|
||||
if p.DominantHeadingLevel() != 2 {
|
||||
t.Errorf("expected fallback to level 2 (most frequent), got %d", p.DominantHeadingLevel())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_NumberedSections(t *testing.T) {
|
||||
doc := `1. Introduction
|
||||
text
|
||||
|
||||
2. Methodology
|
||||
text
|
||||
|
||||
3. Results
|
||||
text`
|
||||
p := ProfileDocument(doc)
|
||||
if p.NumberedSectionCount < 3 {
|
||||
t.Errorf("expected ≥3 numbered sections, got %d", p.NumberedSectionCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_GermanChapters(t *testing.T) {
|
||||
doc := "Kapitel 1: Einführung\n\nText\n\nKapitel 2: Hauptteil\n\nText"
|
||||
p := ProfileDocument(doc)
|
||||
if p.GermanChapterCount != 2 {
|
||||
t.Errorf("expected 2 German chapters, got %d", p.GermanChapterCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_ChineseChapters(t *testing.T) {
|
||||
doc := "第一章 引言\n\n内容\n\n第二章 方法\n\n内容"
|
||||
p := ProfileDocument(doc)
|
||||
if p.ChineseChapterCount != 2 {
|
||||
t.Errorf("expected 2 Chinese chapters, got %d", p.ChineseChapterCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_FormFeed(t *testing.T) {
|
||||
doc := "page 1 content\f\npage 2 content\f\npage 3 content"
|
||||
p := ProfileDocument(doc)
|
||||
if p.FormFeedCount != 2 {
|
||||
t.Errorf("expected 2 form feeds, got %d", p.FormFeedCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DetectsCodeBlock(t *testing.T) {
|
||||
doc := "Some prose.\n\n```go\nfunc main() {}\n```\n\nMore prose."
|
||||
p := ProfileDocument(doc)
|
||||
if !p.HasCode {
|
||||
t.Error("expected HasCode=true for fenced block")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DetectsTable(t *testing.T) {
|
||||
doc := "Intro.\n\n| col a | col b |\n| --- | --- |\n| 1 | 2 |\n"
|
||||
p := ProfileDocument(doc)
|
||||
if !p.HasTables {
|
||||
t.Error("expected HasTables=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_LineStatistics(t *testing.T) {
|
||||
doc := "short\nthis is a longer line of text\nanother line here"
|
||||
p := ProfileDocument(doc)
|
||||
if p.TotalLines != 3 {
|
||||
t.Errorf("expected 3 lines, got %d", p.TotalLines)
|
||||
}
|
||||
if p.AvgLineLen <= 0 {
|
||||
t.Error("expected positive avg line len")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_HeadingDoc(t *testing.T) {
|
||||
doc := "# A\nbody\n## B\nbody\n## C\nbody\n## D\nbody"
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[0] != TierHeading {
|
||||
t.Errorf("expected heading tier first, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_HeuristicDoc(t *testing.T) {
|
||||
doc := strings.Repeat("Kapitel 1: Foo\nbody body body\n\n", 1) +
|
||||
strings.Repeat("Kapitel 2: Bar\nbody body body\n\n", 1)
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
// no markdown headings → heuristic must come first (heading tier skipped)
|
||||
if chain[0] != TierHeuristic {
|
||||
t.Errorf("expected heuristic tier first, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_PlainDoc(t *testing.T) {
|
||||
doc := "just a paragraph of plain text without any structure indicators at all here"
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[0] != TierLegacy {
|
||||
t.Errorf("expected legacy tier first for unstructured doc, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_AlwaysFallsBackToLegacy(t *testing.T) {
|
||||
for _, doc := range []string{"", "simple", "# H1\nbody"} {
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[len(chain)-1] != TierLegacy {
|
||||
t.Errorf("chain must end with legacy, got %v for doc=%q", chain, doc)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,914 @@
|
||||
// Package chunker implements text splitting for document chunking.
|
||||
//
|
||||
// Ported from the Python docreader/splitter/splitter.py recursive text splitter.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser"
|
||||
)
|
||||
|
||||
// Chunk represents a piece of split text with position tracking.
|
||||
//
|
||||
// Content holds exactly the text from the original document between Start
|
||||
// and End (rune offsets), so End-Start == utf8.RuneCountInString(Content).
|
||||
// This invariant is relied on by document-reconstruction code paths
|
||||
// (knowledge.go:2278+ for summary generation, UI highlighting, etc.).
|
||||
//
|
||||
// ContextHeader is a separately-tracked context string (e.g. a Markdown
|
||||
// heading breadcrumb) that should be prepended at embedding/retrieval time
|
||||
// but is NOT part of Content. Keeping the two apart preserves the
|
||||
// position invariant while still letting embedding pipelines see the
|
||||
// section context.
|
||||
type Chunk struct {
|
||||
Content string
|
||||
ContextHeader string
|
||||
Seq int
|
||||
Start int
|
||||
End int
|
||||
}
|
||||
|
||||
// EmbeddingContent returns the text that should be fed to the embedding
|
||||
// model — the ContextHeader prepended (when set) plus the chunk content.
|
||||
// Use this where Content alone would lose semantic context (Tier-1 chunks).
|
||||
//
|
||||
// Content is returned verbatim from the source document (the End-Start
|
||||
// rune-count invariant requires that), but for embedding we trim the
|
||||
// surrounding whitespace so leading/trailing newlines from boundary slices
|
||||
// don't dilute the embedded vector or waste tokens. Inner whitespace is
|
||||
// preserved.
|
||||
func (c Chunk) EmbeddingContent() string {
|
||||
body := strings.TrimSpace(c.Content)
|
||||
if c.ContextHeader == "" {
|
||||
return body
|
||||
}
|
||||
return c.ContextHeader + "\n\n" + body
|
||||
}
|
||||
|
||||
// ImageRef is an image reference found within a chunk's content.
|
||||
type ImageRef struct {
|
||||
OriginalRef string
|
||||
AltText string
|
||||
Start int // offset within the chunk content
|
||||
End int
|
||||
}
|
||||
|
||||
// SplitterConfig configures the text splitter. Strategy and TokenLimit are
|
||||
// honored by the strategy entry point in strategy.go; the legacy SplitText
|
||||
// path uses only ChunkSize/Overlap/Separators.
|
||||
type SplitterConfig struct {
|
||||
ChunkSize int
|
||||
ChunkOverlap int
|
||||
Separators []string
|
||||
|
||||
// Strategy selects an adaptive tier. Empty = legacy (backwards-compatible).
|
||||
// See strategy.go for valid values.
|
||||
Strategy string
|
||||
// TokenLimit caps chunk size in approximate tokens. 0 = use ChunkSize chars.
|
||||
TokenLimit int
|
||||
// Languages hints multilingual heuristic patterns. Empty = auto-detect.
|
||||
Languages []string
|
||||
}
|
||||
|
||||
// Default chunk sizing constants. Single source of truth for the entire
|
||||
// chunker package and (via knowledge.go::buildSplitterConfig) the
|
||||
// knowledge service. The frontend KnowledgeBaseEditorModal mirrors these
|
||||
// numbers in its initial form state — keep them in sync if you change
|
||||
// either value here.
|
||||
//
|
||||
// DefaultChunkSize = 512 chars: ~100–130 English tokens / ~300 Chinese
|
||||
// tokens. Validated as a strong baseline by the Vecta Feb-2026 benchmark
|
||||
// across 50 academic papers. Use 200–400 for FAQ-style atomic content,
|
||||
// 1000–2000 for narrative / argumentative documents.
|
||||
//
|
||||
// DefaultChunkOverlap = 80 chars (≈15% of DefaultChunkSize): community-
|
||||
// recommended sweet spot between recall (an answer split across a
|
||||
// boundary needs overlap to be retrievable) and storage cost. Use 0 for
|
||||
// strictly atomic data (FAQ, JSON records), 150–200 for long narratives
|
||||
// where reasoning crosses chunks.
|
||||
//
|
||||
// MIGRATION NOTE: Prior versions had three different overlap defaults
|
||||
// (Go DefaultConfig: 64, knowledge.go buildSplitterConfig: 50, Python
|
||||
// docreader: 100). All consolidated to 80 here.
|
||||
//
|
||||
// Existing knowledge bases that stored ChunkOverlap=0 in the DB pick
|
||||
// this 80 up on next re-index; their previously-indexed embeddings will
|
||||
// not match new ones bit-for-bit. Recall stays similar but search
|
||||
// ranking can shift slightly. To freeze the old behavior on a per-KB
|
||||
// basis, explicitly set ChunkingConfig.ChunkOverlap to 64 before
|
||||
// re-indexing.
|
||||
const (
|
||||
DefaultChunkSize = 512
|
||||
DefaultChunkOverlap = 80
|
||||
)
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() SplitterConfig {
|
||||
return SplitterConfig{
|
||||
ChunkSize: DefaultChunkSize,
|
||||
ChunkOverlap: DefaultChunkOverlap,
|
||||
Separators: []string{"\n\n", "\n", "。"},
|
||||
}
|
||||
}
|
||||
|
||||
// protectedPatterns are regex patterns for content that must not be split.
|
||||
var protectedPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?s)\$\$.*?\$\$`), // LaTeX block math
|
||||
regexp.MustCompile(`!\[[^\]]*\]\([^)]+\)`), // Markdown images
|
||||
regexp.MustCompile(`\[[^\]]*\]\([^)]+\)`), // Markdown links
|
||||
regexp.MustCompile("(?m)[ ]*(?:\\|[^|\\n]*)+\\|[\\r\\n]+\\s*(?:\\|\\s*:?-{3,}:?\\s*)+\\|[\\r\\n]+"), // Table header+separator
|
||||
regexp.MustCompile("(?m)[ ]*(?:\\|[^|\\n]*)+\\|[\\r\\n]+"), // Table rows
|
||||
regexp.MustCompile("(?s)```(?:\\w+)?[\\r\\n].*?```"), // Fenced code blocks
|
||||
regexp.MustCompile("`[^`\\r\\n]+`"), // Markdown inline code
|
||||
}
|
||||
|
||||
type span struct {
|
||||
start, end int
|
||||
}
|
||||
|
||||
// protectedSpansRune converts byte-offset protected spans to rune offsets
|
||||
// in a single forward pass over text. Used by callers that work in rune
|
||||
// space (e.g. the heuristic splitter) to avoid choosing chunk boundaries
|
||||
// that cut through protected content. byteSpans must be sorted by start
|
||||
// (protectedSpans guarantees this).
|
||||
func protectedSpansRune(text string, byteSpans []span) []span {
|
||||
if len(byteSpans) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]span, 0, len(byteSpans))
|
||||
runeIdx := 0
|
||||
byteIdx := 0
|
||||
for _, s := range byteSpans {
|
||||
for byteIdx < s.start && byteIdx < len(text) {
|
||||
_, size := utf8.DecodeRuneInString(text[byteIdx:])
|
||||
byteIdx += size
|
||||
runeIdx++
|
||||
}
|
||||
startRune := runeIdx
|
||||
for byteIdx < s.end && byteIdx < len(text) {
|
||||
_, size := utf8.DecodeRuneInString(text[byteIdx:])
|
||||
byteIdx += size
|
||||
runeIdx++
|
||||
}
|
||||
out = append(out, span{start: startRune, end: runeIdx})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// protectedSpans finds all non-overlapping protected regions in text.
|
||||
func protectedSpans(text string) []span {
|
||||
type match struct {
|
||||
start, end int
|
||||
}
|
||||
var all []match
|
||||
for _, pat := range protectedPatterns {
|
||||
locs := pat.FindAllStringIndex(text, -1)
|
||||
for _, loc := range locs {
|
||||
if loc[1]-loc[0] > 0 {
|
||||
all = append(all, match{loc[0], loc[1]})
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(all) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Sort by start, then by length descending
|
||||
for i := 1; i < len(all); i++ {
|
||||
for j := i; j > 0; j-- {
|
||||
if all[j].start < all[j-1].start ||
|
||||
(all[j].start == all[j-1].start && (all[j].end-all[j].start) > (all[j-1].end-all[j-1].start)) {
|
||||
all[j], all[j-1] = all[j-1], all[j]
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove overlaps
|
||||
var result []span
|
||||
lastEnd := 0
|
||||
for _, m := range all {
|
||||
if m.start >= lastEnd {
|
||||
result = append(result, span{m.start, m.end})
|
||||
lastEnd = m.end
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// splitUnit is a piece of text with its original position.
|
||||
type splitUnit struct {
|
||||
text string
|
||||
start, end int
|
||||
}
|
||||
|
||||
// splitBySeparators splits text by separators in priority order, recursively
|
||||
// applying the next separator to any piece that is still larger than
|
||||
// chunkSize. Mirrors the recursive priority semantics of the Python
|
||||
// reference splitter (docreader/splitter/splitter.py:_split): if `\n\n`
|
||||
// produces a piece that's still too big, `\n` (and subsequent separators)
|
||||
// are applied within that piece — not to the whole text.
|
||||
//
|
||||
// chunkSize == 0 disables the recursion guard; callers that don't care
|
||||
// about size budget (e.g. a final mergeUnits-style pass) pass 0.
|
||||
func splitBySeparators(text string, separators []string, chunkSize int) []string {
|
||||
if text == "" || len(separators) == 0 {
|
||||
return []string{text}
|
||||
}
|
||||
if chunkSize > 0 && runeLen(text) <= chunkSize {
|
||||
return []string{text}
|
||||
}
|
||||
|
||||
for i, sep := range separators {
|
||||
if sep == "" {
|
||||
continue
|
||||
}
|
||||
re := regexp.MustCompile("(" + regexp.QuoteMeta(sep) + ")")
|
||||
splits := re.Split(text, -1)
|
||||
matches := re.FindAllString(text, -1)
|
||||
if len(matches) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var pieces []string
|
||||
for j, s := range splits {
|
||||
if s != "" {
|
||||
pieces = append(pieces, s)
|
||||
}
|
||||
if j < len(matches) && matches[j] != "" {
|
||||
pieces = append(pieces, matches[j])
|
||||
}
|
||||
}
|
||||
if len(pieces) <= 1 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Recursively split any piece that is still too large with the
|
||||
// remaining (lower-priority) separators.
|
||||
var out []string
|
||||
remaining := separators[i+1:]
|
||||
for _, p := range pieces {
|
||||
if chunkSize > 0 && runeLen(p) > chunkSize && len(remaining) > 0 {
|
||||
out = append(out, splitBySeparators(p, remaining, chunkSize)...)
|
||||
} else {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
return []string{text}
|
||||
}
|
||||
|
||||
// runeLen returns the number of runes in s.
|
||||
func runeLen(s string) int {
|
||||
return utf8.RuneCountInString(s)
|
||||
}
|
||||
|
||||
// SplitText splits text into chunks with overlap, respecting protected patterns.
|
||||
func SplitText(text string, cfg SplitterConfig) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
chunkSize := cfg.ChunkSize
|
||||
chunkOverlap := cfg.ChunkOverlap
|
||||
separators := cfg.Separators
|
||||
|
||||
if chunkSize <= 0 {
|
||||
chunkSize = 512
|
||||
}
|
||||
if chunkOverlap < 0 {
|
||||
chunkOverlap = 0
|
||||
}
|
||||
|
||||
// Step 1: Find protected spans
|
||||
protected := protectedSpans(text)
|
||||
|
||||
// Step 2: Split non-protected regions by separators, keep protected as atomic units.
|
||||
// chunkSize is forwarded so splitBySeparators can recursively apply lower-priority
|
||||
// separators to oversize pieces (Python-parity recursive split).
|
||||
units := buildUnitsWithProtection(text, protected, separators, chunkSize)
|
||||
|
||||
// Step 3: Merge units into chunks with overlap
|
||||
return mergeUnits(units, chunkSize, chunkOverlap)
|
||||
}
|
||||
|
||||
// buildUnitsWithProtection splits text into units, preserving protected spans as atomic.
|
||||
// Start/End positions in the returned units are rune offsets (not byte offsets),
|
||||
// because downstream merge logic indexes content via []rune slicing.
|
||||
// If a protected span exceeds maxProtectedSize, it will be forcibly split to prevent
|
||||
// creating chunks that are too large for downstream processing (e.g., embedding APIs).
|
||||
// chunkSize is forwarded to splitBySeparators so recursive splitting can keep pieces
|
||||
// under the budget when one separator alone leaves a piece oversize.
|
||||
func buildUnitsWithProtection(text string, protected []span, separators []string, chunkSize int) []splitUnit {
|
||||
const maxProtectedSize = 7500 // Maximum size for a protected unit (留余量给标题等)
|
||||
|
||||
var units []splitUnit
|
||||
bytePos := 0
|
||||
runePos := 0
|
||||
|
||||
for _, p := range protected {
|
||||
if p.start > bytePos {
|
||||
pre := text[bytePos:p.start]
|
||||
parts := splitBySeparators(pre, separators, chunkSize)
|
||||
runeOffset := runePos
|
||||
for _, part := range parts {
|
||||
partRuneLen := runeLen(part)
|
||||
units = append(units, splitUnit{
|
||||
text: part,
|
||||
start: runeOffset,
|
||||
end: runeOffset + partRuneLen,
|
||||
})
|
||||
runeOffset += partRuneLen
|
||||
}
|
||||
runePos += runeLen(pre)
|
||||
}
|
||||
|
||||
protText := text[p.start:p.end]
|
||||
protRuneLen := runeLen(protText)
|
||||
|
||||
// If protected content is too large, forcibly split it
|
||||
if protRuneLen > maxProtectedSize {
|
||||
// Split into smaller chunks at line breaks or spaces
|
||||
runes := []rune(protText)
|
||||
offset := 0
|
||||
for offset < len(runes) {
|
||||
chunkEnd := offset + maxProtectedSize
|
||||
if chunkEnd > len(runes) {
|
||||
chunkEnd = len(runes)
|
||||
} else {
|
||||
// Try to break at a newline or space
|
||||
for i := chunkEnd - 1; i > offset && i > chunkEnd-200; i-- {
|
||||
if runes[i] == '\n' || runes[i] == ' ' {
|
||||
chunkEnd = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
chunkText := string(runes[offset:chunkEnd])
|
||||
chunkLen := chunkEnd - offset
|
||||
units = append(units, splitUnit{
|
||||
text: chunkText,
|
||||
start: runePos + offset,
|
||||
end: runePos + offset + chunkLen,
|
||||
})
|
||||
offset = chunkEnd
|
||||
}
|
||||
} else {
|
||||
// Normal case: keep protected content as a single unit
|
||||
units = append(units, splitUnit{
|
||||
text: protText,
|
||||
start: runePos,
|
||||
end: runePos + protRuneLen,
|
||||
})
|
||||
}
|
||||
runePos += protRuneLen
|
||||
bytePos = p.end
|
||||
}
|
||||
|
||||
if bytePos < len(text) {
|
||||
remaining := text[bytePos:]
|
||||
parts := splitBySeparators(remaining, separators, chunkSize)
|
||||
runeOffset := runePos
|
||||
for _, part := range parts {
|
||||
partRuneLen := runeLen(part)
|
||||
units = append(units, splitUnit{
|
||||
text: part,
|
||||
start: runeOffset,
|
||||
end: runeOffset + partRuneLen,
|
||||
})
|
||||
runeOffset += partRuneLen
|
||||
}
|
||||
}
|
||||
|
||||
return units
|
||||
}
|
||||
|
||||
// mergeUnits combines split units into chunks with overlap tracking.
|
||||
// Enforces an absolute maximum chunk size to prevent exceeding downstream limits (e.g., embedding APIs).
|
||||
// Active contextual headers (e.g., Markdown table headers) are prepended to new
|
||||
// chunks so that every chunk carries its own header context.
|
||||
func mergeUnits(units []splitUnit, chunkSize, chunkOverlap int) []Chunk {
|
||||
if len(units) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
const absoluteMaxSize = 7500
|
||||
|
||||
ht := newHeaderTracker()
|
||||
|
||||
var chunks []Chunk
|
||||
var current []splitUnit
|
||||
curLen := 0
|
||||
|
||||
for _, u := range units {
|
||||
uLen := runeLen(u.text)
|
||||
|
||||
// If this single unit exceeds absolute max, force split it further
|
||||
if uLen > absoluteMaxSize {
|
||||
// Flush current chunk if any
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
|
||||
// Update header state even for oversized units
|
||||
ht.update(u.text)
|
||||
|
||||
// Split this oversized unit into smaller chunks
|
||||
runes := []rune(u.text)
|
||||
offset := 0
|
||||
for offset < len(runes) {
|
||||
chunkEnd := offset + absoluteMaxSize
|
||||
if chunkEnd > len(runes) {
|
||||
chunkEnd = len(runes)
|
||||
} else {
|
||||
for i := chunkEnd - 1; i > offset && i > chunkEnd-200; i-- {
|
||||
if runes[i] == '\n' || runes[i] == ' ' {
|
||||
chunkEnd = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
chunkText := string(runes[offset:chunkEnd])
|
||||
chunks = append(chunks, Chunk{
|
||||
Content: chunkText,
|
||||
Seq: len(chunks),
|
||||
Start: u.start + offset,
|
||||
End: u.start + chunkEnd,
|
||||
})
|
||||
offset = chunkEnd
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Update header tracking
|
||||
ht.update(u.text)
|
||||
// Flush at table boundary so the next table is not merged into a chunk
|
||||
// that still carries the previous table's prepended header context.
|
||||
if ht.headerEndedThisUnit && len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
headers := ht.getHeaders()
|
||||
headersLen := runeLen(headers)
|
||||
if headersLen > chunkSize {
|
||||
headers = ""
|
||||
headersLen = 0
|
||||
}
|
||||
|
||||
// If adding this unit (plus reserving space for headers in a potential
|
||||
// next chunk) would exceed chunk size, flush the current chunk.
|
||||
if curLen+uLen+headersLen > chunkSize && len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
|
||||
// Keep overlap from the end of current
|
||||
current, curLen = computeOverlap(current, chunkOverlap, chunkSize, uLen)
|
||||
|
||||
// Shrink overlap further if needed to fit headers + next unit
|
||||
if headers != "" && headersLen+uLen <= chunkSize {
|
||||
for len(current) > 0 && curLen+uLen+headersLen > chunkSize {
|
||||
curLen -= runeLen(current[0].text)
|
||||
current = current[1:]
|
||||
}
|
||||
|
||||
// Prepend headers if the column-name context is not already present
|
||||
// in the overlap or the next unit being added.
|
||||
overlapText := unitsText(current)
|
||||
if !headerAlreadyPresent(headers, overlapText, u.text) &&
|
||||
!headerColumnMismatch(headers, u.text) {
|
||||
startPos := u.start
|
||||
if len(current) > 0 {
|
||||
startPos = current[0].start
|
||||
}
|
||||
hUnit := splitUnit{text: headers, start: startPos, end: startPos}
|
||||
current = append([]splitUnit{hUnit}, current...)
|
||||
curLen += headersLen
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if adding this unit would exceed absolute max
|
||||
if curLen+uLen > absoluteMaxSize {
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
}
|
||||
|
||||
current = append(current, u)
|
||||
curLen += uLen
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// unitsText concatenates the text of all units.
|
||||
func unitsText(units []splitUnit) string {
|
||||
var sb strings.Builder
|
||||
for _, u := range units {
|
||||
sb.WriteString(u.text)
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// headerAlreadyPresent returns true if the column-name row from the header
|
||||
// is already present in the overlap or the next unit, preventing duplication.
|
||||
func headerAlreadyPresent(headers, overlapText, unitText string) bool {
|
||||
// Fast path: full header already in overlap or unit
|
||||
if strings.Contains(overlapText, headers) || strings.Contains(unitText, headers) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Extract the column-name row (first meaningful non-separator line).
|
||||
// For a rewritten header like "| col1 | col2 |\n| --- | --- |\n",
|
||||
// the first line is the column names.
|
||||
colRow := headerColumnRow(headers)
|
||||
if colRow == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
return strings.Contains(overlapText, colRow) || strings.Contains(unitText, colRow)
|
||||
}
|
||||
|
||||
// headerColumnRow extracts the column-name line from a header string.
|
||||
// Returns empty string if the header has no meaningful column names.
|
||||
func headerColumnRow(header string) string {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.Contains(line, "---") {
|
||||
continue
|
||||
}
|
||||
// Skip lines that are only pipes/whitespace (empty header rows)
|
||||
onlyPipes := true
|
||||
for _, r := range line {
|
||||
if r != '|' && r != ' ' && r != '\t' {
|
||||
onlyPipes = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if !onlyPipes {
|
||||
return line
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func buildChunk(units []splitUnit, seq int) Chunk {
|
||||
var sb strings.Builder
|
||||
for _, u := range units {
|
||||
sb.WriteString(u.text)
|
||||
}
|
||||
return Chunk{
|
||||
Content: sb.String(),
|
||||
Seq: seq,
|
||||
Start: units[0].start,
|
||||
End: units[len(units)-1].end,
|
||||
}
|
||||
}
|
||||
|
||||
// computeOverlap returns the semantic suffix to keep for overlap and its total
|
||||
// rune length.
|
||||
//
|
||||
// The configured overlap is a hard upper bound, not a raw character slice.
|
||||
// Boundary detection may inspect up to four additional runes immediately
|
||||
// before that tail window so a separator cut by the window boundary remains
|
||||
// visible. A candidate is eligible only when its last rune is at position -1
|
||||
// relative to the original window, or later, so retained content never exceeds
|
||||
// the overlap limit. Eligible boundaries use this priority:
|
||||
//
|
||||
// 1. paragraph break (\n\n / \r\n\r\n)
|
||||
// 2. line break (\n / \r\n)
|
||||
// 3. sentence end (。?!, ". ", "? ", "! ")
|
||||
//
|
||||
// Priority wins first; for boundaries with the same priority, the earliest
|
||||
// one in the window wins so the useful overlap is as large as possible. The
|
||||
// next chunk starts immediately after the selected separator. If the window
|
||||
// has no valid semantic boundary, no overlap is retained. This keeps overlap
|
||||
// deterministic and bounded while allowing a boundary inside a large split
|
||||
// unit (the previous implementation could only retain whole units and often
|
||||
// collapsed to zero overlap for ordinary paragraphs).
|
||||
const semanticOverlapLookbehind = 4 // rune length of the longest separator: \r\n\r\n
|
||||
|
||||
func computeOverlap(current []splitUnit, chunkOverlap, chunkSize, nextLen int) ([]splitUnit, int) {
|
||||
if chunkOverlap <= 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
// Overlap is part of the next chunk's total size. When the incoming unit
|
||||
// is already close to the chunk budget, shrink the search window rather
|
||||
// than creating an oversized chunk.
|
||||
maxOverlap := chunkOverlap
|
||||
if remaining := chunkSize - nextLen; remaining < maxOverlap {
|
||||
maxOverlap = remaining
|
||||
}
|
||||
if maxOverlap <= 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
window := semanticOverlapWindow(current, maxOverlap+semanticOverlapLookbehind)
|
||||
if len(window) == 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
windowText := unitsText(window)
|
||||
// Coordinates before originalWindowStart are the lookbehind region. With
|
||||
// end-exclusive offsets, requiring boundaryEnd >= originalWindowStart is
|
||||
// equivalent to requiring the separator's final rune to be at -1 or later.
|
||||
originalWindowStart := runeLen(windowText) - maxOverlap
|
||||
if originalWindowStart < 0 {
|
||||
originalWindowStart = 0
|
||||
}
|
||||
boundaryEnd, ok := findSemanticOverlapBoundaryEndingAtOrAfter(windowText, originalWindowStart)
|
||||
if !ok {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
overlap := trimUnitsPrefix(window, boundaryEnd)
|
||||
overlapLen := 0
|
||||
for _, u := range overlap {
|
||||
overlapLen += runeLen(u.text)
|
||||
}
|
||||
if overlapLen <= 0 || overlapLen > maxOverlap || strings.TrimSpace(unitsText(overlap)) == "" {
|
||||
return nil, 0
|
||||
}
|
||||
return overlap, overlapLen
|
||||
}
|
||||
|
||||
// semanticOverlapWindow returns at most maxLen source-backed runes from the
|
||||
// tail of current. It may slice inside the first retained splitUnit so a
|
||||
// semantic boundary inside a large paragraph remains discoverable. Synthetic
|
||||
// zero-width units (for example repeated table headers) form a hard barrier:
|
||||
// crossing them would break the Start/End-to-Content source invariant.
|
||||
func semanticOverlapWindow(current []splitUnit, maxLen int) []splitUnit {
|
||||
if maxLen <= 0 || len(current) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
remaining := maxLen
|
||||
reversed := make([]splitUnit, 0, len(current))
|
||||
for i := len(current) - 1; i >= 0 && remaining > 0; i-- {
|
||||
u := current[i]
|
||||
uLen := runeLen(u.text)
|
||||
if uLen == 0 {
|
||||
continue
|
||||
}
|
||||
// Header markers contain generated text but occupy no source range.
|
||||
// Do not include or cross them when deriving an overlap suffix.
|
||||
if u.start == u.end || u.end-u.start != uLen {
|
||||
break
|
||||
}
|
||||
|
||||
if uLen <= remaining {
|
||||
reversed = append(reversed, u)
|
||||
remaining -= uLen
|
||||
continue
|
||||
}
|
||||
|
||||
runes := []rune(u.text)
|
||||
start := uLen - remaining
|
||||
reversed = append(reversed, splitUnit{
|
||||
text: string(runes[start:]),
|
||||
start: u.start + start,
|
||||
end: u.end,
|
||||
})
|
||||
remaining = 0
|
||||
}
|
||||
|
||||
if len(reversed) == 0 {
|
||||
return nil
|
||||
}
|
||||
window := make([]splitUnit, len(reversed))
|
||||
for i := range reversed {
|
||||
window[len(reversed)-1-i] = reversed[i]
|
||||
}
|
||||
return window
|
||||
}
|
||||
|
||||
type semanticOverlapBoundary struct {
|
||||
start int
|
||||
end int
|
||||
priority int
|
||||
}
|
||||
|
||||
// findSemanticOverlapBoundary returns the rune offset immediately after the
|
||||
// selected separator. Boundaries inside protected Markdown/code/math regions
|
||||
// are ignored, and a boundary is invalid when only whitespace follows it.
|
||||
func findSemanticOverlapBoundary(text string) (int, bool) {
|
||||
return findSemanticOverlapBoundaryEndingAtOrAfter(text, 0)
|
||||
}
|
||||
|
||||
// findSemanticOverlapBoundaryEndingAtOrAfter applies the normal priority and
|
||||
// earliest-position rules only to candidates whose end-exclusive rune offset
|
||||
// is at least minEnd. Filtering before comparison prevents an earlier but
|
||||
// ineligible lookbehind separator from hiding a later eligible boundary.
|
||||
func findSemanticOverlapBoundaryEndingAtOrAfter(text string, minEnd int) (int, bool) {
|
||||
runes := []rune(text)
|
||||
if len(runes) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
if minEnd < 0 {
|
||||
minEnd = 0
|
||||
}
|
||||
if minEnd > len(runes) {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
protected := protectedSpansRune(text, protectedSpans(text))
|
||||
insideProtected := func(pos int) bool {
|
||||
for _, p := range protected {
|
||||
if pos < p.start {
|
||||
return false
|
||||
}
|
||||
if pos >= p.start && pos < p.end {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
hasMeaningfulTail := func(end int) bool {
|
||||
return end >= 0 && end < len(runes) && strings.TrimSpace(string(runes[end:])) != ""
|
||||
}
|
||||
|
||||
var best semanticOverlapBoundary
|
||||
found := false
|
||||
consider := func(start, end, priority int) {
|
||||
if start < 0 || end <= start || end < minEnd || end > len(runes) ||
|
||||
insideProtected(start) || !hasMeaningfulTail(end) {
|
||||
return
|
||||
}
|
||||
candidate := semanticOverlapBoundary{start: start, end: end, priority: priority}
|
||||
if !found || candidate.priority < best.priority ||
|
||||
(candidate.priority == best.priority && candidate.start < best.start) {
|
||||
best = candidate
|
||||
found = true
|
||||
}
|
||||
}
|
||||
|
||||
// Mark paragraph-break runes so their component newlines are not also
|
||||
// emitted as lower-priority line-break candidates.
|
||||
paragraphRune := make([]bool, len(runes))
|
||||
for i := 0; i < len(runes); i++ {
|
||||
switch {
|
||||
case i+3 < len(runes) && runes[i] == '\r' && runes[i+1] == '\n' &&
|
||||
runes[i+2] == '\r' && runes[i+3] == '\n':
|
||||
consider(i, i+4, 1)
|
||||
for j := i; j < i+4; j++ {
|
||||
paragraphRune[j] = true
|
||||
}
|
||||
i += 3
|
||||
case i+1 < len(runes) && runes[i] == '\n' && runes[i+1] == '\n':
|
||||
consider(i, i+2, 1)
|
||||
paragraphRune[i], paragraphRune[i+1] = true, true
|
||||
i++
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(runes); i++ {
|
||||
if paragraphRune[i] {
|
||||
continue
|
||||
}
|
||||
if runes[i] == '\r' && i+1 < len(runes) && runes[i+1] == '\n' && !paragraphRune[i+1] {
|
||||
consider(i, i+2, 2)
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if runes[i] == '\n' {
|
||||
consider(i, i+1, 2)
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(runes); i++ {
|
||||
switch runes[i] {
|
||||
case '。', '?', '!':
|
||||
consider(i, i+1, 3)
|
||||
case '.', '?', '!':
|
||||
if i+1 < len(runes) && runes[i+1] == ' ' {
|
||||
consider(i, i+2, 3)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !found {
|
||||
return 0, false
|
||||
}
|
||||
return best.end, true
|
||||
}
|
||||
|
||||
// trimUnitsPrefix removes prefixLen source runes while preserving the
|
||||
// remaining units' original positions. prefixLen is relative to unitsText.
|
||||
func trimUnitsPrefix(units []splitUnit, prefixLen int) []splitUnit {
|
||||
if prefixLen <= 0 {
|
||||
out := make([]splitUnit, len(units))
|
||||
copy(out, units)
|
||||
return out
|
||||
}
|
||||
|
||||
remaining := prefixLen
|
||||
out := make([]splitUnit, 0, len(units))
|
||||
for _, u := range units {
|
||||
uLen := runeLen(u.text)
|
||||
if remaining >= uLen {
|
||||
remaining -= uLen
|
||||
continue
|
||||
}
|
||||
if remaining > 0 {
|
||||
runes := []rune(u.text)
|
||||
u.text = string(runes[remaining:])
|
||||
u.start += remaining
|
||||
remaining = 0
|
||||
}
|
||||
out = append(out, u)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ParentChildResult holds the two-level chunking output.
|
||||
// Parent chunks provide context (large window), child chunks are used for
|
||||
// embedding/retrieval (small window). Each child carries its ParentIndex so
|
||||
// the caller can wire up ParentChunkID after DB insertion.
|
||||
type ParentChildResult struct {
|
||||
Parents []Chunk
|
||||
Children []ChildChunk
|
||||
}
|
||||
|
||||
// ChildChunk extends Chunk with a reference to its parent.
|
||||
type ChildChunk struct {
|
||||
Chunk
|
||||
ParentIndex int // index into ParentChildResult.Parents
|
||||
}
|
||||
|
||||
// SplitTextParentChild performs two-level chunking:
|
||||
// 1. Split text into large parent chunks (parentCfg).
|
||||
// 2. Split each parent into smaller child chunks (childCfg) for embedding.
|
||||
//
|
||||
// The child Seq is globally unique across the entire document.
|
||||
func SplitTextParentChild(text string, parentCfg, childCfg SplitterConfig) ParentChildResult {
|
||||
parents := SplitText(text, parentCfg)
|
||||
if len(parents) == 0 {
|
||||
return ParentChildResult{}
|
||||
}
|
||||
|
||||
var newParents []Chunk
|
||||
var children []ChildChunk
|
||||
childSeq := 0
|
||||
for _, parent := range parents {
|
||||
subs := SplitText(parent.Content, childCfg)
|
||||
|
||||
parentIndex := -1
|
||||
if len(subs) > 1 || (len(subs) == 1 && subs[0].Content != parent.Content) {
|
||||
parentIndex = len(newParents)
|
||||
newParents = append(newParents, parent)
|
||||
}
|
||||
|
||||
for _, sub := range subs {
|
||||
// Adjust offsets: sub positions are relative to parent content,
|
||||
// shift to document-level offsets.
|
||||
// Use additive shift (not Content-length based) so that chunks with
|
||||
// prepended context headers keep correct positional tracking.
|
||||
sub.Seq = childSeq
|
||||
sub.Start += parent.Start
|
||||
sub.End += parent.Start
|
||||
children = append(children, ChildChunk{
|
||||
Chunk: sub,
|
||||
ParentIndex: parentIndex,
|
||||
})
|
||||
childSeq++
|
||||
}
|
||||
}
|
||||
return ParentChildResult{Parents: newParents, Children: children}
|
||||
}
|
||||
|
||||
// ExtractImageRefs extracts markdown image references from text.
|
||||
// The URL group supports one level of balanced parentheses so that URLs
|
||||
// like https://example.com/item_(abc)/123 are captured in full.
|
||||
var imageRefPattern = regexp.MustCompile(`!\[([^\]]*)\]\(([^()\s]*(?:\([^)]*\)[^()\s]*)*)\)`)
|
||||
|
||||
func ExtractImageRefs(text string) []ImageRef {
|
||||
text = docparser.UnwrapLinkedImages(text)
|
||||
matches := imageRefPattern.FindAllStringSubmatchIndex(text, -1)
|
||||
var refs []ImageRef
|
||||
for _, m := range matches {
|
||||
refs = append(refs, ImageRef{
|
||||
OriginalRef: text[m[4]:m[5]], // group 2: URL
|
||||
AltText: text[m[2]:m[3]], // group 1: alt text
|
||||
Start: m[0],
|
||||
End: m[1],
|
||||
})
|
||||
}
|
||||
return refs
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,363 @@
|
||||
// Package chunker - strategy.go is the public entry point for adaptive
|
||||
// chunking. Callers invoke Split / SplitParentChild instead of the legacy
|
||||
// SplitText / SplitTextParentChild functions; the strategy resolver picks
|
||||
// a tier based on document profile and the SplitterConfig.Strategy hint.
|
||||
//
|
||||
// The legacy entry points still exist in splitter.go for backwards
|
||||
// compatibility — strategy.go simply layers a tier-selector on top.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// Strategy values for SplitterConfig.Strategy.
|
||||
const (
|
||||
StrategyAuto = "auto"
|
||||
StrategyHeading = "heading"
|
||||
StrategyHeuristic = "heuristic"
|
||||
StrategyRecursive = "recursive"
|
||||
StrategyLegacy = "legacy"
|
||||
)
|
||||
|
||||
// Split chunks text using the strategy configured in cfg. When cfg.Strategy
|
||||
// is empty or "auto" the document profiler picks the tier. The function
|
||||
// always returns a non-nil result: on tier failure the chain falls through
|
||||
// to the legacy splitter, which is the original Tier 3 implementation.
|
||||
//
|
||||
// Hot path: avoids the Diagnostics struct allocation that
|
||||
// SplitWithDiagnostics performs (matters in SplitParentChild where
|
||||
// Split is called once per parent).
|
||||
func Split(text string, cfg SplitterConfig) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
cfg = ensureDefaults(cfg)
|
||||
chain, profile := resolveChainWithProfile(text, cfg)
|
||||
totalChars := len([]rune(text))
|
||||
|
||||
var lastOut []Chunk
|
||||
for i, tier := range chain {
|
||||
out := runTier(tier, text, cfg, profile)
|
||||
if v := ValidateChunks(out, totalChars, cfg.ChunkSize); v.OK {
|
||||
return out
|
||||
} else {
|
||||
logger.Debugf(context.Background(), "chunker: tier %s rejected: %s", tier, v.Reason)
|
||||
}
|
||||
if tier == TierLegacy && i == len(chain)-1 {
|
||||
lastOut = out
|
||||
}
|
||||
}
|
||||
if lastOut != nil {
|
||||
return lastOut
|
||||
}
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// TierRejection records why a tier was rejected by the validator and the
|
||||
// chain advanced to the next tier. Surfaced by SplitWithDiagnostics for
|
||||
// the debug/preview endpoint.
|
||||
type TierRejection struct {
|
||||
Tier StrategyTier `json:"tier"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// Diagnostics captures which tier produced the returned chunks plus the
|
||||
// chain that was attempted, any rejected tiers along the way, and the
|
||||
// document profile that drove tier selection.
|
||||
//
|
||||
// Useful for surfacing in a debug UI; not produced by the normal Split
|
||||
// path. The JSON shape is part of the public preview-endpoint API — keep
|
||||
// field names stable.
|
||||
type Diagnostics struct {
|
||||
SelectedTier StrategyTier `json:"selected_tier"`
|
||||
TierChain []StrategyTier `json:"tier_chain"`
|
||||
Rejected []TierRejection `json:"rejected"`
|
||||
// Profile is set when the auto strategy resolved the chain via the
|
||||
// document profiler. nil when an explicit Strategy bypassed profiling.
|
||||
Profile *DocProfile `json:"profile,omitempty"`
|
||||
}
|
||||
|
||||
// SplitWithDiagnostics is the same as Split but also returns the
|
||||
// diagnostic trace (selected tier, full chain, rejection reasons,
|
||||
// profile when available). Use this for the chunker preview endpoint
|
||||
// where the caller wants to know which tier won and why others lost.
|
||||
func SplitWithDiagnostics(text string, cfg SplitterConfig) ([]Chunk, *Diagnostics) {
|
||||
// Default selected tier to legacy so an empty diag never carries the
|
||||
// zero string — that would render as a blank tag in the debug UI.
|
||||
diag := &Diagnostics{SelectedTier: TierLegacy}
|
||||
if text == "" {
|
||||
return nil, diag
|
||||
}
|
||||
cfg = ensureDefaults(cfg)
|
||||
chain, profile := resolveChainWithProfile(text, cfg)
|
||||
diag.TierChain = chain
|
||||
diag.Profile = profile
|
||||
totalChars := len([]rune(text))
|
||||
|
||||
var lastOut []Chunk
|
||||
var lastTier StrategyTier
|
||||
for i, tier := range chain {
|
||||
out := runTier(tier, text, cfg, profile)
|
||||
v := ValidateChunks(out, totalChars, cfg.ChunkSize)
|
||||
if v.OK {
|
||||
diag.SelectedTier = tier
|
||||
return out, diag
|
||||
}
|
||||
diag.Rejected = append(diag.Rejected, TierRejection{Tier: tier, Reason: v.Reason})
|
||||
logger.Debugf(context.Background(), "chunker: tier %s rejected: %s", tier, v.Reason)
|
||||
if tier == TierLegacy && i == len(chain)-1 {
|
||||
lastOut = out
|
||||
lastTier = tier
|
||||
}
|
||||
}
|
||||
if lastOut != nil {
|
||||
diag.SelectedTier = lastTier
|
||||
return lastOut, diag
|
||||
}
|
||||
// Defensive last-ditch fallback.
|
||||
return SplitText(text, cfg), diag
|
||||
}
|
||||
|
||||
// SplitParentChild is the strategy-aware analog of SplitTextParentChild.
|
||||
// It runs the tier selector for parent splitting, then re-splits each
|
||||
// parent into children with the small-chunk config.
|
||||
//
|
||||
// Child splitting honours childCfg.Strategy. If it is empty/auto and a
|
||||
// parent has its own internal structure (sub-headings, numbered sub-
|
||||
// sections), the appropriate tier picks it up so child chunks carry a
|
||||
// finer-grained breadcrumb than the parent's. Re-profiling each parent
|
||||
// is bounded by O(sum(parent_size)) ≈ O(N) total, which is the same
|
||||
// order as the original parent profiling pass.
|
||||
func SplitParentChild(text string, parentCfg, childCfg SplitterConfig) ParentChildResult {
|
||||
result, _ := splitParentChild(text, parentCfg, childCfg, false)
|
||||
return result
|
||||
}
|
||||
|
||||
// SplitParentChildWithDiagnostics is the diagnostics-aware form of
|
||||
// SplitParentChild. The result is identical to SplitParentChild; diagnostics
|
||||
// describe the strategy selected for the parent split over the full document.
|
||||
// It is intended for debug surfaces rather than the ingestion hot path.
|
||||
func SplitParentChildWithDiagnostics(text string, parentCfg, childCfg SplitterConfig) (ParentChildResult, *Diagnostics) {
|
||||
return splitParentChild(text, parentCfg, childCfg, true)
|
||||
}
|
||||
|
||||
func splitParentChild(text string, parentCfg, childCfg SplitterConfig, withDiagnostics bool) (ParentChildResult, *Diagnostics) {
|
||||
parentCfg = ensureDefaults(parentCfg)
|
||||
childCfg = ensureDefaults(childCfg)
|
||||
|
||||
var (
|
||||
parents []Chunk
|
||||
diag *Diagnostics
|
||||
)
|
||||
if withDiagnostics {
|
||||
parents, diag = SplitWithDiagnostics(text, parentCfg)
|
||||
} else {
|
||||
parents = Split(text, parentCfg)
|
||||
}
|
||||
if len(parents) == 0 {
|
||||
return ParentChildResult{}, diag
|
||||
}
|
||||
|
||||
var newParents []Chunk
|
||||
var children []ChildChunk
|
||||
childSeq := 0
|
||||
for _, parent := range parents {
|
||||
subs := Split(parent.Content, childCfg)
|
||||
|
||||
parentIndex := -1
|
||||
if len(subs) > 1 || (len(subs) == 1 && subs[0].Content != parent.Content) {
|
||||
parentIndex = len(newParents)
|
||||
newParents = append(newParents, parent)
|
||||
}
|
||||
for _, sub := range subs {
|
||||
sub.Seq = childSeq
|
||||
sub.Start += parent.Start
|
||||
sub.End += parent.Start
|
||||
sub.ContextHeader = mergeBreadcrumbs(parent.ContextHeader, sub.ContextHeader)
|
||||
children = append(children, ChildChunk{Chunk: sub, ParentIndex: parentIndex})
|
||||
childSeq++
|
||||
}
|
||||
}
|
||||
return ParentChildResult{Parents: newParents, Children: children}, diag
|
||||
}
|
||||
|
||||
// NormalizeSplitterConfig applies the same base splitter defaults used by
|
||||
// knowledge ingestion before parent-child derivation or single-pass splitting.
|
||||
func NormalizeSplitterConfig(cfg SplitterConfig) SplitterConfig {
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = DefaultChunkSize
|
||||
}
|
||||
if cfg.ChunkOverlap <= 0 {
|
||||
cfg.ChunkOverlap = DefaultChunkOverlap
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
cfg.Separators = []string{"\n\n", "\n", "。"}
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// NormalizeLineEndings canonicalizes text before chunking. Browser textareas
|
||||
// normalize pasted CRLF text to LF, while uploaded files preserve their
|
||||
// original line endings; without this, the same document produces different
|
||||
// character counts and chunk boundaries.
|
||||
func NormalizeLineEndings(text string) string {
|
||||
if !strings.Contains(text, "\r") {
|
||||
return text
|
||||
}
|
||||
text = strings.ReplaceAll(text, "\r\n", "\n")
|
||||
return strings.ReplaceAll(text, "\r", "\n")
|
||||
}
|
||||
|
||||
// DeriveParentChildConfigs produces the exact parent and child splitter
|
||||
// configurations used by knowledge ingestion. Keeping this here lets preview
|
||||
// and ingestion remain in lockstep as the parent-child defaults evolve.
|
||||
// Languages are copied to both levels so parent and child splitters use the same boundary rules.
|
||||
// TokenLimit is copied only to children because parents keep the configured context window.
|
||||
func DeriveParentChildConfigs(base SplitterConfig, parentSize, childSize int) (parent, child SplitterConfig) {
|
||||
if parentSize <= 0 {
|
||||
parentSize = 4096
|
||||
}
|
||||
if childSize <= 0 {
|
||||
childSize = 384
|
||||
}
|
||||
parent = SplitterConfig{
|
||||
ChunkSize: parentSize,
|
||||
ChunkOverlap: base.ChunkOverlap,
|
||||
Separators: base.Separators,
|
||||
Strategy: base.Strategy,
|
||||
Languages: base.Languages,
|
||||
}
|
||||
child = SplitterConfig{
|
||||
ChunkSize: childSize,
|
||||
ChunkOverlap: childSize / 5,
|
||||
Separators: base.Separators,
|
||||
Strategy: base.Strategy,
|
||||
TokenLimit: base.TokenLimit,
|
||||
Languages: base.Languages,
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// mergeBreadcrumbs combines the parent and child heading breadcrumbs into a
|
||||
// single ContextHeader. When the child re-runs heading detection on parent
|
||||
// content, its first breadcrumb line typically duplicates the parent's last
|
||||
// line (the parent's leading heading sits at the top of the child's input);
|
||||
// drop that duplicate so the embedding context isn't redundant.
|
||||
func mergeBreadcrumbs(parent, child string) string {
|
||||
if parent == "" {
|
||||
return child
|
||||
}
|
||||
if child == "" {
|
||||
return parent
|
||||
}
|
||||
parentLines := strings.Split(parent, "\n")
|
||||
childLines := strings.Split(child, "\n")
|
||||
if len(parentLines) > 0 && len(childLines) > 0 &&
|
||||
strings.TrimSpace(parentLines[len(parentLines)-1]) == strings.TrimSpace(childLines[0]) {
|
||||
childLines = childLines[1:]
|
||||
}
|
||||
if len(childLines) == 0 {
|
||||
return parent
|
||||
}
|
||||
return parent + "\n" + strings.Join(childLines, "\n")
|
||||
}
|
||||
|
||||
// resolveChainWithProfile returns the strategy chain to attempt and, when
|
||||
// the chain was selected by the profiler (auto strategy), the DocProfile
|
||||
// that drove the selection. Profile is nil for explicit non-auto strategies
|
||||
// so callers don't pay for an unused profiling pass.
|
||||
func resolveChainWithProfile(text string, cfg SplitterConfig) ([]StrategyTier, *DocProfile) {
|
||||
switch cfg.Strategy {
|
||||
case StrategyHeading:
|
||||
return []StrategyTier{TierHeading, TierLegacy}, nil
|
||||
case StrategyHeuristic:
|
||||
return []StrategyTier{TierHeuristic, TierLegacy}, nil
|
||||
case StrategyRecursive:
|
||||
// "recursive" is a public-API alias for "legacy": both invoke
|
||||
// SplitText. Kept for backwards compatibility with stored configs.
|
||||
return []StrategyTier{TierLegacy}, nil
|
||||
case StrategyLegacy, "":
|
||||
// Empty == legacy preserves backwards compatibility with stored
|
||||
// ChunkingConfig rows that pre-date the Strategy field.
|
||||
return []StrategyTier{TierLegacy}, nil
|
||||
case StrategyAuto:
|
||||
fallthrough
|
||||
default:
|
||||
profile := ProfileDocument(text)
|
||||
return SelectStrategy(profile), profile
|
||||
}
|
||||
}
|
||||
|
||||
// runTier dispatches the splitter implementation for the given tier.
|
||||
// splitByHeadings / splitByHeuristics are package-level vars overridden
|
||||
// from heading_splitter.go / heuristic_splitter.go via init(); legacy
|
||||
// runs SplitText. The default branch is defensive for future
|
||||
// StrategyTier additions.
|
||||
//
|
||||
// profile may be nil when the caller did not run the document profiler
|
||||
// (explicit non-auto strategies skip profiling); splitters that need a
|
||||
// profile compute one on demand.
|
||||
func runTier(tier StrategyTier, text string, cfg SplitterConfig, profile *DocProfile) []Chunk {
|
||||
switch tier {
|
||||
case TierHeading:
|
||||
return splitByHeadings(text, cfg, profile)
|
||||
case TierHeuristic:
|
||||
return splitByHeuristics(text, cfg, profile)
|
||||
case TierLegacy:
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// ensureDefaults fills in zero-value config fields with sane defaults.
|
||||
// Mirrors buildSplitterConfig in internal/application/service/knowledge.go
|
||||
// so direct callers of this package get the same numbers.
|
||||
//
|
||||
// When cfg.TokenLimit is set, ChunkSize is clamped to the character budget
|
||||
// that fits within that token limit (with a 10% safety factor). This makes
|
||||
// chunks safe for embedding APIs that have hard token caps.
|
||||
func ensureDefaults(cfg SplitterConfig) SplitterConfig {
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = DefaultChunkSize
|
||||
}
|
||||
if cfg.ChunkOverlap <= 0 {
|
||||
cfg.ChunkOverlap = DefaultChunkOverlap
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
cfg.Separators = []string{"\n\n", "\n", "。"}
|
||||
}
|
||||
if cfg.TokenLimit > 0 {
|
||||
lang := LangMixed
|
||||
if len(cfg.Languages) > 0 {
|
||||
lang = cfg.Languages[0]
|
||||
}
|
||||
charBudget := CharsForTokenLimit(cfg.TokenLimit, lang)
|
||||
if charBudget > 0 && (cfg.ChunkSize == 0 || charBudget < cfg.ChunkSize) {
|
||||
cfg.ChunkSize = charBudget
|
||||
}
|
||||
}
|
||||
// Guard against pathological overlap configurations: if Overlap exceeds
|
||||
// half of ChunkSize, almost every chunk is duplicate content. Cap it at
|
||||
// ChunkSize/2 so Overlap stays a useful smoothing band rather than a
|
||||
// near-clone of the previous chunk.
|
||||
if cfg.ChunkOverlap > cfg.ChunkSize/2 && cfg.ChunkSize > 0 {
|
||||
cfg.ChunkOverlap = cfg.ChunkSize / 2
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// splitByHeadings is overridden by heading_splitter.go. The default no-op
|
||||
// fallback ignores the profile (profile may be nil; the override computes
|
||||
// one on demand when so).
|
||||
var splitByHeadings = func(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// splitByHeuristics is overridden by heuristic_splitter.go. profile may be nil.
|
||||
var splitByHeuristics = func(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitWithDiagnostics_LegacyStrategy_ReportsLegacyTier(t *testing.T) {
|
||||
// Splittable input so the validator accepts the legacy output cleanly.
|
||||
text := strings.Repeat("Hello world.\n\nNext paragraph here.\n\n", 50)
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{"\n\n", "\n"}, Strategy: StrategyLegacy}
|
||||
chunks, diag := SplitWithDiagnostics(text, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
if diag.SelectedTier != TierLegacy {
|
||||
t.Errorf("expected SelectedTier=legacy, got %s", diag.SelectedTier)
|
||||
}
|
||||
if len(diag.TierChain) != 1 || diag.TierChain[0] != TierLegacy {
|
||||
t.Errorf("expected single-tier chain [legacy], got %v", diag.TierChain)
|
||||
}
|
||||
if len(diag.Rejected) != 0 {
|
||||
t.Errorf("expected no rejections for splittable input, got %v", diag.Rejected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitWithDiagnostics_AutoOnHeadingDoc_PicksHeading(t *testing.T) {
|
||||
doc := strings.Repeat("# Top\nintro paragraph here.\n\n## Section A\nbody A here.\n\n## Section B\nbody B here.\n\n## Section C\nbody C here.\n\n", 1)
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Strategy: StrategyAuto}
|
||||
_, diag := SplitWithDiagnostics(doc, cfg)
|
||||
if len(diag.TierChain) == 0 {
|
||||
t.Fatal("expected non-empty tier chain")
|
||||
}
|
||||
// Heading tier should be tried first for this doc.
|
||||
if diag.TierChain[0] != TierHeading {
|
||||
t.Errorf("expected heading tier first, got chain %v", diag.TierChain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitWithDiagnostics_EmptyText(t *testing.T) {
|
||||
chunks, diag := SplitWithDiagnostics("", DefaultConfig())
|
||||
if chunks != nil {
|
||||
t.Errorf("expected nil chunks for empty text, got %v", chunks)
|
||||
}
|
||||
if diag == nil {
|
||||
t.Fatal("diag must never be nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplit_AndDiagnostics_AgreeOnChunks ensures Split (no diagnostics)
|
||||
// and SplitWithDiagnostics produce the same chunk set for a given input.
|
||||
// They run independent loops as of the post-audit refactor — this test
|
||||
// is the regression wall against them drifting.
|
||||
func TestSplit_AndDiagnostics_AgreeOnChunks(t *testing.T) {
|
||||
text := "para one.\n\npara two.\n\npara three."
|
||||
cfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 10}
|
||||
a := Split(text, cfg)
|
||||
b, diag := SplitWithDiagnostics(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Fatalf("chunk count disagrees: Split=%d Diagnostics=%d", len(a), len(b))
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].Content != b[i].Content || a[i].Start != b[i].Start || a[i].End != b[i].End {
|
||||
t.Errorf("chunk %d differs:\n Split: %+v\n Diag : %+v", i, a[i], b[i])
|
||||
}
|
||||
}
|
||||
if diag == nil {
|
||||
t.Error("diagnostics must not be nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitWithDiagnostics_ProfileSetForAuto verifies that auto-strategy
|
||||
// returns the DocProfile that drove tier selection — required by the
|
||||
// preview endpoint to avoid double-profiling.
|
||||
func TestSplitWithDiagnostics_ProfileSetForAuto(t *testing.T) {
|
||||
doc := "# Top\nintro.\n\n## A\nbody A.\n\n## B\nbody B."
|
||||
_, diag := SplitWithDiagnostics(doc, SplitterConfig{ChunkSize: 200, Strategy: StrategyAuto})
|
||||
if diag.Profile == nil {
|
||||
t.Fatal("auto strategy must populate diag.Profile")
|
||||
}
|
||||
if diag.Profile.MdHeadingTotal == 0 {
|
||||
t.Errorf("profile should have detected headings, got %+v", diag.Profile)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitWithDiagnostics_ProfileNilForExplicit verifies the inverse:
|
||||
// explicit strategies bypass profiling and leave Profile nil so the
|
||||
// preview handler knows to materialize one if it needs stats.
|
||||
func TestSplitWithDiagnostics_ProfileNilForExplicit(t *testing.T) {
|
||||
for _, strat := range []string{StrategyHeading, StrategyHeuristic, StrategyRecursive, StrategyLegacy} {
|
||||
t.Run(strat, func(t *testing.T) {
|
||||
_, diag := SplitWithDiagnostics("plain text", SplitterConfig{ChunkSize: 200, Strategy: strat})
|
||||
if diag.Profile != nil {
|
||||
t.Errorf("strategy %q should leave Profile nil, got %+v", strat, diag.Profile)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitParentChildWithDiagnostics_MatchesSplitParentChild(t *testing.T) {
|
||||
text := strings.Repeat("## Record\n"+strings.Repeat("A sufficiently long entry body. ", 10)+"\n\n", 12)
|
||||
parentCfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{"\n\n", "\n"}, Strategy: StrategyHeading}
|
||||
childCfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Separators: []string{"\n\n", "\n"}, Strategy: StrategyHeading}
|
||||
|
||||
want := SplitParentChild(text, parentCfg, childCfg)
|
||||
got, diag := SplitParentChildWithDiagnostics(text, parentCfg, childCfg)
|
||||
|
||||
if len(got.Parents) != len(want.Parents) || len(got.Children) != len(want.Children) {
|
||||
t.Fatalf("parent-child result differs: got %d parents/%d children, want %d parents/%d children",
|
||||
len(got.Parents), len(got.Children), len(want.Parents), len(want.Children))
|
||||
}
|
||||
for i := range want.Children {
|
||||
if got.Children[i] != want.Children[i] {
|
||||
t.Errorf("child %d differs:\n got: %+v\n want: %+v", i, got.Children[i], want.Children[i])
|
||||
}
|
||||
}
|
||||
if diag == nil || diag.SelectedTier == "" {
|
||||
t.Fatal("expected diagnostics for parent split")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplit_EmptyText(t *testing.T) {
|
||||
if got := Split("", DefaultConfig()); got != nil {
|
||||
t.Errorf("empty text should return nil, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_LegacyStrategy_MatchesSplitText(t *testing.T) {
|
||||
text := strings.Repeat("Hello world.\n\n", 30)
|
||||
cfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Separators: []string{"\n\n"}, Strategy: StrategyLegacy}
|
||||
a := Split(text, cfg)
|
||||
b := SplitText(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Errorf("legacy strategy should match SplitText: got %d vs %d chunks", len(a), len(b))
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].Content != b[i].Content {
|
||||
t.Errorf("chunk %d differs", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_EmptyStrategyEqualsLegacy(t *testing.T) {
|
||||
text := strings.Repeat("Sentence one. Sentence two.\n", 20)
|
||||
cfg := SplitterConfig{ChunkSize: 80, ChunkOverlap: 10}
|
||||
a := Split(text, cfg)
|
||||
cfg.Strategy = StrategyLegacy
|
||||
b := Split(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Errorf("empty Strategy should equal legacy: %d vs %d", len(a), len(b))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_AutoStrategy_PicksHeadingForMarkdownDoc(t *testing.T) {
|
||||
doc := strings.Repeat("# A\nbody\n## B\nbody\n## C\nbody\n## D\nbody\n", 1)
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Strategy: StrategyAuto}
|
||||
// Until heading_splitter is wired, this falls through to SplitText —
|
||||
// just assert we get a valid result.
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Error("auto strategy should produce chunks")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_HeadingStrategyKeepsDistinctTopLevelHeadings(t *testing.T) {
|
||||
doc := `# Intro
|
||||
short intro.
|
||||
|
||||
# Usage
|
||||
short usage.
|
||||
|
||||
# FAQ
|
||||
short faq.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0, Strategy: StrategyHeading}
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("expected one chunk per top-level heading, got %d:\n%v", len(chunks), chunks)
|
||||
}
|
||||
for i, heading := range []string{"# Intro", "# Usage", "# FAQ"} {
|
||||
if !strings.Contains(chunks[i].Content, heading) {
|
||||
t.Errorf("chunk %d should contain heading %q, got:\n%s", i, heading, chunks[i].Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplit_PreservesPositionInvariantAcrossTiers ensures every chunk's
|
||||
// (Start, End, Content) triple stays consistent — End-Start must equal the
|
||||
// rune length of Content, and runes[Start:End] must equal Content. This is
|
||||
// the contract that knowledge.go:2278+ relies on for document reconstruction
|
||||
// during summary generation.
|
||||
func TestSplit_PreservesPositionInvariantAcrossTiers(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"heading-tier": "# Top\nintro paragraph here.\n\n## Section A\nbody A here.\n\n## Section B\nbody B here.\n\n## Section C\nbody C.",
|
||||
"heuristic-tier": strings.Repeat("Kapitel 1: Einleitung\n", 1) + strings.Repeat("Beispieltext. ", 50) +
|
||||
"\n\n" + strings.Repeat("Kapitel 2: Hauptteil\n", 1) + strings.Repeat("Mehr Text. ", 50),
|
||||
"recursive-tier": strings.Repeat("plain prose without structure. ", 100),
|
||||
}
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{"\n\n", "\n", "。", ". "}, Strategy: StrategyAuto}
|
||||
|
||||
for name, doc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
runes := []rune(doc)
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
spanLen := c.End - c.Start
|
||||
if spanLen != contentRuneLen {
|
||||
t.Errorf("chunk %d: End(%d)-Start(%d)=%d but Content has %d runes:\n%q",
|
||||
i, c.End, c.Start, spanLen, contentRuneLen, c.Content)
|
||||
}
|
||||
if c.Start < 0 || c.End > len(runes) {
|
||||
t.Errorf("chunk %d: position out of range Start=%d End=%d totalRunes=%d",
|
||||
i, c.Start, c.End, len(runes))
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(runes) {
|
||||
sliced := string(runes[c.Start:c.End])
|
||||
if sliced != c.Content {
|
||||
t.Errorf("chunk %d: runes[Start:End] differs from Content", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Children should pick up sub-heading context when childCfg uses auto
|
||||
// strategy on a heading-rich document. Previously child splitting was
|
||||
// pinned to the recursive tier, so sub-chunks all shared the parent's
|
||||
// top-level breadcrumb regardless of inner H2/H3 structure.
|
||||
func TestSplitParentChild_AutoStrategyEnrichesChildBreadcrumbs(t *testing.T) {
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet. ", 40)
|
||||
doc := "# Chapter\n" + body +
|
||||
"\n\n## Section A\n" + body +
|
||||
"\n\n## Section B\n" + body +
|
||||
"\n\n## Section C\n" + body
|
||||
seps := []string{"\n\n", "\n", ". "}
|
||||
parentCfg := SplitterConfig{ChunkSize: 800, ChunkOverlap: 80, Strategy: StrategyAuto, Separators: seps}
|
||||
childCfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Strategy: StrategyAuto, Separators: seps}
|
||||
res := SplitParentChild(doc, parentCfg, childCfg)
|
||||
if len(res.Children) == 0 {
|
||||
t.Fatal("expected children")
|
||||
}
|
||||
|
||||
// At least one child should carry a breadcrumb that includes a
|
||||
// sub-heading (## Section …) — that's only possible if the child
|
||||
// splitter re-detected headings inside the parent.
|
||||
saw := false
|
||||
for _, c := range res.Children {
|
||||
if strings.Contains(c.ContextHeader, "## Section") {
|
||||
saw = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Errorf("no child carries a sub-heading breadcrumb. samples:\n %q\n %q",
|
||||
firstHeader(res.Children), lastHeader(res.Children))
|
||||
}
|
||||
|
||||
// And no breadcrumb should contain the same line twice in a row
|
||||
// (mergeBreadcrumbs deduplicates the parent/child seam).
|
||||
for i, c := range res.Children {
|
||||
lines := strings.Split(c.ContextHeader, "\n")
|
||||
for j := 1; j < len(lines); j++ {
|
||||
if strings.TrimSpace(lines[j]) != "" && lines[j] == lines[j-1] {
|
||||
t.Errorf("child[%d] has duplicate breadcrumb lines: %q", i, c.ContextHeader)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func firstHeader(cs []ChildChunk) string {
|
||||
if len(cs) == 0 {
|
||||
return ""
|
||||
}
|
||||
return cs[0].ContextHeader
|
||||
}
|
||||
|
||||
func lastHeader(cs []ChildChunk) string {
|
||||
if len(cs) == 0 {
|
||||
return ""
|
||||
}
|
||||
return cs[len(cs)-1].ContextHeader
|
||||
}
|
||||
|
||||
func TestMergeBreadcrumbs(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
parent, child string
|
||||
want string
|
||||
}{
|
||||
{"empty parent", "", "## Sub", "## Sub"},
|
||||
{"empty child", "# Top", "", "# Top"},
|
||||
{"both empty", "", "", ""},
|
||||
{"disjoint", "# Top", "## Other", "# Top\n## Other"},
|
||||
{"duplicate seam", "# Top\n## A", "## A\n### A1", "# Top\n## A\n### A1"},
|
||||
{"deep duplicate", "# Top", "# Top", "# Top"},
|
||||
{"only whitespace differs", "# Top\n## A", " ## A \n### A1", "# Top\n## A\n### A1"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := mergeBreadcrumbs(tc.parent, tc.child)
|
||||
if got != tc.want {
|
||||
t.Errorf("mergeBreadcrumbs(%q, %q) = %q, want %q", tc.parent, tc.child, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitParentChild_LegacyStrategy(t *testing.T) {
|
||||
text := strings.Repeat("This is a sentence. Another one.\n\n", 50)
|
||||
parentCfg := SplitterConfig{ChunkSize: 400, ChunkOverlap: 40, Strategy: StrategyLegacy}
|
||||
childCfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Strategy: StrategyLegacy}
|
||||
res := SplitParentChild(text, parentCfg, childCfg)
|
||||
if len(res.Children) == 0 {
|
||||
t.Fatal("expected children chunks")
|
||||
}
|
||||
for i, c := range res.Children {
|
||||
if c.ParentIndex >= 0 && c.ParentIndex >= len(res.Parents) {
|
||||
t.Errorf("child[%d] has invalid ParentIndex %d (parents=%d)", i, c.ParentIndex, len(res.Parents))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults(t *testing.T) {
|
||||
cfg := ensureDefaults(SplitterConfig{})
|
||||
if cfg.ChunkSize != DefaultChunkSize {
|
||||
t.Errorf("expected default ChunkSize %d, got %d", DefaultChunkSize, cfg.ChunkSize)
|
||||
}
|
||||
if cfg.ChunkOverlap != DefaultChunkOverlap {
|
||||
t.Errorf("expected default ChunkOverlap %d, got %d", DefaultChunkOverlap, cfg.ChunkOverlap)
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
t.Error("expected default separators")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeSplitterConfig_MatchesIngestionDefaults(t *testing.T) {
|
||||
cfg := NormalizeSplitterConfig(SplitterConfig{})
|
||||
if cfg.ChunkSize != DefaultChunkSize {
|
||||
t.Errorf("chunk size: got %d want %d", cfg.ChunkSize, DefaultChunkSize)
|
||||
}
|
||||
if cfg.ChunkOverlap != DefaultChunkOverlap {
|
||||
t.Errorf("chunk overlap: got %d want %d", cfg.ChunkOverlap, DefaultChunkOverlap)
|
||||
}
|
||||
if len(cfg.Separators) != 3 {
|
||||
t.Fatalf("separators: got %v", cfg.Separators)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeLineEndings(t *testing.T) {
|
||||
input := "first\r\nsecond\rthird\nfourth"
|
||||
if got, want := NormalizeLineEndings(input), "first\nsecond\nthird\nfourth"; got != want {
|
||||
t.Errorf("NormalizeLineEndings() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeriveParentChildConfigs_DefaultSizes(t *testing.T) {
|
||||
base := SplitterConfig{
|
||||
ChunkSize: 1000,
|
||||
ChunkOverlap: 100,
|
||||
Separators: []string{"\n\n", "\n"},
|
||||
Strategy: StrategyHeading,
|
||||
}
|
||||
parent, child := DeriveParentChildConfigs(base, 0, 0)
|
||||
if parent.ChunkSize != 4096 || child.ChunkSize != 384 {
|
||||
t.Fatalf("default sizes: parent=%d child=%d", parent.ChunkSize, child.ChunkSize)
|
||||
}
|
||||
if parent.Strategy != StrategyHeading || child.Strategy != StrategyHeading {
|
||||
t.Fatalf("strategy not propagated: parent=%q child=%q", parent.Strategy, child.Strategy)
|
||||
}
|
||||
if child.ChunkOverlap != 384/5 {
|
||||
t.Fatalf("child overlap: got %d want %d", child.ChunkOverlap, 384/5)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeriveParentChildConfigs_AppliesEmbeddingTokenLimitOnlyToChildren(t *testing.T) {
|
||||
const (
|
||||
parentSize = 4096
|
||||
childSize = 1024
|
||||
tokenLimit = 200
|
||||
)
|
||||
base := SplitterConfig{
|
||||
ChunkSize: DefaultChunkSize,
|
||||
ChunkOverlap: DefaultChunkOverlap,
|
||||
Separators: []string{"。"},
|
||||
Strategy: StrategyLegacy,
|
||||
TokenLimit: tokenLimit,
|
||||
Languages: []string{LangChinese},
|
||||
}
|
||||
|
||||
parent, child := DeriveParentChildConfigs(base, parentSize, childSize)
|
||||
if parent.TokenLimit != 0 {
|
||||
t.Fatalf("parent token limit: got %d want 0", parent.TokenLimit)
|
||||
}
|
||||
if len(parent.Languages) != 1 || parent.Languages[0] != LangChinese {
|
||||
t.Fatalf("parent languages: got %v want [%s]", parent.Languages, LangChinese)
|
||||
}
|
||||
if child.TokenLimit != tokenLimit {
|
||||
t.Fatalf("child token limit: got %d want %d", child.TokenLimit, tokenLimit)
|
||||
}
|
||||
if len(child.Languages) != 1 || child.Languages[0] != LangChinese {
|
||||
t.Fatalf("child languages: got %v want [%s]", child.Languages, LangChinese)
|
||||
}
|
||||
|
||||
text := strings.Repeat("这是用于验证中文分块预算的句子。", 100)
|
||||
result := SplitParentChild(text, parent, child)
|
||||
if len(result.Parents) != 1 {
|
||||
t.Fatalf("parents: got %d want 1", len(result.Parents))
|
||||
}
|
||||
if got := len([]rune(result.Parents[0].Content)); got != len([]rune(text)) {
|
||||
t.Fatalf("parent length: got %d want %d", got, len([]rune(text)))
|
||||
}
|
||||
budget := CharsForTokenLimit(tokenLimit, LangChinese)
|
||||
for i, c := range result.Children {
|
||||
if got := len([]rune(c.Content)); got > budget {
|
||||
t.Fatalf("child[%d] length: got %d want <= %d", i, got, budget)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_Empty(t *testing.T) {
|
||||
if v := ValidateChunks(nil, 1000, 500); v.OK {
|
||||
t.Error("nil chunks should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_SingleChunkLargeDoc(t *testing.T) {
|
||||
c := []Chunk{{Content: strings.Repeat("a", 5000)}}
|
||||
if v := ValidateChunks(c, 5000, 500); v.OK {
|
||||
t.Error("single 10x-too-large chunk should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_AcceptsReasonableOutput(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 480)},
|
||||
{Content: strings.Repeat("b", 510)},
|
||||
{Content: strings.Repeat("c", 460)},
|
||||
}
|
||||
if v := ValidateChunks(chunks, 1500, 512); !v.OK {
|
||||
t.Errorf("reasonable chunks should validate, got: %s", v.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_RejectsOversized(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 100)},
|
||||
{Content: strings.Repeat("b", 5000)}, // > 2x chunkSize
|
||||
}
|
||||
if v := ValidateChunks(chunks, 5100, 1000); v.OK {
|
||||
t.Error("chunk >2x size should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_TolerantTinyTail(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 480)},
|
||||
{Content: strings.Repeat("b", 510)},
|
||||
{Content: "tail"},
|
||||
}
|
||||
if v := ValidateChunks(chunks, 994, 512); !v.OK {
|
||||
t.Errorf("tiny last chunk should be tolerated, got: %s", v.Reason)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestEnsureDefaults_TokenLimitClampsChunkSize(t *testing.T) {
|
||||
cfg := SplitterConfig{
|
||||
ChunkSize: 10000, // huge
|
||||
TokenLimit: 100,
|
||||
Languages: []string{LangEnglish},
|
||||
}
|
||||
out := ensureDefaults(cfg)
|
||||
// 100 tokens * 4 chars/token * 0.9 ≈ 360 chars
|
||||
if out.ChunkSize >= 1000 {
|
||||
t.Errorf("expected ChunkSize clamped by TokenLimit, got %d", out.ChunkSize)
|
||||
}
|
||||
if out.ChunkOverlap >= out.ChunkSize {
|
||||
t.Errorf("overlap should be smaller than clamped chunk size: overlap=%d size=%d", out.ChunkOverlap, out.ChunkSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults_TokenLimitChineseTighter(t *testing.T) {
|
||||
cfgEN := ensureDefaults(SplitterConfig{TokenLimit: 200, Languages: []string{LangEnglish}})
|
||||
cfgZH := ensureDefaults(SplitterConfig{TokenLimit: 200, Languages: []string{LangChinese}})
|
||||
if cfgZH.ChunkSize >= cfgEN.ChunkSize {
|
||||
t.Errorf("Chinese char budget should be tighter than English: zh=%d en=%d", cfgZH.ChunkSize, cfgEN.ChunkSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults_NoTokenLimitKeepsChunkSize(t *testing.T) {
|
||||
cfg := SplitterConfig{ChunkSize: 800}
|
||||
out := ensureDefaults(cfg)
|
||||
if out.ChunkSize != 800 {
|
||||
t.Errorf("ChunkSize should stay 800, got %d", out.ChunkSize)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
// Package chunker - tokens.go provides language-aware token count approximation.
|
||||
//
|
||||
// We avoid pulling in a tokenizer dependency (e.g. tiktoken) and instead use
|
||||
// per-language chars-per-token ratios derived from common embedding model
|
||||
// vocabularies. The numbers are conservative — they tend to slightly
|
||||
// over-estimate token counts so that chunks stay safely under model limits.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Language identifiers used by the token estimator and the heuristic splitter.
|
||||
const (
|
||||
LangEnglish = "en"
|
||||
LangGerman = "de"
|
||||
LangChinese = "zh"
|
||||
LangMixed = "mixed"
|
||||
)
|
||||
|
||||
// charsPerToken holds approximate chars/token ratios per language.
|
||||
// Numbers err on the conservative side so estimates over-shoot a little.
|
||||
var charsPerToken = map[string]float64{
|
||||
LangEnglish: 4.0,
|
||||
LangGerman: 4.5,
|
||||
LangChinese: 1.7,
|
||||
LangMixed: 3.0,
|
||||
}
|
||||
|
||||
// ApproxTokenCount returns a conservative token estimate for s in the given
|
||||
// language. An empty or unknown lang falls back to "mixed".
|
||||
func ApproxTokenCount(s string, lang string) int {
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
return ApproxTokenCountFromRuneLen(utf8.RuneCountInString(s), lang)
|
||||
}
|
||||
|
||||
// ApproxTokenCountFromRuneLen is the allocation-free variant of
|
||||
// ApproxTokenCount when the caller has already computed the rune length.
|
||||
// Use this in hot loops where the same content's rune count would
|
||||
// otherwise be recomputed multiple times (e.g. preview endpoint emitting
|
||||
// per-chunk stats).
|
||||
func ApproxTokenCountFromRuneLen(runeLen int, lang string) int {
|
||||
if runeLen <= 0 {
|
||||
return 0
|
||||
}
|
||||
ratio, ok := charsPerToken[lang]
|
||||
if !ok {
|
||||
ratio = charsPerToken[LangMixed]
|
||||
}
|
||||
approx := float64(runeLen) / ratio
|
||||
if approx < 1 {
|
||||
return 1
|
||||
}
|
||||
return int(approx + 0.5)
|
||||
}
|
||||
|
||||
// DetectLanguage returns a coarse language label by counting CJK runes vs.
|
||||
// Latin runes. The result is one of LangChinese, LangGerman, LangEnglish or
|
||||
// LangMixed. Detection is cheap and meant only for heuristic dispatch — it
|
||||
// is NOT a replacement for proper language identification.
|
||||
func DetectLanguage(s string) string {
|
||||
if s == "" {
|
||||
return LangMixed
|
||||
}
|
||||
var cjk, latin, umlaut int
|
||||
for _, r := range s {
|
||||
switch {
|
||||
case unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hangul, r) || unicode.Is(unicode.Hiragana, r) || unicode.Is(unicode.Katakana, r):
|
||||
cjk++
|
||||
case isGermanUmlaut(r):
|
||||
umlaut++
|
||||
latin++
|
||||
case (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z'):
|
||||
latin++
|
||||
}
|
||||
}
|
||||
total := cjk + latin
|
||||
if total == 0 {
|
||||
return LangMixed
|
||||
}
|
||||
cjkRatio := float64(cjk) / float64(total)
|
||||
latinRatio := float64(latin) / float64(total)
|
||||
// Mixed: meaningful presence of both scripts (>=15% each).
|
||||
if cjkRatio >= 0.15 && latinRatio >= 0.15 {
|
||||
return LangMixed
|
||||
}
|
||||
if cjkRatio > 0.3 {
|
||||
return LangChinese
|
||||
}
|
||||
if umlaut > 0 || hasGermanWords(s) {
|
||||
return LangGerman
|
||||
}
|
||||
return LangEnglish
|
||||
}
|
||||
|
||||
func isGermanUmlaut(r rune) bool {
|
||||
switch r {
|
||||
case 'ä', 'ö', 'ü', 'Ä', 'Ö', 'Ü', 'ß':
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// hasGermanWords does a tiny stop-word check to bias towards "de" when the
|
||||
// text uses common German function words. Cheap heuristic — false positives
|
||||
// on borrowed terms are acceptable.
|
||||
func hasGermanWords(s string) bool {
|
||||
const sample = 512
|
||||
if len(s) > sample {
|
||||
s = s[:sample]
|
||||
}
|
||||
for _, w := range []string{" der ", " die ", " das ", " und ", " ist ", " nicht ", " mit ", " auf "} {
|
||||
if containsLower(s, w) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func containsLower(haystack, needle string) bool {
|
||||
if len(haystack) < len(needle) {
|
||||
return false
|
||||
}
|
||||
for i := 0; i+len(needle) <= len(haystack); i++ {
|
||||
match := true
|
||||
for j := 0; j < len(needle); j++ {
|
||||
h := haystack[i+j]
|
||||
if h >= 'A' && h <= 'Z' {
|
||||
h += 'a' - 'A'
|
||||
}
|
||||
if h != needle[j] {
|
||||
match = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if match {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// CharsForTokenLimit converts a token limit into an approximate character
|
||||
// budget for a given language. Used to size chunks so they fit within an
|
||||
// embedding model's max-token window with a small safety margin.
|
||||
func CharsForTokenLimit(tokens int, lang string) int {
|
||||
if tokens <= 0 {
|
||||
return 0
|
||||
}
|
||||
ratio, ok := charsPerToken[lang]
|
||||
if !ok {
|
||||
ratio = charsPerToken[LangMixed]
|
||||
}
|
||||
// 0.9 safety factor so we under-shoot the model limit instead of overshooting.
|
||||
return int(float64(tokens) * ratio * 0.9)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestApproxTokenCount_English(t *testing.T) {
|
||||
got := ApproxTokenCount("The quick brown fox jumps over the lazy dog.", LangEnglish)
|
||||
// 44 chars / 4 ≈ 11 tokens
|
||||
if got < 9 || got > 13 {
|
||||
t.Errorf("English token estimate out of range: got %d, want 9..13", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_Chinese(t *testing.T) {
|
||||
got := ApproxTokenCount("这是一段中文测试内容用于检验分词估算", LangChinese)
|
||||
// 18 runes / 1.7 ≈ 10
|
||||
if got < 9 || got > 12 {
|
||||
t.Errorf("Chinese token estimate out of range: got %d, want 9..12", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_Empty(t *testing.T) {
|
||||
if got := ApproxTokenCount("", LangEnglish); got != 0 {
|
||||
t.Errorf("empty string should return 0 tokens, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_UnknownLang(t *testing.T) {
|
||||
got := ApproxTokenCount("Hello world hello world", "xx")
|
||||
if got <= 0 {
|
||||
t.Errorf("unknown lang should fall back to mixed, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_English(t *testing.T) {
|
||||
if got := DetectLanguage("The quick brown fox jumps over the lazy dog."); got != LangEnglish {
|
||||
t.Errorf("expected English, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_German(t *testing.T) {
|
||||
if got := DetectLanguage("Der schnelle braune Fuchs springt über den faulen Hund."); got != LangGerman {
|
||||
t.Errorf("expected German, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_GermanByStopwords(t *testing.T) {
|
||||
// No umlauts but plenty of German function words.
|
||||
if got := DetectLanguage("Das ist ein Test und nicht mit Umlauten."); got != LangGerman {
|
||||
t.Errorf("expected German via stopwords, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_Chinese(t *testing.T) {
|
||||
if got := DetectLanguage("这是一段中文测试内容"); got != LangChinese {
|
||||
t.Errorf("expected Chinese, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_Mixed(t *testing.T) {
|
||||
got := DetectLanguage("This 这是 mixed 测试 content with 多语言 inside")
|
||||
if got != LangMixed {
|
||||
t.Errorf("expected Mixed, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCharsForTokenLimit_AppliesSafetyMargin(t *testing.T) {
|
||||
got := CharsForTokenLimit(1000, LangEnglish)
|
||||
// 1000 * 4 * 0.9 = 3600
|
||||
if got < 3500 || got > 3700 {
|
||||
t.Errorf("char budget for 1000 EN tokens out of range: got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCharsForTokenLimit_ZeroTokens(t *testing.T) {
|
||||
if got := CharsForTokenLimit(0, LangEnglish); got != 0 {
|
||||
t.Errorf("zero tokens should give zero chars, got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
// Package chunker - validator.go inspects a tier's output and decides whether
|
||||
// it is good enough to ship or whether the strategy chain should fall through
|
||||
// to the next tier. The validator is intentionally permissive: a single
|
||||
// "obviously broken" output is rejected, but plausible-looking variation is
|
||||
// accepted so we don't oscillate between tiers.
|
||||
package chunker
|
||||
|
||||
import "math"
|
||||
|
||||
// ValidationResult captures the verdict and reason for a chunk-set.
|
||||
type ValidationResult struct {
|
||||
OK bool
|
||||
Reason string
|
||||
}
|
||||
|
||||
// ValidateChunks checks whether the given chunks form a usable result for a
|
||||
// document of totalChars characters with a target chunkSize. Returns OK=true
|
||||
// when no broken-output indicator triggers.
|
||||
func ValidateChunks(chunks []Chunk, totalChars, chunkSize int) ValidationResult {
|
||||
if len(chunks) == 0 {
|
||||
return ValidationResult{Reason: "no chunks produced"}
|
||||
}
|
||||
|
||||
// A single chunk for a document much larger than chunkSize means the
|
||||
// strategy did not actually split — fail so the next tier runs.
|
||||
if len(chunks) == 1 && totalChars > 2*chunkSize {
|
||||
return ValidationResult{Reason: "single chunk for large document"}
|
||||
}
|
||||
|
||||
// Compute size statistics.
|
||||
var sum, sumSq float64
|
||||
maxLen, minLen := 0, math.MaxInt32
|
||||
for _, c := range chunks {
|
||||
l := len([]rune(c.Content))
|
||||
sum += float64(l)
|
||||
sumSq += float64(l * l)
|
||||
if l > maxLen {
|
||||
maxLen = l
|
||||
}
|
||||
if l < minLen {
|
||||
minLen = l
|
||||
}
|
||||
}
|
||||
avg := sum / float64(len(chunks))
|
||||
|
||||
// All but the last chunk should carry meaningful content. We allow the
|
||||
// last chunk to be tiny because tail residue is normal.
|
||||
tinyCount := 0
|
||||
for i, c := range chunks {
|
||||
if i == len(chunks)-1 {
|
||||
continue
|
||||
}
|
||||
if len([]rune(c.Content)) < 50 {
|
||||
tinyCount++
|
||||
}
|
||||
}
|
||||
if tinyCount > len(chunks)/4 && tinyCount > 2 {
|
||||
return ValidationResult{Reason: "too many tiny chunks"}
|
||||
}
|
||||
|
||||
// Reject when no chunk reached at least 25% of the target — the splitter
|
||||
// is fragmenting too aggressively to be useful.
|
||||
if maxLen < chunkSize/4 && totalChars > chunkSize {
|
||||
return ValidationResult{Reason: "all chunks far below target size"}
|
||||
}
|
||||
|
||||
// Sanity check on absolute upper bound. Anything past 2x chunkSize is a
|
||||
// red flag — the splitter ignored its size budget.
|
||||
if maxLen > 2*chunkSize && chunkSize > 0 {
|
||||
return ValidationResult{Reason: "chunk exceeds 2x target size"}
|
||||
}
|
||||
|
||||
_ = avg
|
||||
return ValidationResult{OK: true}
|
||||
}
|
||||
Reference in New Issue
Block a user