feat: add OpenAI embedding wrapper and in-memory vector index

internal/index/hybrid.go

@@ -0,0 +1,77 @@
+package index
+
+import (
+	"context"
+
+	"eino-test/internal/embedding"
+	"eino-test/internal/store"
+
+	"github.com/cloudwego/eino/components/retriever"
+	"github.com/cloudwego/eino/schema"
+)
+
+type HybridRetriever struct {
+	vectorIdx *InMemoryVectorIndex
+	embedder  *embedding.Embedder
+	noteStore store.NoteStore
+	topK      int
+}
+
+func NewHybridRetriever(vectorIdx *InMemoryVectorIndex, embedder *embedding.Embedder, noteStore store.NoteStore, topK int) *HybridRetriever {
+	return &HybridRetriever{
+		vectorIdx: vectorIdx,
+		embedder:  embedder,
+		noteStore: noteStore,
+		topK:      topK,
+	}
+}
+
+func (r *HybridRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {
+	options := &retriever.Options{
+		TopK: &r.topK,
+	}
+	options = retriever.GetCommonOptions(options, opts...)
+
+	topK := r.topK
+	if options.TopK != nil {
+		topK = *options.TopK
+	}
+
+	queryVecs, err := r.embedder.Embed(ctx, []string{query})
+	if err != nil {
+		return nil, err
+	}
+
+	hits := r.vectorIdx.Search(queryVecs[0], topK)
+
+	var docs []*schema.Document
+	for _, hit := range hits {
+		if options.ScoreThreshold != nil && hit.Score < *options.ScoreThreshold {
+			continue
+		}
+		note, err := r.noteStore.GetByID(ctx, noteIDFromChunkID(hit.ID))
+		if err != nil {
+			continue
+		}
+		docs = append(docs, &schema.Document{
+			ID:      hit.ID,
+			Content: note.Content,
+			MetaData: map[string]any{
+				"note_id": note.ID,
+				"title":   note.Title,
+				"tags":    note.Tags,
+				"score":   hit.Score,
+			},
+		})
+	}
+	return docs, nil
+}
+
+func noteIDFromChunkID(chunkID string) string {
+	for i := len(chunkID) - 1; i >= 0; i-- {
+		if chunkID[i] == '_' {
+			return chunkID[:i]
+		}
+	}
+	return chunkID
+}

internal/index/vector.go

@@ -0,0 +1,119 @@
+package index
+
+import (
+	"math"
+	"sort"
+	"sync"
+)
+
+type SearchHit struct {
+	ID    string
+	Score float64
+}
+
+type InMemoryVectorIndex struct {
+	mu      sync.RWMutex
+	ids     []string
+	vectors [][]float64
+	idToIdx map[string]int
+}
+
+func NewInMemoryVectorIndex() *InMemoryVectorIndex {
+	return &InMemoryVectorIndex{
+		idToIdx: make(map[string]int),
+	}
+}
+
+func (idx *InMemoryVectorIndex) Add(id string, vector []float64) {
+	idx.mu.Lock()
+	defer idx.mu.Unlock()
+	if i, exists := idx.idToIdx[id]; exists {
+		idx.vectors[i] = vector
+		return
+	}
+	idx.idToIdx[id] = len(idx.ids)
+	idx.ids = append(idx.ids, id)
+	idx.vectors = append(idx.vectors, vector)
+}
+
+func (idx *InMemoryVectorIndex) Remove(id string) {
+	idx.mu.Lock()
+	defer idx.mu.Unlock()
+	i, exists := idx.idToIdx[id]
+	if !exists {
+		return
+	}
+	delete(idx.idToIdx, id)
+	last := len(idx.ids) - 1
+	if i < last {
+		idx.ids[i] = idx.ids[last]
+		idx.vectors[i] = idx.vectors[last]
+		idx.idToIdx[idx.ids[i]] = i
+	}
+	idx.ids = idx.ids[:last]
+	idx.vectors = idx.vectors[:last]
+}
+
+func (idx *InMemoryVectorIndex) RemoveByPrefix(prefix string) {
+	idx.mu.Lock()
+	defer idx.mu.Unlock()
+	n := 0
+	for i, id := range idx.ids {
+		if len(id) >= len(prefix) && id[:len(prefix)] == prefix {
+			delete(idx.idToIdx, id)
+			continue
+		}
+		if i != n {
+			idx.ids[n] = idx.ids[i]
+			idx.vectors[n] = idx.vectors[i]
+			idx.idToIdx[idx.ids[n]] = n
+		}
+		n++
+	}
+	idx.ids = idx.ids[:n]
+	idx.vectors = idx.vectors[:n]
+}
+
+func (idx *InMemoryVectorIndex) Search(query []float64, topK int) []SearchHit {
+	idx.mu.RLock()
+	defer idx.mu.RUnlock()
+
+	scores := make([]SearchHit, len(idx.ids))
+	for i, vec := range idx.vectors {
+		scores[i] = SearchHit{ID: idx.ids[i], Score: cosine(query, vec)}
+	}
+	sort.Slice(scores, func(i, j int) bool { return scores[i].Score > scores[j].Score })
+	if topK > len(scores) {
+		topK = len(scores)
+	}
+	return scores[:topK]
+}
+
+func (idx *InMemoryVectorIndex) GetVector(id string) ([]float64, bool) {
+	idx.mu.RLock()
+	defer idx.mu.RUnlock()
+	i, exists := idx.idToIdx[id]
+	if !exists {
+		return nil, false
+	}
+	return idx.vectors[i], true
+}
+
+func (idx *InMemoryVectorIndex) Len() int {
+	idx.mu.RLock()
+	defer idx.mu.RUnlock()
+	return len(idx.ids)
+}
+
+func cosine(a, b []float64) float64 {
+	var dot, normA, normB float64
+	for i := range a {
+		dot += a[i] * b[i]
+		normA += a[i] * a[i]
+		normB += b[i] * b[i]
+	}
+	if normA == 0 || normB == 0 {
+		return 0
+	}
+	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
+}