glpi-ai-agent/internal/knowledge/store.go

package knowledge

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"unicode"

	"github.com/example/glpi-ai-agent/internal/model"
)

type Embedder interface {
	Embed(context.Context, []string) ([][]float64, error)
}
type Store struct {
	mu        sync.RWMutex
	docs      []model.KnowledgeDoc
	vectors   map[string][]float64
	embedder  Embedder
	rag       bool
	cachePath string
}
type cacheFile struct {
	Hashes  map[string]string    `json:"hashes"`
	Vectors map[string][]float64 `json:"vectors"`
}

func Load(ctx context.Context, dir, dataDir string, embedder Embedder, rag bool, allowedSources []string) (*Store, error) {
	s := &Store{vectors: map[string][]float64{}, embedder: embedder, rag: rag, cachePath: filepath.Join(dataDir, "embeddings.json")}
	allowed := make(map[string]struct{}, len(allowedSources))
	for _, source := range allowedSources {
		allowed[strings.ToLower(strings.TrimSpace(source))] = struct{}{}
	}
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("read knowledge directory %q: %w", dir, err)
	}
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(strings.ToLower(e.Name()), ".json") {
			continue
		}
		b, err := os.ReadFile(filepath.Join(dir, e.Name()))
		if err != nil {
			return nil, err
		}
		var d model.KnowledgeDoc
		if err := json.Unmarshal(b, &d); err != nil {
			return nil, fmt.Errorf("%s: %w", e.Name(), err)
		}
		if d.ID == "" || d.Title == "" {
			return nil, fmt.Errorf("%s: id/title required", e.Name())
		}
		d.Source = strings.ToLower(strings.TrimSpace(d.Source))
		if d.Source == "" {
			return nil, fmt.Errorf("%s: source required", e.Name())
		}
		if _, ok := allowed[d.Source]; !ok {
			continue
		}
		d.Language = strings.TrimSpace(d.Language)
		d.CommunicationStyle = strings.ToLower(strings.TrimSpace(d.CommunicationStyle))
		s.docs = append(s.docs, d)
	}
	if rag && len(s.docs) > 0 {
		if s.embedder == nil {
			return nil, fmt.Errorf("RAG is enabled but no embedding provider is configured")
		}
		if err := s.index(ctx); err != nil {
			return s, err
		}
	}
	return s, nil
}
func (s *Store) Count() int {
	if s == nil {
		return 0
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.docs)
}
func (s *Store) ByID(id string) (model.KnowledgeDoc, bool) {
	if s == nil {
		return model.KnowledgeDoc{}, false
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, d := range s.docs {
		if d.ID == id {
			return d, true
		}
	}
	return model.KnowledgeDoc{}, false
}
func (s *Store) Search(ctx context.Context, text string, topK int) ([]model.KnowledgeHit, error) {
	if s == nil {
		return nil, fmt.Errorf("knowledge store is not initialized")
	}
	s.mu.RLock()
	docs := append([]model.KnowledgeDoc(nil), s.docs...)
	vecs := make(map[string][]float64, len(s.vectors))
	for k, v := range s.vectors {
		vecs[k] = v
	}
	s.mu.RUnlock()
	if len(docs) == 0 {
		return nil, nil
	}
	scores := map[string]float64{}
	if s.rag && s.embedder != nil && len(vecs) > 0 {
		q, err := s.embedder.Embed(ctx, []string{text})
		if err != nil {
			return nil, err
		}
		if len(q) > 0 {
			for _, d := range docs {
				scores[d.ID] = cosine(q[0], vecs[d.ID])
			}
		}
	} else {
		for _, d := range docs {
			scores[d.ID] = lexical(text, d)
		}
	}
	hits := make([]model.KnowledgeHit, 0, len(docs))
	for _, d := range docs {
		hits = append(hits, model.KnowledgeHit{Doc: d, Score: scores[d.ID]})
	}
	sort.Slice(hits, func(i, j int) bool { return hits[i].Score > hits[j].Score })
	if topK > 0 && len(hits) > topK {
		hits = hits[:topK]
	}
	return hits, nil
}
func (s *Store) index(ctx context.Context) error {
	_ = os.MkdirAll(filepath.Dir(s.cachePath), 0o750)
	cf := cacheFile{Hashes: map[string]string{}, Vectors: map[string][]float64{}}
	if b, err := os.ReadFile(s.cachePath); err == nil {
		_ = json.Unmarshal(b, &cf)
	}
	var need []model.KnowledgeDoc
	for _, d := range s.docs {
		h := hashDoc(d)
		if cf.Hashes[d.ID] == h && len(cf.Vectors[d.ID]) > 0 {
			s.vectors[d.ID] = cf.Vectors[d.ID]
		} else {
			need = append(need, d)
		}
	}
	if len(need) > 0 {
		texts := make([]string, len(need))
		for i, d := range need {
			texts[i] = d.Title + "\n" + d.Text + "\n" + strings.Join(d.Keywords, " ")
		}
		vv, err := s.embedder.Embed(ctx, texts)
		if err != nil {
			return err
		}
		for i, d := range need {
			s.vectors[d.ID] = vv[i]
			cf.Hashes[d.ID] = hashDoc(d)
			cf.Vectors[d.ID] = vv[i]
		}
		b, _ := json.MarshalIndent(cf, "", "  ")
		tmp := s.cachePath + ".tmp"
		if err := os.WriteFile(tmp, b, 0o640); err != nil {
			return err
		}
		if err := os.Rename(tmp, s.cachePath); err != nil {
			return err
		}
	}
	return nil
}
func hashDoc(d model.KnowledgeDoc) string {
	b, _ := json.Marshal(d)
	h := sha256.Sum256(b)
	return hex.EncodeToString(h[:])
}
func cosine(a, b []float64) float64 {
	if len(a) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, aa, bb float64
	for i := range a {
		dot += a[i] * b[i]
		aa += a[i] * a[i]
		bb += b[i] * b[i]
	}
	if aa == 0 || bb == 0 {
		return 0
	}
	return dot / (math.Sqrt(aa) * math.Sqrt(bb))
}
func lexical(text string, d model.KnowledgeDoc) float64 {
	q := tokens(text)
	hay := tokens(d.Title + " " + d.Text + " " + strings.Join(d.Keywords, " "))
	if len(q) == 0 {
		return 0
	}
	hits := 0
	for t := range q {
		if _, ok := hay[t]; ok {
			hits++
		}
	}
	return float64(hits) / float64(len(q))
}
func tokens(s string) map[string]struct{} {
	m := map[string]struct{}{}
	for _, p := range strings.FieldsFunc(strings.ToLower(s), func(r rune) bool { return !unicode.IsLetter(r) && !unicode.IsDigit(r) }) {
		if len(p) >= 3 {
			m[p] = struct{}{}
		}
	}
	return m
}