feat: add RAG pipeline with markdown chunker
This commit is contained in:
@@ -0,0 +1,68 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user