Agent-System/internal/controlcenter/docs/chunk.go

package docs

import (
	"strings"
)

// maxChunkTokens bounds one chunk; oversized sections are bisected by
// paragraph until they fit (or cannot be split further).
const maxChunkTokens = 1800

// Chunk is one embeddable piece of a document.
type Chunk struct {
	Seq     int
	Heading string
	Content string
}

// chunkMarkdown splits md at ATX heading lines; content before the first
// heading becomes a chunk with an empty heading. Oversized sections are then
// bisected by paragraph. Heading text is kept inside the content so the
// embedding carries it.
func chunkMarkdown(md string) []Chunk {
	var chunks []Chunk
	for _, sec := range splitSections(md) {
		chunks = append(chunks, bisect(sec.heading, sec.content)...)
	}
	for i := range chunks {
		chunks[i].Seq = i
	}
	return chunks
}

type section struct {
	heading string
	content string
}

func splitSections(md string) []section {
	var secs []section
	var cur section
	var lines []string
	flush := func() {
		body := strings.TrimSpace(strings.Join(lines, "\n"))
		if cur.heading == "" && body == "" {
			lines = nil
			return
		}
		cur.content = strings.TrimSpace(cur.heading + "\n" + body)
		if cur.content != "" {
			secs = append(secs, cur)
		}
		lines = nil
	}
	for _, line := range strings.Split(md, "\n") {
		if h, ok := atxHeading(line); ok {
			flush()
			cur = section{heading: h}
			continue
		}
		lines = append(lines, line)
	}
	flush()
	return secs
}

// atxHeading returns the heading line (with # marks) for ATX headings.
func atxHeading(line string) (string, bool) {
	t := strings.TrimSpace(line)
	if !strings.HasPrefix(t, "#") {
		return "", false
	}
	i := 0
	for i < len(t) && t[i] == '#' {
		i++
	}
	if i > 6 || i >= len(t) || t[i] != ' ' {
		return "", false
	}
	return t, true
}

// bisect splits an oversized section by paragraphs, halving recursively; a
// single oversized paragraph is cut by rune count.
func bisect(heading, content string) []Chunk {
	if estimateTokens(content) <= maxChunkTokens {
		return []Chunk{{Heading: heading, Content: content}}
	}
	paras := strings.Split(content, "\n\n")
	if len(paras) > 1 {
		mid := len(paras) / 2
		left := bisect(heading, strings.TrimSpace(strings.Join(paras[:mid], "\n\n")))
		right := bisect(heading, strings.TrimSpace(strings.Join(paras[mid:], "\n\n")))
		return append(left, right...)
	}
	// single huge paragraph: hard cut on a rune boundary
	runes := []rune(content)
	half := len(runes) / 2
	for half > 0 && half < len(runes)-1 && runes[half] != ' ' && runes[half] != '\n' && runes[half] != '。' {
		half++
	}
	if half <= 0 || half >= len(runes) {
		half = len(runes) / 2
	}
	left := bisect(heading, strings.TrimSpace(string(runes[:half])))
	right := bisect(heading, strings.TrimSpace(string(runes[half:])))
	return append(left, right...)
}

// estimateTokens approximates token count: CJK runes cost 2, ASCII runs cost
// ceil(len/4). Same heuristic as internal/kb.
func estimateTokens(s string) int {
	s = strings.TrimSpace(s)
	if s == "" {
		return 1
	}
	tokens, asciiRun := 0, 0
	flush := func() {
		if asciiRun > 0 {
			tokens += (asciiRun + 3) / 4
			asciiRun = 0
		}
	}
	for _, r := range s {
		if r <= 127 {
			asciiRun++
			continue
		}
		flush()
		tokens += 2
	}
	flush()
	if tokens < 1 {
		return 1
	}
	return tokens
}