121 lines
3.1 KiB
Go
121 lines
3.1 KiB
Go
package domain
|
||
|
||
import (
|
||
"strings"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
const (
|
||
// chunkMaxChars is the target upper bound for one chunk. ~1,800 chars
|
||
// is roughly 450–600 tokens of code: small enough that a similarity
|
||
// hit points at a specific function rather than a whole file, large
|
||
// enough that a function usually fits in one chunk. It is well under
|
||
// nomic-embed-text's 2,048-token context, so nothing is truncated.
|
||
chunkMaxChars = 1800
|
||
|
||
// chunkOverlapChars of trailing context (whole lines) is repeated at
|
||
// the start of the next chunk, so a match that straddles a boundary
|
||
// still lands in at least one chunk with its surroundings.
|
||
chunkOverlapChars = 240
|
||
)
|
||
|
||
// ChunkText splits source text into overlapping, line-aligned chunks of
|
||
// at most ~chunkMaxChars characters. Whitespace-only input yields nil.
|
||
//
|
||
// Chunks break on line boundaries (never mid-line, except for a single
|
||
// line longer than the limit, e.g. minified JS, which is split at rune
|
||
// boundaries) so each chunk reads as valid-looking code.
|
||
func ChunkText(content string) []string {
|
||
if strings.TrimSpace(content) == "" {
|
||
return nil
|
||
}
|
||
content = strings.ReplaceAll(content, "\r\n", "\n")
|
||
|
||
lines := strings.Split(strings.TrimRight(content, "\n"), "\n")
|
||
|
||
var (
|
||
chunks []string
|
||
cur []string
|
||
curSize int
|
||
)
|
||
|
||
flush := func() {
|
||
text := strings.Join(cur, "\n")
|
||
if strings.TrimSpace(text) != "" {
|
||
chunks = append(chunks, text)
|
||
}
|
||
}
|
||
|
||
// startNext seeds the next chunk with up to chunkOverlapChars of the
|
||
// previous chunk's trailing lines.
|
||
startNext := func() {
|
||
var keep []string
|
||
size := 0
|
||
for i := len(cur) - 1; i >= 0; i-- {
|
||
l := len(cur[i]) + 1
|
||
if size+l > chunkOverlapChars {
|
||
break
|
||
}
|
||
keep = append([]string{cur[i]}, keep...)
|
||
size += l
|
||
}
|
||
cur, curSize = keep, size
|
||
}
|
||
|
||
for _, line := range lines {
|
||
for _, piece := range splitLongLine(line, chunkMaxChars) {
|
||
add := len(piece) + 1
|
||
if curSize+add > chunkMaxChars && len(cur) > 0 {
|
||
flush()
|
||
startNext()
|
||
// If the overlap plus this piece still overflows, drop the overlap:
|
||
// guarantees forward progress for pathological inputs.
|
||
if curSize+add > chunkMaxChars {
|
||
cur, curSize = nil, 0
|
||
}
|
||
}
|
||
cur = append(cur, piece)
|
||
curSize += add
|
||
}
|
||
}
|
||
if len(cur) > 0 {
|
||
flush()
|
||
}
|
||
return dedupeTail(chunks)
|
||
}
|
||
|
||
// splitLongLine breaks a single line longer than max bytes into pieces of
|
||
// at most max bytes, cutting only at rune boundaries.
|
||
func splitLongLine(line string, max int) []string {
|
||
if len(line) <= max {
|
||
return []string{line}
|
||
}
|
||
var out []string
|
||
for len(line) > max {
|
||
cut := max
|
||
for cut > 0 && !utf8.RuneStart(line[cut]) {
|
||
cut--
|
||
}
|
||
if cut == 0 { // no rune boundary found (invalid UTF-8): hard cut
|
||
cut = max
|
||
}
|
||
out = append(out, line[:cut])
|
||
line = line[cut:]
|
||
}
|
||
if line != "" {
|
||
out = append(out, line)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// dedupeTail drops a final chunk that is nothing but overlap already
|
||
// contained in the previous chunk (which happens when the last flush
|
||
// carried only re-seeded lines).
|
||
func dedupeTail(chunks []string) []string {
|
||
n := len(chunks)
|
||
if n >= 2 && strings.HasSuffix(chunks[n-2], chunks[n-1]) {
|
||
return chunks[:n-1]
|
||
}
|
||
return chunks
|
||
}
|