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
	dir            string
	managedDir     string
	docs           []model.KnowledgeDoc
	files          map[string]string
	managed        map[string]bool
	external       map[string]string
	staticDocs     map[string]model.KnowledgeDoc
	vectors        map[string][]float64
	embedder       Embedder
	rag            bool
	cachePath      string
	allowedSources map[string]struct{}
}
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) {
	managedDir := filepath.Join(dataDir, "knowledge-managed")
	if err := os.MkdirAll(managedDir, 0o750); err != nil {
		return nil, fmt.Errorf("create managed knowledge directory: %w", err)
	}
	s := &Store{dir: dir, managedDir: managedDir, vectors: map[string][]float64{}, files: map[string]string{}, managed: map[string]bool{}, external: map[string]string{}, staticDocs: map[string]model.KnowledgeDoc{}, embedder: embedder, rag: rag, cachePath: filepath.Join(dataDir, "embeddings.json"), allowedSources: map[string]struct{}{}}
	for _, source := range allowedSources {
		s.allowedSources[strings.ToLower(strings.TrimSpace(source))] = struct{}{}
	}
	static, staticFiles, err := readDocs(dir, s.allowedSources)
	if err != nil {
		return nil, err
	}
	for i, d := range static {
		s.staticDocs[d.ID] = d
		s.files[d.ID] = staticFiles[i]
	}
	managed, managedFiles, err := readDocs(managedDir, s.allowedSources)
	if err != nil {
		return nil, err
	}
	merged := map[string]model.KnowledgeDoc{}
	order := []string{}
	for _, d := range static {
		if _, ok := merged[d.ID]; !ok {
			order = append(order, d.ID)
		}
		merged[d.ID] = d
	}
	for i, d := range managed {
		if _, ok := merged[d.ID]; !ok {
			order = append(order, d.ID)
		}
		merged[d.ID] = d
		s.files[d.ID] = managedFiles[i]
		s.managed[d.ID] = true
	}
	for _, id := range order {
		s.docs = append(s.docs, merged[id])
	}
	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 readDocs(dir string, allowed map[string]struct{}) ([]model.KnowledgeDoc, []string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, nil, fmt.Errorf("read knowledge directory %q: %w", dir, err)
	}
	var docs []model.KnowledgeDoc
	var files []string
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(strings.ToLower(e.Name()), ".json") {
			continue
		}
		path := filepath.Join(dir, e.Name())
		b, err := os.ReadFile(path)
		if err != nil {
			return nil, nil, err
		}
		var d model.KnowledgeDoc
		if err := json.Unmarshal(b, &d); err != nil {
			return nil, nil, fmt.Errorf("%s: %w", e.Name(), err)
		}
		if d.ID == "" || d.Title == "" {
			return nil, nil, fmt.Errorf("%s: id/title required", e.Name())
		}
		if !safeID(d.ID) {
			return nil, nil, fmt.Errorf("%s: invalid id %q", e.Name(), d.ID)
		}
		d.Source = strings.ToLower(strings.TrimSpace(d.Source))
		if d.Source == "" {
			return nil, 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))
		docs = append(docs, d)
		files = append(files, path)
	}
	return docs, files, 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) List() []model.KnowledgeDoc {
	if s == nil {
		return nil
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := append([]model.KnowledgeDoc(nil), s.docs...)
	sort.SliceStable(out, func(i, j int) bool { return strings.ToLower(out[i].Title) < strings.ToLower(out[j].Title) })
	return out
}

func (s *Store) Upsert(ctx context.Context, d model.KnowledgeDoc) error {
	if s == nil {
		return fmt.Errorf("knowledge store is not initialized")
	}
	d.ID = strings.TrimSpace(d.ID)
	d.Title = strings.TrimSpace(d.Title)
	d.Text = strings.TrimSpace(d.Text)
	d.Answer = strings.TrimSpace(d.Answer)
	d.Source = strings.ToLower(strings.TrimSpace(d.Source))
	d.Language = strings.TrimSpace(d.Language)
	d.CommunicationStyle = strings.ToLower(strings.TrimSpace(d.CommunicationStyle))
	if d.ID == "" || d.Title == "" {
		return fmt.Errorf("id/title required")
	}
	if !safeID(d.ID) {
		return fmt.Errorf("knowledge id may contain only letters, digits, dot, dash and underscore")
	}
	if d.Source == "" {
		return fmt.Errorf("source required")
	}
	if _, ok := s.allowedSources[d.Source]; !ok {
		return fmt.Errorf("source %q is not allowed", d.Source)
	}
	if d.Language == "" || d.CommunicationStyle == "" {
		return fmt.Errorf("language and communication_style required")
	}
	if d.MinScore < 0 || d.MinScore > 1 {
		return fmt.Errorf("min_score must be between 0 and 1")
	}
	s.mu.RLock()
	_, exists := s.files[d.ID]
	isManaged := s.managed[d.ID]
	externalSource := s.external[d.ID]
	s.mu.RUnlock()
	if externalSource != "" {
		return fmt.Errorf("externally synchronized knowledge entry %q from %q is read-only", d.ID, externalSource)
	}
	if exists && !isManaged {
		return fmt.Errorf("static knowledge entry %q is read-only; use a new id for a managed entry", d.ID)
	}
	var vector []float64
	if s.rag {
		if s.embedder == nil {
			return fmt.Errorf("RAG is enabled but no embedding provider is configured")
		}
		vv, err := s.embedder.Embed(ctx, []string{d.Title + "\n" + d.Text + "\n" + strings.Join(d.Keywords, " ")})
		if err != nil {
			return err
		}
		if len(vv) != 1 || len(vv[0]) == 0 {
			return fmt.Errorf("embedding provider returned no vector")
		}
		vector = vv[0]
	}
	path := filepath.Join(s.managedDir, d.ID+".json")
	b, err := json.MarshalIndent(d, "", "  ")
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, b, 0o640); err != nil {
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	replaced := false
	for i := range s.docs {
		if s.docs[i].ID == d.ID {
			s.docs[i] = d
			replaced = true
			break
		}
	}
	if !replaced {
		s.docs = append(s.docs, d)
	}
	s.files[d.ID] = path
	s.managed[d.ID] = true
	if s.rag {
		s.vectors[d.ID] = vector
	}
	return nil
}

func (s *Store) Delete(id string) error {
	if s == nil {
		return fmt.Errorf("knowledge store is not initialized")
	}
	id = strings.TrimSpace(id)
	if !safeID(id) {
		return fmt.Errorf("invalid knowledge id")
	}
	s.mu.RLock()
	path := s.files[id]
	isManaged := s.managed[id]
	s.mu.RUnlock()
	if path == "" {
		return os.ErrNotExist
	}
	if !isManaged {
		return fmt.Errorf("static knowledge entry %q is read-only", id)
	}
	if err := os.Remove(path); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	out := s.docs[:0]
	for _, d := range s.docs {
		if d.ID != id {
			out = append(out, d)
		}
	}
	s.docs = append([]model.KnowledgeDoc(nil), out...)
	delete(s.files, id)
	delete(s.managed, id)
	delete(s.vectors, id)
	return nil
}

func (s *Store) IsManaged(id string) bool { s.mu.RLock(); defer s.mu.RUnlock(); return s.managed[id] }
func (s *Store) Origin(id string) string {
	if s == nil {
		return ""
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.managed[id] {
		return "managed"
	}
	if src := s.external[id]; src != "" {
		return src
	}
	if _, ok := s.staticDocs[id]; ok {
		return "static"
	}
	return ""
}

// ReplaceExternalSource atomically replaces all read-only documents imported
// from one connector source. Existing vectors are reused when the normalized
// document did not change, so periodic synchronization does not re-embed the
// whole GLPI knowledge base on every run.
func (s *Store) ReplaceExternalSource(ctx context.Context, source string, docs []model.KnowledgeDoc) error {
	if s == nil {
		return fmt.Errorf("knowledge store is not initialized")
	}
	source = strings.ToLower(strings.TrimSpace(source))
	if _, ok := s.allowedSources[source]; !ok {
		return fmt.Errorf("source %q is not allowed", source)
	}

	s.mu.RLock()
	oldDocs := make(map[string]model.KnowledgeDoc, len(s.docs))
	oldVectors := make(map[string][]float64, len(s.vectors))
	for _, d := range s.docs {
		oldDocs[d.ID] = d
	}
	for id, v := range s.vectors {
		oldVectors[id] = append([]float64(nil), v...)
	}
	s.mu.RUnlock()
	cached := cacheFile{Hashes: map[string]string{}, Vectors: map[string][]float64{}}
	if b, err := os.ReadFile(s.cachePath); err == nil {
		_ = json.Unmarshal(b, &cached)
	}

	changed := make([]model.KnowledgeDoc, 0)
	seen := map[string]struct{}{}
	for i := range docs {
		d := &docs[i]
		d.ID = strings.TrimSpace(d.ID)
		d.Title = strings.TrimSpace(d.Title)
		d.Source = strings.ToLower(strings.TrimSpace(d.Source))
		if d.Source == "" {
			d.Source = source
		}
		if d.Source != source {
			return fmt.Errorf("external document %q has source %q, expected %q", d.ID, d.Source, source)
		}
		if d.ID == "" || d.Title == "" || !safeID(d.ID) {
			return fmt.Errorf("invalid external knowledge document id/title")
		}
		if _, dup := seen[d.ID]; dup {
			return fmt.Errorf("duplicate external knowledge id %q", d.ID)
		}
		seen[d.ID] = struct{}{}
		h := hashDoc(*d)
		old, ok := oldDocs[d.ID]
		same := ok && hashDoc(old) == h && len(oldVectors[d.ID]) > 0
		if !same && cached.Hashes[d.ID] == h && len(cached.Vectors[d.ID]) > 0 {
			oldVectors[d.ID] = append([]float64(nil), cached.Vectors[d.ID]...)
			same = true
		}
		if !same {
			changed = append(changed, *d)
		}
	}
	newVectors := map[string][]float64{}
	if s.rag && len(changed) > 0 {
		if s.embedder == nil {
			return fmt.Errorf("RAG is enabled but no embedding provider is configured")
		}
		texts := make([]string, len(changed))
		for i, d := range changed {
			texts[i] = d.Title + "\n" + d.Text + "\n" + strings.Join(d.Keywords, " ")
		}
		vv, err := s.embedder.Embed(ctx, texts)
		if err != nil {
			return err
		}
		if len(vv) != len(changed) {
			return fmt.Errorf("embedding provider returned %d vectors for %d documents", len(vv), len(changed))
		}
		for i, d := range changed {
			if len(vv[i]) == 0 {
				return fmt.Errorf("embedding provider returned empty vector for %s", d.ID)
			}
			newVectors[d.ID] = vv[i]
		}
	}

	s.mu.Lock()
	// Reject collisions with local/static documents.
	for _, d := range docs {
		if src := s.external[d.ID]; src == "" {
			if _, exists := oldDocs[d.ID]; exists {
				s.mu.Unlock()
				return fmt.Errorf("external knowledge id %q collides with local knowledge", d.ID)
			}
		} else if src != source {
			s.mu.Unlock()
			return fmt.Errorf("external knowledge id %q belongs to source %q", d.ID, src)
		}
	}
	rebuilt := make([]model.KnowledgeDoc, 0, len(s.docs)+len(docs))
	for _, d := range s.docs {
		if s.external[d.ID] != source {
			rebuilt = append(rebuilt, d)
		}
	}
	for id, src := range s.external {
		if src == source {
			delete(s.external, id)
			delete(s.vectors, id)
		}
	}
	for _, d := range docs {
		rebuilt = append(rebuilt, d)
		s.external[d.ID] = source
		if v := newVectors[d.ID]; len(v) > 0 {
			s.vectors[d.ID] = v
		} else if v := oldVectors[d.ID]; len(v) > 0 {
			s.vectors[d.ID] = v
		}
	}
	s.docs = rebuilt
	s.mu.Unlock()
	return s.persistVectorCache()
}

func (s *Store) persistVectorCache() error {
	if s == nil || !s.rag {
		return nil
	}
	s.mu.RLock()
	cf := cacheFile{Hashes: map[string]string{}, Vectors: map[string][]float64{}}
	for _, d := range s.docs {
		if v := s.vectors[d.ID]; len(v) > 0 {
			cf.Hashes[d.ID] = hashDoc(d)
			cf.Vectors[d.ID] = append([]float64(nil), v...)
		}
	}
	s.mu.RUnlock()
	b, err := json.MarshalIndent(cf, "", "  ")
	if err != nil {
		return err
	}
	tmp := s.cachePath + ".tmp"
	if err := os.WriteFile(tmp, b, 0o640); err != nil {
		return err
	}
	return os.Rename(tmp, s.cachePath)
}
func (s *Store) ManagedDir() string {
	if s == nil {
		return ""
	}
	return s.managedDir
}
func safeID(v string) bool {
	if v == "" {
		return false
	}
	for _, r := range v {
		if !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' || r == '.') {
			return false
		}
	}
	return !strings.Contains(v, "..")
}

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
}