llmgw/internal/kb/kbsync/sync.go

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
	// EmbedHere embeds the chunks on this machine even when the gateway
	// would: a vector up and down for each, rather than the text alone.
	EmbedHere bool
	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]))
	}
	// Sent as text alone when the gateway embeds it; else about ten bytes
	// of JSON a dimension with each chunk.
	serverEmbeds := s.coll.ServerEmbeds && !s.o.EmbedHere
	if !serverEmbeds {
		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
	}
	out := DocOut{Name: name, Digest: digest, Bytes: f.size, Chunks: make([]ChunkOut, len(cs))}
	if serverEmbeds {
		for i := range cs {
			out.Chunks[i] = ChunkOut{Text: texts[i]}
		}
		s.put(ctx, f, name, digest, out, len(cs), 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
	}
	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]}
	}
	s.put(ctx, f, name, digest, out, len(cs), size)
}

// put sends a document and records it.
func (s *syncer) put(ctx context.Context, f file, name, digest string, out DocOut, chunks int, size int64) {
	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, "", chunks))
	s.emit(Event{Kind: "sent", Name: name, Chunks: chunks, 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")