chunk.go

package chunky

import (
	"strings"
)

type Chunker interface {
	Chunk(string) []string
}

type SentenceChunker struct {
	separators []byte
	chunkSize  int
}

func NewSentenceChunker(chunkSize int) *SentenceChunker {
	return &SentenceChunker{
		separators: []byte{'\n', '.', '!', '?'},
		chunkSize:  chunkSize,
	}
}

func (c *SentenceChunker) Chunk(text string) []string {
	if text == "" {
		return nil
	}
	data := []byte(text)
	return chunkBytes(data, c.chunkSize, c.separators)
}

func chunkBytes(data []byte, blockSize int, separators []byte) []string {
	var chunks [][]byte
	i := 0
	for i < len(data) {
		end := i + blockSize
		if end > len(data) {
			end = len(data)
		}

		block := data[i:end]

		found := -1
		for j := len(block) - 1; j >= 0; j-- {
			for _, d := range separators {
				if block[j] == d {
					found = j + 1
					break
				}
			}
			if found > 0 {
				break
			}
		}

		if found <= 0 {
			found = len(block)
		}

		chunks = append(chunks, block[:found])
		i += found
	}

	// convert to strings
	out := make([]string, len(chunks))
	for i, c := range chunks {
		out[i] = string(c)
	}
	return out
}

type RecursiveCharacterTextSplitter struct {
	separators   []string
	chunkSize    int
	chunkOverlap int
}

func NewRecursiveCharacterTextSplitter(chunkSize, chunkOverlap int) *RecursiveCharacterTextSplitter {
	return &RecursiveCharacterTextSplitter{
		separators:   []string{"\n\n", "\n", " ", ""},
		chunkSize:    chunkSize,
		chunkOverlap: chunkOverlap,
	}
}

func (c *RecursiveCharacterTextSplitter) Chunk(text string) []string {
	return c.splitText(text, c.separators)
}

func (c *RecursiveCharacterTextSplitter) splitText(text string, separators []string) []string {
	sep, remaining := c.chooseSeparator(text, separators)
	splits := filterEmpty(splitOnSeparator(text, sep))

	var goodSplits []string
	var chunks []string

	for _, piece := range splits {
		if len(piece) <= c.chunkSize {
			goodSplits = append(goodSplits, piece)
			continue
		}

		if len(goodSplits) > 0 {
			chunks = append(chunks, c.mergeSplits(goodSplits, sep)...)
			goodSplits = nil
		}

		if len(remaining) == 0 {
			chunks = append(chunks, piece)
		} else {
			chunks = append(chunks, c.splitText(piece, remaining)...)
		}
	}

	if len(goodSplits) > 0 {
		chunks = append(chunks, c.mergeSplits(goodSplits, sep)...)
	}

	return chunks
}

func (c *RecursiveCharacterTextSplitter) chooseSeparator(text string, separators []string) (string, []string) {
	for i, sep := range separators {
		if sep == "" || strings.Contains(text, sep) {
			return sep, separators[i+1:]
		}
	}
	last := separators[len(separators)-1]
	return last, nil
}

func joinedLen(pieces []string, sep string) int {
	if len(pieces) == 0 {
		return 0
	}
	total := 0
	for _, p := range pieces {
		total += len(p)
	}
	total += len(sep) * (len(pieces) - 1)
	return total
}

func (c *RecursiveCharacterTextSplitter) mergeSplits(splits []string, sep string) []string {
	joinSep := cleanJoinSep(sep)

	var chunks []string
	var current []string

	for _, piece := range splits {
		addLen := len(piece)
		if len(current) > 0 {
			addLen += len(joinSep)
		}

		if joinedLen(current, joinSep)+addLen > c.chunkSize && len(current) > 0 {
			chunks = append(chunks, strings.Join(current, joinSep))
			for len(current) > 0 && joinedLen(current, joinSep) > c.chunkOverlap {
				current = current[1:]
			}
		}

		current = append(current, piece)
	}

	if len(current) > 0 {
		chunks = append(chunks, strings.Join(current, joinSep))
	}

	return chunks
}

func cleanJoinSep(sep string) string {
	if strings.Contains(sep, "\n") {
		return "\n"
	}
	if sep == " " {
		return " "
	}
	return ""
}

func filterEmpty(ss []string) []string {
	out := ss[:0]
	for _, s := range ss {
		if strings.TrimSpace(s) != "" {
			out = append(out, s)
		}
	}
	return out
}

func splitOnSeparator(text, sep string) []string {
	if sep == "" {
		runes := []rune(text)
		out := make([]string, len(runes))
		for i, r := range runes {
			out[i] = string(r)
		}
		return out
	}
	return strings.Split(text, sep)
}