jarvis-home/internal/rag/index.go

package rag

import (
	"encoding/json"
	"errors"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

type Chunk struct {
	ID         string    `json:"id"`
	DocumentID string    `json:"documentId"`
	Name       string    `json:"name"`
	Text       string    `json:"text"`
	Embedding  []float64 `json:"embedding"`
}

type Hit struct {
	DocumentID string  `json:"documentId"`
	Name       string  `json:"name"`
	Text       string  `json:"text"`
	Score      float64 `json:"score"`
}

type Index struct {
	mu     sync.RWMutex
	path   string
	chunks []Chunk
}

func New(path string) (*Index, error) {
	i := &Index{path: path, chunks: []Chunk{}}
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return nil, err
	}
	b, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return i, i.saveLocked()
	}
	if err != nil {
		return nil, err
	}
	if len(b) > 0 {
		if err := json.Unmarshal(b, &i.chunks); err != nil {
			return nil, err
		}
	}
	return i, nil
}

func ChunkText(text string, size, overlap int) []string {
	text = strings.TrimSpace(text)
	if text == "" {
		return nil
	}
	if size < 200 {
		size = 1200
	}
	if overlap < 0 {
		overlap = 0
	}
	if overlap >= size {
		overlap = size / 5
	}
	r := []rune(text)
	var out []string
	for start := 0; start < len(r); {
		end := start + size
		if end > len(r) {
			end = len(r)
		}
		// Prefer a natural boundary near the end.
		if end < len(r) {
			floor := start + size*2/3
			for j := end; j > floor; j-- {
				if r[j-1] == '\n' || r[j-1] == '.' || r[j-1] == '!' || r[j-1] == '?' {
					end = j
					break
				}
			}
		}
		piece := strings.TrimSpace(string(r[start:end]))
		if piece != "" {
			out = append(out, piece)
		}
		if end == len(r) {
			break
		}
		start = end - overlap
		if start < 0 {
			start = 0
		}
	}
	return out
}

func (i *Index) ReplaceDocument(docID, name string, texts []string, vectors [][]float64) error {
	if len(texts) != len(vectors) {
		return errors.New("chunk/vector count mismatch")
	}
	i.mu.Lock()
	defer i.mu.Unlock()
	kept := i.chunks[:0]
	for _, c := range i.chunks {
		if c.DocumentID != docID {
			kept = append(kept, c)
		}
	}
	i.chunks = append([]Chunk(nil), kept...)
	for idx, text := range texts {
		i.chunks = append(i.chunks, Chunk{ID: docID + "#" + itoa(idx), DocumentID: docID, Name: name, Text: text, Embedding: vectors[idx]})
	}
	return i.saveLocked()
}

func (i *Index) DeleteDocument(docID string) error {
	i.mu.Lock()
	defer i.mu.Unlock()
	out := i.chunks[:0]
	for _, c := range i.chunks {
		if c.DocumentID != docID {
			out = append(out, c)
		}
	}
	i.chunks = append([]Chunk(nil), out...)
	return i.saveLocked()
}

func (i *Index) HasDocument(docID string) bool {
	i.mu.RLock()
	defer i.mu.RUnlock()
	for _, c := range i.chunks {
		if c.DocumentID == docID {
			return true
		}
	}
	return false
}

func (i *Index) Count() int { i.mu.RLock(); defer i.mu.RUnlock(); return len(i.chunks) }

func (i *Index) Search(query []float64, k int) []Hit {
	i.mu.RLock()
	defer i.mu.RUnlock()
	if k <= 0 {
		k = 6
	}
	hits := make([]Hit, 0, len(i.chunks))
	for _, c := range i.chunks {
		if len(c.Embedding) == 0 || len(c.Embedding) != len(query) {
			continue
		}
		hits = append(hits, Hit{DocumentID: c.DocumentID, Name: c.Name, Text: c.Text, Score: cosine(query, c.Embedding)})
	}
	sort.Slice(hits, func(a, b int) bool { return hits[a].Score > hits[b].Score })
	if len(hits) > k {
		hits = hits[:k]
	}
	return hits
}

func (i *Index) saveLocked() error {
	b, err := json.Marshal(i.chunks)
	if err != nil {
		return err
	}
	tmp := i.path + ".tmp"
	if err := os.WriteFile(tmp, b, 0600); err != nil {
		return err
	}
	return os.Rename(tmp, i.path)
}

func cosine(a, b []float64) float64 {
	var dot, aa, bb float64
	for idx := range a {
		dot += a[idx] * b[idx]
		aa += a[idx] * a[idx]
		bb += b[idx] * b[idx]
	}
	if aa == 0 || bb == 0 {
		return 0
	}
	return dot / (math.Sqrt(aa) * math.Sqrt(bb))
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	const d = "0123456789"
	var b [20]byte
	p := len(b)
	for n > 0 {
		p--
		b[p] = d[n%10]
		n /= 10
	}
	return string(b[p:])
}