eino-test/internal/rag/pipeline.go

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
}