package kbsync import ( "context" "crypto/sha256" "encoding/hex" "errors" "fmt" "io" "io/fs" "os" "path" "path/filepath" "regexp" "sort" "strings" "sync" "time" "golang.org/x/text/unicode/norm" "gitea.mynodes.xyz/ilker/llmgw/internal/kb/chunk" "gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract" ) // Options say what to sync and how. type Options struct { // Collection is the collection's name or id. Collection string // Roots are the directories and files to sync. Roots []string // Include, when set, are the only files read (globs; one without a "/" // matches a file's name, one with it the path under its root, "**" // crossing directories). Exclude are never read. Include, Exclude []string // Hidden includes files and directories whose names start with a dot. Hidden bool // DryRun reads and cuts everything that would be sent, and sends // nothing. DryRun bool // Delete removes from the collection documents this machine sent whose // files under these roots are gone. Delete bool // Retry reads again files skipped before, as when a helper has been // installed since. Retry bool // Workers is how many files are read and sent at once; 0 is 4. Workers int Extract extract.Options Chunk chunk.Options } // Event is something a sync did to one document, for showing progress. type Event struct { // Kind is sent, unchanged, skipped, failed, deleted, would-send, // would-delete, or warning. Kind string Name string Detail string Chunks int Bytes int64 } // Report is what a sync did in all. type Report struct { Collection Collection Files int Sent int Unchanged int Skipped int Failed int Deleted int // Chunks and Bytes are what was sent, or with DryRun what would be. Chunks int Bytes int64 } // maxDocJSON is the most one document's request is, under the gateway's // limit with room to spare. const maxDocJSON = 60 << 20 // defaultExcludes are directories nobody means to put in a knowledge base. var defaultExcludes = map[string]bool{ ".git": true, ".svn": true, ".hg": true, "node_modules": true, ".venv": true, "venv": true, "__pycache__": true, ".idea": true, ".vscode": true, "$RECYCLE.BIN": true, "System Volume Information": true, ".Trash": true, ".Trashes": true, ".cache": true, } type file struct { abs string name string size int64 mtime time.Time } // Sync brings a collection up to date with the files under the roots. func Sync(ctx context.Context, c *Client, l *Ledger, o Options, progress func(Event)) (Report, error) { var mu sync.Mutex var rep Report emit := func(e Event) { mu.Lock() defer mu.Unlock() switch e.Kind { case "sent", "would-send": rep.Sent++ rep.Chunks += e.Chunks rep.Bytes += e.Bytes case "unchanged": rep.Unchanged++ case "skipped": rep.Skipped++ case "failed": rep.Failed++ case "deleted", "would-delete": rep.Deleted++ } if progress != nil { progress(e) } } m, err := c.Manifest(ctx, o.Collection) if err != nil { return rep, fmt.Errorf("reading collection %s: %w", o.Collection, err) } rep.Collection = m.Collection if !m.Collection.Writable && !o.DryRun { return rep, fmt.Errorf("this key may read collection %s but not add to it", m.Collection.Name) } remote, chunks := map[string]string{}, map[string]int{} for _, d := range m.Documents { remote[d.Name], chunks[d.Name] = d.Digest, d.Chunks } gw := strings.TrimRight(c.Base, "/") coll := m.Collection.ID files, roots, err := walk(o, func(name, why string) { emit(Event{Kind: "warning", Name: name, Detail: why}) }) if err != nil { return rep, err } rep.Files = len(files) if len(files) == 0 { return rep, ErrNothing } runID, err := l.StartRun(ctx, gw, coll) if err != nil { return rep, err } previous, err := l.Entries(ctx, gw, coll) if err != nil { return rep, err } byFile := map[string][]Entry{} for _, e := range previous { byFile[e.File] = append(byFile[e.File], e) } s := &syncer{c: c, l: l, o: o, gw: gw, coll: m.Collection, remote: remote, chunks: chunks, byFile: byFile, emit: emit} workers := o.Workers if workers <= 0 { workers = 4 } jobs := make(chan file) var wg sync.WaitGroup for range workers { wg.Add(1) go func() { defer wg.Done() for f := range jobs { s.one(ctx, f) } }() } for _, f := range files { select { case jobs <- f: case <-ctx.Done(): } if ctx.Err() != nil { break } } close(jobs) wg.Wait() if o.Delete && ctx.Err() == nil { seen := map[string]bool{} for _, f := range files { seen[f.abs] = true } for _, e := range previous { if e.State != Sent || seen[e.File] || !under(e.File, roots) { continue } if o.DryRun { emit(Event{Kind: "would-delete", Name: e.Name}) continue } if err := c.Delete(ctx, coll, e.Name); err != nil { emit(Event{Kind: "failed", Name: e.Name, Detail: "deleting: " + err.Error()}) continue } _ = l.Remove(ctx, gw, coll, e.Name) emit(Event{Kind: "deleted", Name: e.Name}) } } _ = l.EndRun(context.WithoutCancel(ctx), runID, Run{Sent: rep.Sent, Unchanged: rep.Unchanged, Skipped: rep.Skipped, Failed: rep.Failed, Delete: rep.Deleted}) return rep, ctx.Err() } type syncer struct { c *Client l *Ledger o Options gw string coll Collection remote map[string]string chunks map[string]int byFile map[string][]Entry emit func(Event) } func (s *syncer) entry(f file, name, digest string, st State, reason string, chunks int) Entry { return Entry{Gateway: s.gw, Collection: s.coll.ID, Name: name, File: f.abs, Size: f.size, MTime: f.mtime, Digest: digest, State: st, Reason: reason, Chunks: chunks} } // one syncs one file: every document in it. func (s *syncer) one(ctx context.Context, f file) { if ctx.Err() != nil { return } digest, ok := s.l.Known(ctx, f.abs, f.size, f.mtime) if !ok { d, err := hashFile(f.abs) if err != nil { s.emit(Event{Kind: "skipped", Name: f.name, Detail: err.Error()}) return } digest = d } prev := s.byFile[f.abs] // The collection has this file as it is: nothing to read. if s.remote[f.name] == digest { s.record(ctx, s.entry(f, f.name, digest, Sent, "", s.chunks[f.name])) s.emit(Event{Kind: "unchanged", Name: f.name}) return } // An archive or a mail is several documents; all of them as they are // is nothing to read either. if len(prev) > 0 && allSame(prev, digest, s.remote) && !(s.o.Retry && anySkipped(prev)) { for _, e := range prev { if e.State == Skipped { s.emit(Event{Kind: "skipped", Name: e.Name, Detail: e.Reason + " (as before; --retry reads it again)"}) } else { s.emit(Event{Kind: "unchanged", Name: e.Name}) } } return } parts, err := extract.File(ctx, f.abs, s.o.Extract) if err != nil { s.skip(ctx, f, f.name, digest, err) return } current := map[string]bool{} for _, p := range parts { name := f.name if p.Name != "" { name = f.name + "!/" + p.Name } current[name] = true if p.Err != nil { s.skip(ctx, f, name, digest, p.Err) continue } s.send(ctx, f, name, digest, p.Text) } // Members an archive no longer has go with it. for _, e := range prev { if current[e.Name] || e.Name == f.name { continue } if s.o.DryRun { s.emit(Event{Kind: "would-delete", Name: e.Name, Detail: "no longer in " + f.name}) continue } if e.State == Sent { if err := s.c.Delete(ctx, s.coll.ID, e.Name); err != nil { s.emit(Event{Kind: "failed", Name: e.Name, Detail: "deleting: " + err.Error()}) continue } s.emit(Event{Kind: "deleted", Name: e.Name, Detail: "no longer in " + f.name}) } _ = s.l.Remove(ctx, s.gw, s.coll.ID, e.Name) } } // allSame reports whether every document from a file was left with this // digest, and the collection still holds those that were sent. func allSame(prev []Entry, digest string, remote map[string]string) bool { for _, e := range prev { if e.Digest != digest { return false } switch e.State { case Sent: if remote[e.Name] != digest { return false } case Failed: return false } } return true } func anySkipped(es []Entry) bool { for _, e := range es { if e.State == Skipped { return true } } return false } func (s *syncer) skip(ctx context.Context, f file, name, digest string, err error) { s.record(ctx, s.entry(f, name, digest, Skipped, err.Error(), 0)) s.emit(Event{Kind: "skipped", Name: name, Detail: err.Error()}) } // record writes the ledger; a dry run leaves it as it was. func (s *syncer) record(ctx context.Context, e Entry) { if s.o.DryRun { return } _ = s.l.Put(ctx, e) } // send cuts a document's text, embeds it and sends it. func (s *syncer) send(ctx context.Context, f file, name, digest, text string) { if s.remote[name] == digest { s.record(ctx, s.entry(f, name, digest, Sent, "", s.chunks[name])) s.emit(Event{Kind: "unchanged", Name: name}) return } cs := chunk.Split(text, s.o.Chunk) if len(cs) == 0 { s.skip(ctx, f, name, digest, extract.ErrNoText) return } texts := make([]string, len(cs)) var size int64 for i, c := range cs { texts[i] = c.Text() size += int64(len(texts[i])) } // About ten bytes of JSON a dimension. size += int64(len(cs) * s.coll.Dims * 10) if size > maxDocJSON { s.fail(ctx, f, name, digest, fmt.Errorf("%d MiB to send is over the gateway's limit for one document; split the file", size>>20)) return } if s.o.DryRun { s.emit(Event{Kind: "would-send", Name: name, Chunks: len(cs), Bytes: size}) return } vecs, err := s.c.Embed(ctx, s.coll.EmbedModel, texts) if err != nil { s.fail(ctx, f, name, digest, fmt.Errorf("embedding with %s: %w", s.coll.EmbedModel, err)) return } out := DocOut{Name: name, Digest: digest, Bytes: f.size, Chunks: make([]ChunkOut, len(cs))} for i := range cs { if len(vecs[i]) != s.coll.Dims { s.fail(ctx, f, name, digest, fmt.Errorf("%s gave %d dimensions and the collection takes %d", s.coll.EmbedModel, len(vecs[i]), s.coll.Dims)) return } out.Chunks[i] = ChunkOut{Text: texts[i], Embedding: vecs[i]} } if _, err := s.c.Put(ctx, s.coll.ID, out); err != nil { s.fail(ctx, f, name, digest, err) return } s.record(ctx, s.entry(f, name, digest, Sent, "", len(cs))) s.emit(Event{Kind: "sent", Name: name, Chunks: len(cs), Bytes: size}) } func (s *syncer) fail(ctx context.Context, f file, name, digest string, err error) { s.record(ctx, s.entry(f, name, digest, Failed, err.Error(), 0)) s.emit(Event{Kind: "failed", Name: name, Detail: err.Error()}) } func hashFile(p string) (string, error) { fh, err := os.Open(p) if err != nil { return "", err } defer fh.Close() h := sha256.New() if _, err := io.Copy(h, fh); err != nil { return "", err } return "sha256:" + hex.EncodeToString(h.Sum(nil)), nil } // under reports whether a path is one of the roots or inside one. func under(p string, roots []string) bool { for _, r := range roots { if p == r || strings.HasPrefix(p, strings.TrimRight(r, string(filepath.Separator))+string(filepath.Separator)) { return true } } return false } // walk lists the files to sync, named as documents: the path under its root // with forward slashes, in Unicode NFC, prefixed by the root's own name when // there are several roots so that two roots' "README.md" stay apart. func walk(o Options, warn func(name, why string)) ([]file, []string, error) { inc, err := globs(o.Include) if err != nil { return nil, nil, err } exc, err := globs(o.Exclude) if err != nil { return nil, nil, err } var out []file var roots []string seen := map[string]string{} for _, root := range o.Roots { abs, err := filepath.Abs(root) if err != nil { return nil, nil, err } st, err := os.Stat(abs) if err != nil { return nil, nil, err } roots = append(roots, abs) prefix := "" if len(o.Roots) > 1 && st.IsDir() { prefix = norm.NFC.String(filepath.Base(abs)) + "/" } add := func(p string, info fs.FileInfo) { rel := filepath.Base(p) if st.IsDir() { rel, _ = filepath.Rel(abs, p) } name := prefix + norm.NFC.String(filepath.ToSlash(rel)) if !wanted(name, inc, exc, o.Hidden) { return } if prev, ok := seen[strings.ToLower(name)]; ok && prev != name { warn(name, "differs from "+prev+" only in case; on Windows and macOS they are one file") } seen[strings.ToLower(name)] = name out = append(out, file{abs: p, name: name, size: info.Size(), mtime: info.ModTime()}) } if !st.IsDir() { add(abs, st) continue } err = filepath.WalkDir(abs, func(p string, d fs.DirEntry, err error) error { if err != nil { // A directory this user may not read is skipped, not fatal. warn(p, err.Error()) if d != nil && d.IsDir() { return filepath.SkipDir } return nil } base := d.Name() if d.IsDir() { if p != abs && (defaultExcludes[base] || (!o.Hidden && strings.HasPrefix(base, "."))) { return filepath.SkipDir } return nil } if !d.Type().IsRegular() { return nil } info, err := d.Info() if err != nil { return nil } add(p, info) return nil }) if err != nil { return nil, nil, err } } sort.Slice(out, func(i, j int) bool { return out[i].name < out[j].name }) return out, roots, nil } // wanted applies the filters to a document name. func wanted(name string, inc, exc []*regexp.Regexp, hidden bool) bool { base := path.Base(name) if !hidden && strings.HasPrefix(base, ".") { return false } // Office's lock files, LibreOffice's, editors' backups. if strings.HasPrefix(base, "~$") || strings.HasPrefix(base, ".~lock.") || strings.HasSuffix(base, "~") || base == "Thumbs.db" || base == "desktop.ini" || base == ".DS_Store" { return false } for _, re := range exc { if re.MatchString(name) || re.MatchString(base) { return false } } if len(inc) > 0 { for _, re := range inc { if re.MatchString(name) || re.MatchString(base) { return true } } return false } return extract.Known(base) } // globs compiles glob patterns: "*" within a name, "**" across directories, // "?" one character; case-insensitive, as names are on Windows and macOS. func globs(patterns []string) ([]*regexp.Regexp, error) { var out []*regexp.Regexp for _, p := range patterns { for _, one := range strings.Split(p, ",") { one = strings.TrimSpace(filepath.ToSlash(one)) if one == "" { continue } var b strings.Builder b.WriteString("(?i)^") for i := 0; i < len(one); i++ { switch c := one[i]; c { case '*': if i+1 < len(one) && one[i+1] == '*' { b.WriteString(".*") i++ if i+1 < len(one) && one[i+1] == '/' { b.WriteString("/?") i++ } } else { b.WriteString("[^/]*") } case '?': b.WriteString("[^/]") default: b.WriteString(regexp.QuoteMeta(string(c))) } } b.WriteString("$") re, err := regexp.Compile(b.String()) if err != nil { return nil, fmt.Errorf("pattern %q: %w", one, err) } out = append(out, re) } } return out, nil } // ErrNothing is a sync that found no file to read. var ErrNothing = errors.New("no files to read under these paths")