Init
This commit is contained in:
mlogclub
@@ -0,0 +1,44 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"cs-agent/internal/pkg/enums"
|
||||
)
|
||||
|
||||
type fixedProvider struct{}
|
||||
|
||||
func NewFixedProvider() Provider {
|
||||
return &fixedProvider{}
|
||||
}
|
||||
|
||||
func (p *fixedProvider) Name() string {
|
||||
return string(enums.KnowledgeChunkProviderFixed)
|
||||
}
|
||||
|
||||
func (p *fixedProvider) Supports(contentType enums.KnowledgeDocumentContentType) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (p *fixedProvider) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||
text := req.PlainText
|
||||
if text == "" {
|
||||
text = req.Content
|
||||
}
|
||||
parts := splitPlainText(text, req.Options)
|
||||
results := make([]ChunkResult, 0, len(parts))
|
||||
for i, part := range parts {
|
||||
results = append(results, ChunkResult{
|
||||
ChunkNo: i,
|
||||
Title: req.DocumentTitle,
|
||||
Content: part,
|
||||
ChunkType: enums.KnowledgeChunkTypeText,
|
||||
SectionPath: req.DocumentTitle,
|
||||
CharCount: len([]rune(part)),
|
||||
TokenCount: estimateTokenCount(part),
|
||||
Metadata: map[string]any{
|
||||
"provider": enums.KnowledgeChunkProviderFixed,
|
||||
},
|
||||
})
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"cs-agent/internal/pkg/enums"
|
||||
)
|
||||
|
||||
type Provider interface {
|
||||
Name() string
|
||||
Supports(contentType enums.KnowledgeDocumentContentType) bool
|
||||
Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error)
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"cs-agent/internal/pkg/enums"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
type Registry struct {
|
||||
providers map[string]Provider
|
||||
}
|
||||
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{
|
||||
providers: make(map[string]Provider),
|
||||
}
|
||||
}
|
||||
|
||||
func NewDefaultRegistry() *Registry {
|
||||
r := NewRegistry()
|
||||
r.Register(NewFixedProvider())
|
||||
r.Register(NewStructuredProvider())
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Registry) Register(p Provider) {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
r.providers[p.Name()] = p
|
||||
}
|
||||
|
||||
func (r *Registry) Get(name string) Provider {
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
return r.providers[name]
|
||||
}
|
||||
|
||||
func (r *Registry) Resolve(name string, contentType enums.KnowledgeDocumentContentType) Provider {
|
||||
if p := r.Get(name); p != nil && p.Supports(contentType) {
|
||||
return p
|
||||
}
|
||||
if p := r.Get(string(enums.KnowledgeChunkProviderStructured)); p != nil && p.Supports(contentType) {
|
||||
return p
|
||||
}
|
||||
return r.Get(string(enums.KnowledgeChunkProviderFixed))
|
||||
}
|
||||
|
||||
func (r *Registry) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("chunk request is nil")
|
||||
}
|
||||
provider := r.Resolve(req.Options.Provider, req.ContentType)
|
||||
if provider == nil {
|
||||
return nil, fmt.Errorf("chunk provider not found")
|
||||
}
|
||||
return provider.Chunk(ctx, req)
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"cs-agent/internal/pkg/enums"
|
||||
"strings"
|
||||
|
||||
"github.com/gomarkdown/markdown"
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
type structuredProvider struct{}
|
||||
|
||||
type contentBlock struct {
|
||||
Type string
|
||||
Level int
|
||||
Text string
|
||||
Title string
|
||||
SectionPath string
|
||||
}
|
||||
|
||||
func NewStructuredProvider() Provider {
|
||||
return &structuredProvider{}
|
||||
}
|
||||
|
||||
func (p *structuredProvider) Name() string {
|
||||
return string(enums.KnowledgeChunkProviderStructured)
|
||||
}
|
||||
|
||||
func (p *structuredProvider) Supports(contentType enums.KnowledgeDocumentContentType) bool {
|
||||
switch contentType {
|
||||
case enums.KnowledgeDocumentContentTypeHTML, enums.KnowledgeDocumentContentTypeMarkdown:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (p *structuredProvider) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||
content := req.Content
|
||||
if req.ContentType == enums.KnowledgeDocumentContentTypeMarkdown {
|
||||
content = string(markdown.ToHTML([]byte(content), nil, nil))
|
||||
}
|
||||
|
||||
blocks := parseStructuredBlocks(content, req.DocumentTitle)
|
||||
if len(blocks) == 0 {
|
||||
return NewFixedProvider().Chunk(ctx, req)
|
||||
}
|
||||
|
||||
results := make([]ChunkResult, 0)
|
||||
chunkNo := 0
|
||||
for _, block := range blocks {
|
||||
parts := splitPlainText(block.Text, req.Options)
|
||||
for _, part := range parts {
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
results = append(results, ChunkResult{
|
||||
ChunkNo: chunkNo,
|
||||
Title: block.Title,
|
||||
Content: part,
|
||||
ChunkType: mapBlockType(block.Type),
|
||||
SectionPath: block.SectionPath,
|
||||
CharCount: len([]rune(part)),
|
||||
TokenCount: estimateTokenCount(part),
|
||||
Metadata: map[string]any{
|
||||
"provider": enums.KnowledgeChunkProviderStructured,
|
||||
"blockType": block.Type,
|
||||
"sectionPath": block.SectionPath,
|
||||
"sectionTitle": block.Title,
|
||||
},
|
||||
})
|
||||
chunkNo++
|
||||
}
|
||||
}
|
||||
if len(results) == 0 {
|
||||
return NewFixedProvider().Chunk(ctx, req)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func parseStructuredBlocks(content string, documentTitle string) []contentBlock {
|
||||
content = strings.TrimSpace(content)
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
parent := &html.Node{Type: html.ElementNode, Data: "div"}
|
||||
nodes, err := html.ParseFragment(strings.NewReader(content), parent)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var blocks []contentBlock
|
||||
headings := make([]string, 0)
|
||||
var walk func(node *html.Node)
|
||||
walk = func(node *html.Node) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
if node.Type == html.ElementNode {
|
||||
switch node.Data {
|
||||
case "h1", "h2", "h3", "h4", "h5", "h6":
|
||||
title := normalizeText(nodeText(node))
|
||||
if title != "" {
|
||||
level := int(node.Data[1] - '0')
|
||||
if level <= 0 {
|
||||
level = 1
|
||||
}
|
||||
headings = updateHeadingPath(headings, level, title)
|
||||
}
|
||||
return
|
||||
case "p":
|
||||
appendBlock(&blocks, "paragraph", normalizeText(nodeText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||
return
|
||||
case "ul", "ol":
|
||||
appendBlock(&blocks, "list", normalizeText(listText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||
return
|
||||
case "table":
|
||||
appendBlock(&blocks, "table", normalizeText(tableText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||
return
|
||||
case "pre", "code":
|
||||
appendBlock(&blocks, "code", normalizeText(nodeText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||
return
|
||||
}
|
||||
}
|
||||
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||
walk(child)
|
||||
}
|
||||
}
|
||||
for _, node := range nodes {
|
||||
walk(node)
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
func appendBlock(blocks *[]contentBlock, blockType string, text string, title string, sectionPath string) {
|
||||
text = normalizeText(text)
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
if sectionPath == "" {
|
||||
sectionPath = title
|
||||
}
|
||||
*blocks = append(*blocks, contentBlock{
|
||||
Type: blockType,
|
||||
Text: text,
|
||||
Title: title,
|
||||
SectionPath: sectionPath,
|
||||
})
|
||||
}
|
||||
|
||||
func updateHeadingPath(headings []string, level int, title string) []string {
|
||||
if level <= 0 {
|
||||
level = 1
|
||||
}
|
||||
if len(headings) >= level {
|
||||
headings = headings[:level-1]
|
||||
}
|
||||
headings = append(headings, title)
|
||||
return headings
|
||||
}
|
||||
|
||||
func currentTitle(headings []string, documentTitle string) string {
|
||||
if len(headings) == 0 {
|
||||
return documentTitle
|
||||
}
|
||||
return headings[len(headings)-1]
|
||||
}
|
||||
|
||||
func mapBlockType(blockType string) enums.KnowledgeChunkType {
|
||||
switch blockType {
|
||||
case "table":
|
||||
return enums.KnowledgeChunkTypeTable
|
||||
case "code":
|
||||
return enums.KnowledgeChunkTypeCode
|
||||
default:
|
||||
return enums.KnowledgeChunkTypeText
|
||||
}
|
||||
}
|
||||
|
||||
func nodeText(node *html.Node) string {
|
||||
if node == nil {
|
||||
return ""
|
||||
}
|
||||
var builder strings.Builder
|
||||
writeNodeText(&builder, node)
|
||||
return builder.String()
|
||||
}
|
||||
|
||||
func writeNodeText(builder *strings.Builder, node *html.Node) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
switch node.Type {
|
||||
case html.TextNode:
|
||||
builder.WriteString(node.Data)
|
||||
case html.ElementNode:
|
||||
if shouldSeparate(node.Data) {
|
||||
builder.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||
writeNodeText(builder, child)
|
||||
}
|
||||
if node.Type == html.ElementNode && shouldSeparate(node.Data) {
|
||||
builder.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
|
||||
func shouldSeparate(tag string) bool {
|
||||
switch tag {
|
||||
case "p", "div", "br", "li", "ul", "ol", "blockquote", "pre", "table", "tr", "td", "th", "h1", "h2", "h3", "h4", "h5", "h6":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func listText(node *html.Node) string {
|
||||
items := make([]string, 0)
|
||||
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||
if child.Type == html.ElementNode && child.Data == "li" {
|
||||
item := normalizeText(nodeText(child))
|
||||
if item != "" {
|
||||
items = append(items, item)
|
||||
}
|
||||
}
|
||||
}
|
||||
return strings.Join(items, " ")
|
||||
}
|
||||
|
||||
func tableText(node *html.Node) string {
|
||||
rows := make([]string, 0)
|
||||
var walk func(*html.Node)
|
||||
walk = func(n *html.Node) {
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
if n.Type == html.ElementNode && n.Data == "tr" {
|
||||
cells := make([]string, 0)
|
||||
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||
if child.Type == html.ElementNode && (child.Data == "td" || child.Data == "th") {
|
||||
cell := normalizeText(nodeText(child))
|
||||
if cell != "" {
|
||||
cells = append(cells, cell)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(cells) > 0 {
|
||||
rows = append(rows, strings.Join(cells, " | "))
|
||||
}
|
||||
return
|
||||
}
|
||||
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||
walk(child)
|
||||
}
|
||||
}
|
||||
walk(node)
|
||||
return strings.Join(rows, " ")
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package chunk
|
||||
|
||||
import "cs-agent/internal/pkg/enums"
|
||||
|
||||
type ChunkRequest struct {
|
||||
KnowledgeBaseID int64
|
||||
DocumentID int64
|
||||
DocumentTitle string
|
||||
ContentType enums.KnowledgeDocumentContentType
|
||||
Content string
|
||||
PlainText string
|
||||
Options ChunkOptions
|
||||
}
|
||||
|
||||
type ChunkOptions struct {
|
||||
Provider string
|
||||
TargetTokens int
|
||||
MaxTokens int
|
||||
OverlapTokens int
|
||||
EnableFallback bool
|
||||
}
|
||||
|
||||
type ChunkResult struct {
|
||||
ChunkNo int
|
||||
Title string
|
||||
Content string
|
||||
ChunkType enums.KnowledgeChunkType
|
||||
SectionPath string
|
||||
CharCount int
|
||||
TokenCount int
|
||||
Metadata map[string]any
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"cs-agent/internal/pkg/enums"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultTargetTokens = 300
|
||||
defaultMaxTokens = 400
|
||||
defaultOverlapTokens = 40
|
||||
)
|
||||
|
||||
func normalizeOptions(opts ChunkOptions) ChunkOptions {
|
||||
if opts.TargetTokens <= 0 {
|
||||
opts.TargetTokens = defaultTargetTokens
|
||||
}
|
||||
if opts.MaxTokens <= 0 {
|
||||
opts.MaxTokens = defaultMaxTokens
|
||||
}
|
||||
if opts.MaxTokens < opts.TargetTokens {
|
||||
opts.MaxTokens = opts.TargetTokens
|
||||
}
|
||||
if opts.OverlapTokens < 0 {
|
||||
opts.OverlapTokens = 0
|
||||
}
|
||||
if opts.OverlapTokens == 0 {
|
||||
opts.OverlapTokens = defaultOverlapTokens
|
||||
}
|
||||
if opts.Provider == "" {
|
||||
opts.Provider = string(enums.KnowledgeChunkProviderStructured)
|
||||
}
|
||||
return opts
|
||||
}
|
||||
|
||||
func normalizeText(text string) string {
|
||||
return strings.Join(strings.Fields(strings.TrimSpace(text)), " ")
|
||||
}
|
||||
|
||||
func estimateTokenCount(text string) int {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
count := 0
|
||||
inWord := false
|
||||
for _, r := range text {
|
||||
switch {
|
||||
case unicode.IsSpace(r):
|
||||
inWord = false
|
||||
case unicode.Is(unicode.Han, r):
|
||||
count++
|
||||
inWord = false
|
||||
case unicode.IsLetter(r) || unicode.IsDigit(r):
|
||||
if !inWord {
|
||||
count++
|
||||
inWord = true
|
||||
}
|
||||
default:
|
||||
count++
|
||||
inWord = false
|
||||
}
|
||||
}
|
||||
if count == 0 {
|
||||
return utf8.RuneCountInString(text)
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func contentHash(text string) string {
|
||||
sum := sha256.Sum256([]byte(text))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func splitSentences(text string) []string {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
var sentences []string
|
||||
var builder strings.Builder
|
||||
for _, r := range text {
|
||||
builder.WriteRune(r)
|
||||
switch r {
|
||||
case '\n', '。', '！', '？', '!', '?', ';', '；':
|
||||
sentence := normalizeText(builder.String())
|
||||
if sentence != "" {
|
||||
sentences = append(sentences, sentence)
|
||||
}
|
||||
builder.Reset()
|
||||
}
|
||||
}
|
||||
if builder.Len() > 0 {
|
||||
sentence := normalizeText(builder.String())
|
||||
if sentence != "" {
|
||||
sentences = append(sentences, sentence)
|
||||
}
|
||||
}
|
||||
if len(sentences) == 0 {
|
||||
return []string{normalizeText(text)}
|
||||
}
|
||||
return sentences
|
||||
}
|
||||
|
||||
func tailTextByTokens(text string, tokenLimit int) string {
|
||||
if tokenLimit <= 0 {
|
||||
return ""
|
||||
}
|
||||
sentences := splitSentences(text)
|
||||
if len(sentences) == 0 {
|
||||
return ""
|
||||
}
|
||||
var selected []string
|
||||
total := 0
|
||||
for i := len(sentences) - 1; i >= 0; i-- {
|
||||
sentence := sentences[i]
|
||||
tokens := estimateTokenCount(sentence)
|
||||
if total > 0 && total+tokens > tokenLimit {
|
||||
break
|
||||
}
|
||||
selected = append([]string{sentence}, selected...)
|
||||
total += tokens
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(selected, " "))
|
||||
}
|
||||
|
||||
func splitPlainText(text string, opts ChunkOptions) []string {
|
||||
text = normalizeText(text)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
opts = normalizeOptions(opts)
|
||||
sentences := splitSentences(text)
|
||||
if len(sentences) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
chunks := make([]string, 0)
|
||||
current := make([]string, 0)
|
||||
currentTokens := 0
|
||||
|
||||
flush := func() {
|
||||
if len(current) == 0 {
|
||||
return
|
||||
}
|
||||
chunks = append(chunks, strings.Join(current, " "))
|
||||
}
|
||||
|
||||
for _, sentence := range sentences {
|
||||
sentenceTokens := estimateTokenCount(sentence)
|
||||
if sentenceTokens > opts.MaxTokens {
|
||||
if len(current) > 0 {
|
||||
flush()
|
||||
overlap := tailTextByTokens(strings.Join(current, " "), opts.OverlapTokens)
|
||||
current = nil
|
||||
currentTokens = 0
|
||||
if overlap != "" {
|
||||
current = append(current, overlap)
|
||||
currentTokens = estimateTokenCount(overlap)
|
||||
}
|
||||
}
|
||||
for _, piece := range splitLongSentence(sentence, opts.MaxTokens) {
|
||||
piece = normalizeText(piece)
|
||||
if piece != "" {
|
||||
chunks = append(chunks, piece)
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if currentTokens > 0 && currentTokens+sentenceTokens > opts.MaxTokens {
|
||||
flush()
|
||||
overlap := tailTextByTokens(strings.Join(current, " "), opts.OverlapTokens)
|
||||
current = nil
|
||||
currentTokens = 0
|
||||
if overlap != "" {
|
||||
current = append(current, overlap)
|
||||
currentTokens = estimateTokenCount(overlap)
|
||||
}
|
||||
}
|
||||
|
||||
current = append(current, sentence)
|
||||
currentTokens += sentenceTokens
|
||||
}
|
||||
|
||||
flush()
|
||||
return chunks
|
||||
}
|
||||
|
||||
func splitLongSentence(text string, maxTokens int) []string {
|
||||
runes := []rune(strings.TrimSpace(text))
|
||||
if len(runes) == 0 {
|
||||
return nil
|
||||
}
|
||||
if maxTokens <= 0 {
|
||||
return []string{text}
|
||||
}
|
||||
window := maxTokens * 2
|
||||
if window < 50 {
|
||||
window = 50
|
||||
}
|
||||
var result []string
|
||||
for start := 0; start < len(runes); start += window {
|
||||
end := start + window
|
||||
if end > len(runes) {
|
||||
end = len(runes)
|
||||
}
|
||||
part := normalizeText(string(runes[start:end]))
|
||||
if part != "" {
|
||||
result = append(result, part)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
Reference in New Issue
Block a user