From d2c5dafabef5d041a6dfc0add5edc03496547987 Mon Sep 17 00:00:00 2001 From: wonder Date: Tue, 5 May 2026 11:27:48 +0800 Subject: [PATCH] feat: add RAG pipeline with markdown chunker --- internal/rag/chunker.go | 68 ++++++++++++++++ internal/rag/pipeline.go | 169 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 237 insertions(+) create mode 100644 internal/rag/chunker.go create mode 100644 internal/rag/pipeline.go diff --git a/internal/rag/chunker.go b/internal/rag/chunker.go new file mode 100644 index 0000000..b0ff77f --- /dev/null +++ b/internal/rag/chunker.go @@ -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 +} diff --git a/internal/rag/pipeline.go b/internal/rag/pipeline.go new file mode 100644 index 0000000..a60d54e --- /dev/null +++ b/internal/rag/pipeline.go @@ -0,0 +1,169 @@ +package rag + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "eino-test/internal/embedding" + "eino-test/internal/index" + "eino-test/internal/store" +) + +type RAGPipeline struct { + noteStore store.NoteStore + embedder *embedding.Embedder + vectorIdx *index.InMemoryVectorIndex + chunkSize int +} + +func NewRAGPipeline(noteStore store.NoteStore, embedder *embedding.Embedder, vectorIdx *index.InMemoryVectorIndex) *RAGPipeline { + return &RAGPipeline{ + noteStore: noteStore, + embedder: embedder, + vectorIdx: vectorIdx, + chunkSize: 500, + } +} + +func (p *RAGPipeline) IndexNote(ctx context.Context, note *store.Note) error { + chunks := ChunkMarkdown(note.Content, p.chunkSize) + if len(chunks) == 0 { + return nil + } + + chunkIDs := make([]string, len(chunks)) + for i := range chunks { + chunkIDs[i] = fmt.Sprintf("%s_%d", note.ID, i) + } + + vectors, err := p.embedder.Embed(ctx, chunks) + if err != nil { + return fmt.Errorf("embed chunks: %w", err) + } + + for i, vec := range vectors { + p.vectorIdx.Add(chunkIDs[i], vec) + } + return nil +} + +func (p *RAGPipeline) RemoveNote(noteID string) { + p.vectorIdx.RemoveByPrefix(noteID + "_") +} + +func (p *RAGPipeline) IndexAllFromDir(ctx context.Context, dir string) error { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".md") { + continue + } + path := filepath.Join(dir, entry.Name()) + content, err := os.ReadFile(path) + if err != nil { + return fmt.Errorf("read %s: %w", path, err) + } + + title := strings.TrimSuffix(entry.Name(), ".md") + note := &store.Note{ + ID: title, + Title: title, + Content: string(content), + Tags: extractFrontmatterTags(string(content)), + FilePath: path, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + + if err := p.noteStore.Create(ctx, note); err != nil { + continue + } + if err := p.IndexNote(ctx, note); err != nil { + return fmt.Errorf("index %s: %w", path, err) + } + } + return nil +} + +func (p *RAGPipeline) Search(ctx context.Context, query string, topK int) ([]store.SearchResult, error) { + queryVecs, err := p.embedder.Embed(ctx, []string{query}) + if err != nil { + return nil, err + } + + hits := p.vectorIdx.Search(queryVecs[0], topK) + var results []store.SearchResult + for _, hit := range hits { + noteID := noteIDFromChunkID(hit.ID) + note, err := p.noteStore.GetByID(ctx, noteID) + if err != nil { + continue + } + results = append(results, store.SearchResult{Note: note, Score: hit.Score}) + } + return results, nil +} + +func (p *RAGPipeline) NoteStore() store.NoteStore { + return p.noteStore +} + +func (p *RAGPipeline) VectorIndex() *index.InMemoryVectorIndex { + return p.vectorIdx +} + +func (p *RAGPipeline) Embedder() *embedding.Embedder { + return p.embedder +} + +func noteIDFromChunkID(chunkID string) string { + for i := len(chunkID) - 1; i >= 0; i-- { + if chunkID[i] == '_' { + return chunkID[:i] + } + } + return chunkID +} + +func extractFrontmatterTags(content string) []string { + if !strings.HasPrefix(content, "---") { + return nil + } + end := strings.Index(content[3:], "---") + if end < 0 { + return nil + } + frontmatter := content[3 : end+3] + for _, line := range strings.Split(frontmatter, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "tags:") { + tagStr := strings.TrimPrefix(line, "tags:") + tagStr = strings.TrimSpace(tagStr) + tagStr = strings.Trim(tagStr, "[]") + if tagStr == "" { + return nil + } + parts := strings.Split(tagStr, ",") + var tags []string + for _, t := range parts { + t = strings.TrimSpace(t) + t = strings.Trim(t, "\"'") + if t != "" { + tags = append(tags, t) + } + } + return tags + } + } + return nil +}