feat: add RAG pipeline with markdown chunker

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
+}

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
+}