69 lines
1.4 KiB
Go
69 lines
1.4 KiB
Go
|
package rag
|
||
|
|
|
||
|
|
import "strings"
|
||
|
|
|
||
|
|
func ChunkMarkdown(content string, maxChunkSize int) []string {
|
||
|
|
if maxChunkSize <= 0 {
|
||
|
|
maxChunkSize = 500
|
||
|
|
}
|
||
|
|
|
||
|
|
sections := splitByHeaders(content)
|
||
|
|
var chunks []string
|
||
|
|
|
||
|
|
for _, section := range sections {
|
||
|
|
if len(section) <= maxChunkSize {
|
||
|
|
chunks = append(chunks, section)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
paragraphs := splitByParagraphs(section)
|
||
|
|
current := ""
|
||
|
|
for _, p := range paragraphs {
|
||
|
|
if len(current)+len(p)+2 > maxChunkSize && current != "" {
|
||
|
|
chunks = append(chunks, strings.TrimSpace(current))
|
||
|
|
current = ""
|
||
|
|
}
|
||
|
|
if current != "" {
|
||
|
|
current += "\n\n"
|
||
|
|
}
|
||
|
|
current += p
|
||
|
|
}
|
||
|
|
if current != "" {
|
||
|
|
chunks = append(chunks, strings.TrimSpace(current))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return chunks
|
||
|
|
}
|
||
|
|
|
||
|
|
func splitByHeaders(content string) []string {
|
||
|
|
lines := strings.Split(content, "\n")
|
||
|
|
var sections []string
|
||
|
|
current := ""
|
||
|
|
|
||
|
|
for _, line := range lines {
|
||
|
|
if strings.HasPrefix(line, "## ") && current != "" {
|
||
|
|
sections = append(sections, strings.TrimSpace(current))
|
||
|
|
current = ""
|
||
|
|
}
|
||
|
|
if current != "" {
|
||
|
|
current += "\n"
|
||
|
|
}
|
||
|
|
current += line
|
||
|
|
}
|
||
|
|
if current != "" {
|
||
|
|
sections = append(sections, strings.TrimSpace(current))
|
||
|
|
}
|
||
|
|
return sections
|
||
|
|
}
|
||
|
|
|
||
|
|
func splitByParagraphs(content string) []string {
|
||
|
|
paragraphs := strings.Split(content, "\n\n")
|
||
|
|
var result []string
|
||
|
|
for _, p := range paragraphs {
|
||
|
|
p = strings.TrimSpace(p)
|
||
|
|
if p != "" {
|
||
|
|
result = append(result, p)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return result
|
||
|
|
}
|