@@ -0,0 +1,209 @@
|
||||
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, 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 err := s.index(ctx); err != nil {
|
||||
return s, err
|
||||
}
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
func (s *Store) Count() int { s.mu.RLock(); defer s.mu.RUnlock(); return len(s.docs) }
|
||||
func (s *Store) ByID(id string) (model.KnowledgeDoc, bool) {
|
||||
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) {
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadSearchesOnlyAllowedSources(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
data := filepath.Join(dir, "data")
|
||||
if err := os.MkdirAll(data, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
internal := `{"id":"I1","title":"VPN intern","text":"gateway vpn","answer":"x","source":"internal-kb","language":"de-DE","communication_style":"formal"}`
|
||||
vendor := `{"id":"V1","title":"VPN vendor","text":"gateway vpn","answer":"x","source":"vendor-docs","language":"de-DE","communication_style":"formal"}`
|
||||
if err := os.WriteFile(filepath.Join(dir, "internal.json"), []byte(internal), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, "vendor.json"), []byte(vendor), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s, err := Load(context.Background(), dir, data, nil, false, []string{"internal-kb"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s.Count() != 1 {
|
||||
t.Fatalf("count=%d", s.Count())
|
||||
}
|
||||
hits, err := s.Search(context.Background(), "vpn gateway", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) != 1 || hits[0].Doc.Source != "internal-kb" {
|
||||
t.Fatalf("unexpected hits: %+v", hits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRequiresSourceMetadata(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
data := filepath.Join(dir, "data")
|
||||
if err := os.MkdirAll(data, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, "bad.json"), []byte(`{"id":"I1","title":"Missing source"}`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := Load(context.Background(), dir, data, nil, false, []string{"internal-kb"}); err == nil {
|
||||
t.Fatal("expected missing source to fail")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user