diff --git a/cmd/llmgw-kb/main.go b/cmd/llmgw-kb/main.go new file mode 100644 index 0000000..0636c7f --- /dev/null +++ b/cmd/llmgw-kb/main.go @@ -0,0 +1,467 @@ +// Command llmgw-kb puts a person's documents into an llmgw knowledge base +// from their own machine (BACKLOG F2): it reads the files there, cuts them +// into chunks, has the gateway embed them with the person's key, and sends +// the text and vectors — only for what changed since the last run, which a +// ledger on the machine remembers. +// +// Its own binary, for Windows, macOS and Linux: it carries none of the +// server, agent or setup code that llmgw does. +package main + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/signal" + "path/filepath" + "strings" + "text/tabwriter" + "time" + + "github.com/spf13/cobra" + "golang.org/x/term" + + "gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract" + "gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync" + "gitea.mynodes.xyz/ilker/llmgw/internal/version" +) + +func main() { + if err := root().Execute(); err != nil { + fmt.Fprintln(os.Stderr, "error:", err) + os.Exit(1) + } +} + +// config is what `login` keeps: the gateway, the key, and for a personal key +// the project it acts in. +type config struct { + Gateway string `json:"gateway"` + Key string `json:"key"` + Project string `json:"project,omitempty"` + CA string `json:"ca,omitempty"` +} + +func configPath() (string, error) { + dir, err := os.UserConfigDir() + if err != nil { + return "", err + } + return filepath.Join(dir, "llmgw-kb", "config.json"), nil +} + +// loadConfig reads the saved configuration, with the environment over it: +// LLMGW_URL, LLMGW_API_KEY, LLMGW_PROJECT. +func loadConfig() (config, error) { + var c config + if p, err := configPath(); err == nil { + if b, err := os.ReadFile(p); err == nil { + if err := json.Unmarshal(b, &c); err != nil { + return c, fmt.Errorf("%s: %w", p, err) + } + } + } + for env, dst := range map[string]*string{"LLMGW_URL": &c.Gateway, "LLMGW_API_KEY": &c.Key, "LLMGW_PROJECT": &c.Project} { + if v := strings.TrimSpace(os.Getenv(env)); v != "" { + *dst = v + } + } + return c, nil +} + +func client(flags *globalFlags) (*kbsync.Client, error) { + c, err := loadConfig() + if err != nil { + return nil, err + } + if flags.gateway != "" { + c.Gateway = flags.gateway + } + if flags.project != "" { + c.Project = flags.project + } + if flags.ca != "" { + c.CA = flags.ca + } + if c.Gateway == "" || c.Key == "" { + return nil, errors.New("no gateway or key: run `llmgw-kb login`, or set LLMGW_URL and LLMGW_API_KEY") + } + h, err := kbsync.NewHTTP(c.CA) + if err != nil { + return nil, err + } + return &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h}, nil +} + +type globalFlags struct { + gateway, project, ca, ledger string +} + +func (g *globalFlags) openLedger(ctx context.Context) (*kbsync.Ledger, error) { + p := g.ledger + if p == "" { + var err error + if p, err = kbsync.DefaultLedgerPath(); err != nil { + return nil, err + } + } + return kbsync.OpenLedger(ctx, p) +} + +func root() *cobra.Command { + g := &globalFlags{} + cmd := &cobra.Command{ + Use: "llmgw-kb", + Short: "Put your documents into an llmgw knowledge base, from this machine", + Long: `llmgw-kb reads documents on this machine, cuts them into chunks, has the +gateway embed them with your key, and sends the text and vectors to a +knowledge base. Files are never uploaded. Run it as often as you like: a +ledger on this machine remembers what was sent, and only what changed is +read and sent again.`, + Version: version.Version, + SilenceUsage: true, + SilenceErrors: true, + } + cmd.PersistentFlags().StringVar(&g.gateway, "gateway", "", "the gateway's address (default: from login, or LLMGW_URL)") + cmd.PersistentFlags().StringVar(&g.project, "project", "", "the project a personal key acts in (default: from login, or LLMGW_PROJECT)") + cmd.PersistentFlags().StringVar(&g.ca, "ca", "", "a PEM file of certificates to trust besides the system's") + cmd.PersistentFlags().StringVar(&g.ledger, "ledger", "", "the ledger file (default: in your configuration directory)") + cmd.AddCommand(loginCmd(g), collectionsCmd(g), syncCmd(g), statusCmd(g), historyCmd(g), formatsCmd()) + return cmd +} + +func loginCmd(g *globalFlags) *cobra.Command { + return &cobra.Command{ + Use: "login", + Short: "Save the gateway's address and your key", + Long: `Asks for the gateway's address and your API key, checks them, and saves them +in your configuration directory, readable by you only. The key is typed +without being shown, so it does not end up in your shell's history.`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + c, _ := loadConfig() + in := bufio.NewReader(os.Stdin) + ask := func(q, def string) string { + if def != "" { + fmt.Fprintf(cmd.OutOrStdout(), "%s [%s]: ", q, def) + } else { + fmt.Fprintf(cmd.OutOrStdout(), "%s: ", q) + } + s, _ := in.ReadString('\n') + if s = strings.TrimSpace(s); s == "" { + return def + } + return s + } + c.Gateway = strings.TrimRight(ask("Gateway address (https://…)", or(g.gateway, c.Gateway)), "/") + fmt.Fprint(cmd.OutOrStdout(), "API key: ") + var key string + if term.IsTerminal(int(os.Stdin.Fd())) { + b, err := term.ReadPassword(int(os.Stdin.Fd())) + fmt.Fprintln(cmd.OutOrStdout()) + if err != nil { + return err + } + key = strings.TrimSpace(string(b)) + } else { + key, _ = in.ReadString('\n') + key = strings.TrimSpace(key) + } + if key != "" { + c.Key = key + } + c.Project = ask("Project, for a personal key (empty for a project's key)", or(g.project, c.Project)) + if g.ca != "" { + c.CA = g.ca + } + h, err := kbsync.NewHTTP(c.CA) + if err != nil { + return err + } + cl := &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h, Wait: time.Second} + ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second) + defer cancel() + cs, err := cl.Collections(ctx) + if err != nil { + return fmt.Errorf("the gateway refused: %w", err) + } + p, err := configPath() + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil { + return err + } + b, _ := json.MarshalIndent(c, "", " ") + if err := os.WriteFile(p, b, 0o600); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "Saved in %s. This key reaches %d collection(s); `llmgw-kb collections` lists them.\n", p, len(cs)) + return nil + }, + } +} + +func or(a, b string) string { + if a != "" { + return a + } + return b +} + +func collectionsCmd(g *globalFlags) *cobra.Command { + return &cobra.Command{ + Use: "collections", + Short: "List the collections your key may read, and which it may add to", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + c, err := client(g) + if err != nil { + return err + } + cs, err := c.Collections(cmd.Context()) + if err != nil { + return err + } + if len(cs) == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "No collections. An administrator creates them and grants access (llmgw ctl kb).") + return nil + } + tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + fmt.Fprintln(tw, "NAME\tACCESS\tOWNER\tEMBEDDED BY") + for _, x := range cs { + access := "read" + if x.Writable { + access = "read, write" + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%s (%d)\n", x.Name, access, x.Owner, x.EmbedModel, x.Dims) + } + return tw.Flush() + }, + } +} + +func syncCmd(g *globalFlags) *cobra.Command { + var o kbsync.Options + var verbose bool + var ocr string + var maxMiB int64 + cmd := &cobra.Command{ + Use: "sync COLLECTION PATH...", + Short: "Send a collection what changed in these folders and files", + Long: `Reads the files under each PATH that the collection does not already have as +they are, and sends their text and vectors. Run it again whenever you like; +unchanged files are not even read. + +Documents are named by their path under PATH, with forward slashes; with +several PATHs, each folder's own name comes first. + + --dry-run reads and cuts everything, sends nothing, and says what would + leave this machine + --delete removes documents this machine sent whose files are gone + (only under these PATHs; nothing anybody else sent) + --retry reads again files skipped before (after installing LibreOffice + or Tesseract, say) + +Without --include, every file of a known format is read (llmgw-kb formats); +hidden files, version control folders and Office lock files are not.`, + Args: cobra.MinimumNArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + c, err := client(g) + if err != nil { + return err + } + ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt) + defer stop() + l, err := g.openLedger(ctx) + if err != nil { + return err + } + defer l.Close() + o.Collection, o.Roots = args[0], args[1:] + o.Extract = extract.Options{Helpers: extract.FindHelpers(), OCRLanguages: ocr, MaxBytes: maxMiB << 20} + out := cmd.OutOrStdout() + c.Waiting = func(why string, d time.Duration) { + fmt.Fprintf(out, " .. waiting %s: %s\n", d, why) + } + start := time.Now() + rep, err := kbsync.Sync(ctx, c, l, o, func(e kbsync.Event) { + switch e.Kind { + case "unchanged": + if !verbose { + return + } + fmt.Fprintf(out, " = %s\n", e.Name) + case "sent", "would-send": + fmt.Fprintf(out, " %-8s %s (%d chunks, %s)\n", label(e.Kind), e.Name, e.Chunks, size(e.Bytes)) + default: + line := fmt.Sprintf(" %-8s %s", label(e.Kind), e.Name) + if e.Detail != "" { + line += ": " + e.Detail + } + fmt.Fprintln(out, line) + } + }) + if rep.Collection.Name != "" { + verb := "sent" + if o.DryRun { + verb = "would send" + } + fmt.Fprintf(out, "\n%s: %d files; %s %d (%d chunks, %s), %d unchanged, %d skipped, %d failed, %d deleted, in %s\n", + rep.Collection.Name, rep.Files, verb, rep.Sent, rep.Chunks, size(rep.Bytes), rep.Unchanged, + rep.Skipped, rep.Failed, rep.Deleted, time.Since(start).Round(time.Second)) + } + if err != nil { + return err + } + if rep.Failed > 0 { + return fmt.Errorf("%d document(s) failed; run again to retry them", rep.Failed) + } + return nil + }, + } + f := cmd.Flags() + f.StringSliceVar(&o.Include, "include", nil, "read only files matching these globs (\"*.pdf,docs/**/*.md\")") + f.StringSliceVar(&o.Exclude, "exclude", nil, "never read files matching these globs") + f.BoolVar(&o.Hidden, "hidden", false, "read hidden files and folders too") + f.BoolVar(&o.DryRun, "dry-run", false, "send nothing; say what would be sent") + f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone") + f.BoolVar(&o.Retry, "retry", false, "read again files skipped before") + f.IntVar(&o.Workers, "workers", 4, "files read and sent at once") + f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages") + f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB") + f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too") + return cmd +} + +func label(kind string) string { + switch kind { + case "would-send": + return "send" + case "would-delete": + return "delete" + } + return kind +} + +func size(b int64) string { + switch { + case b >= 1<<20: + return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20)) + case b >= 1<<10: + return fmt.Sprintf("%.0f KiB", float64(b)/(1<<10)) + } + return fmt.Sprintf("%d B", b) +} + +func statusCmd(g *globalFlags) *cobra.Command { + var all bool + return &cobra.Command{ + Use: "status COLLECTION", + Short: "Show what this machine has sent to a collection, and what it skipped", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + c, err := client(g) + if err != nil { + return err + } + m, err := c.Manifest(cmd.Context(), args[0]) + if err != nil { + return err + } + l, err := g.openLedger(cmd.Context()) + if err != nil { + return err + } + defer l.Close() + es, err := l.Entries(cmd.Context(), strings.TrimRight(c.Base, "/"), m.Collection.ID) + if err != nil { + return err + } + counts := map[kbsync.State]int{} + tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + fmt.Fprintln(tw, "STATE\tDOCUMENT\tCHUNKS\tWHEN\tWHY") + for _, e := range es { + counts[e.State]++ + if e.State == kbsync.Sent && !all { + continue + } + fmt.Fprintf(tw, "%s\t%s\t%d\t%s\t%s\n", e.State, e.Name, e.Chunks, e.At.Format("2006-01-02 15:04"), e.Reason) + } + _ = tw.Flush() + fmt.Fprintf(cmd.OutOrStdout(), "\n%s holds %d documents; from this machine %d sent, %d skipped, %d failed.", + m.Collection.Name, len(m.Documents), counts[kbsync.Sent], counts[kbsync.Skipped], counts[kbsync.Failed]) + if !all { + fmt.Fprint(cmd.OutOrStdout(), " (--all lists the sent ones too)") + } + fmt.Fprintln(cmd.OutOrStdout()) + return nil + }, + } +} + +func historyCmd(g *globalFlags) *cobra.Command { + return &cobra.Command{ + Use: "history", + Short: "List the latest syncs from this machine", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + l, err := g.openLedger(cmd.Context()) + if err != nil { + return err + } + defer l.Close() + runs, err := l.Runs(cmd.Context(), 20) + if err != nil { + return err + } + tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + fmt.Fprintln(tw, "STARTED\tTOOK\tCOLLECTION\tSENT\tUNCHANGED\tSKIPPED\tFAILED\tDELETED") + for _, r := range runs { + took := "unfinished" + if !r.Finished.IsZero() { + took = r.Finished.Sub(r.Started).Round(time.Second).String() + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%d\t%d\t%d\t%d\t%d\n", r.Started.Format("2006-01-02 15:04"), took, + r.Collection, r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete) + } + return tw.Flush() + }, + } +} + +func formatsCmd() *cobra.Command { + return &cobra.Command{ + Use: "formats", + Short: "List the formats read, and the helpers found on this machine", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + builtIn, helper := extract.Supported() + out := cmd.OutOrStdout() + fmt.Fprintf(out, "Read here:\n%s\n\n", wrap(builtIn)) + fmt.Fprintf(out, "Read with a helper program:\n%s\n\n", wrap(helper)) + fmt.Fprintln(out, "Helpers:") + for _, l := range extract.FindHelpers().Describe() { + fmt.Fprintln(out, " "+l) + } + return nil + }, + } +} + +func wrap(words []string) string { + var b strings.Builder + line := " " + for _, w := range words { + if len(line)+len(w) > 78 { + b.WriteString(strings.TrimRight(line, " ") + "\n") + line = " " + } + line += w + " " + } + b.WriteString(strings.TrimRight(line, " ")) + return b.String() +} diff --git a/cmd/llmgw-kb/main_test.go b/cmd/llmgw-kb/main_test.go new file mode 100644 index 0000000..127ac7a --- /dev/null +++ b/cmd/llmgw-kb/main_test.go @@ -0,0 +1,24 @@ +package main + +import ( + "os/exec" + "strings" + "testing" +) + +// The client goes onto employees' laptops: none of the server, agent, store +// or setup code may come with it, and it must build where llmgw does not. +func TestTheClientCarriesNoServerCode(t *testing.T) { + out, err := exec.Command("go", "list", "-deps", ".").Output() + if err != nil { + t.Skipf("go list: %v", err) + } + for _, dep := range strings.Fields(string(out)) { + for _, banned := range []string{"/internal/store", "/internal/router", "/internal/agent", "/internal/control", + "/internal/cli", "/internal/guard", "github.com/jackc/pgx"} { + if strings.Contains(dep, banned) { + t.Errorf("llmgw-kb depends on %s", dep) + } + } + } +} diff --git a/go.mod b/go.mod index ac921bb..e763318 100644 --- a/go.mod +++ b/go.mod @@ -35,6 +35,7 @@ require ( golang.org/x/net v0.59.0 golang.org/x/oauth2 v0.37.0 golang.org/x/sync v0.23.0 + golang.org/x/term v0.46.0 golang.org/x/text v0.42.0 google.golang.org/protobuf v1.36.12 modernc.org/sqlite v1.59.0 diff --git a/internal/kb/kbsync/api.go b/internal/kb/kbsync/api.go new file mode 100644 index 0000000..2710d9e --- /dev/null +++ b/internal/kb/kbsync/api.go @@ -0,0 +1,245 @@ +package kbsync + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "time" +) + +// Client speaks to a gateway's knowledge-base and embeddings endpoints with +// a person's key. +type Client struct { + // Base is the gateway's address, "https://llm.example.com". + Base string + Key string + // Project is the project a personal key acts in. + Project string + HTTP *http.Client + // Wait is how long a request waits out a model loading or a busy + // gateway (429/503 with Retry-After) before giving up; 0 is 15 minutes. + Wait time.Duration + // Waiting, when set, is told each time a request waits, and why. + Waiting func(why string, d time.Duration) +} + +// NewHTTP is an HTTP client trusting the system's certificates, and those in +// caFile when one is given. +func NewHTTP(caFile string) (*http.Client, error) { + tr := http.DefaultTransport.(*http.Transport).Clone() + if caFile != "" { + pem, err := os.ReadFile(caFile) + if err != nil { + return nil, err + } + pool, err := x509.SystemCertPool() + if err != nil || pool == nil { + pool = x509.NewCertPool() + } + if !pool.AppendCertsFromPEM(pem) { + return nil, fmt.Errorf("no certificate in %s", caFile) + } + tr.TLSClientConfig = &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12} + } + return &http.Client{Transport: tr, Timeout: 10 * time.Minute}, nil +} + +// Collection is a collection as the gateway shows it to this key. +type Collection struct { + ID string `json:"id"` + Name string `json:"name"` + Owner string `json:"owner"` + EmbedModel string `json:"embed_model"` + Dims int `json:"dims"` + Writable bool `json:"writable"` +} + +// RemoteDoc is a document the collection holds. +type RemoteDoc struct { + Name string `json:"name"` + Digest string `json:"digest"` + Bytes int64 `json:"bytes"` + Chunks int `json:"chunks"` + AddedAt time.Time `json:"added_at"` +} + +// Manifest is a collection and its documents. +type Manifest struct { + Collection Collection `json:"collection"` + Documents []RemoteDoc `json:"documents"` +} + +// ChunkOut is one chunk as sent. +type ChunkOut struct { + Text string `json:"text"` + Embedding []float32 `json:"embedding"` +} + +// DocOut is one document as sent. +type DocOut struct { + Name string `json:"name"` + Digest string `json:"digest"` + Bytes int64 `json:"bytes"` + Chunks []ChunkOut `json:"chunks"` +} + +// APIError is an answer that was not a success. +type APIError struct { + Status int + Message string +} + +func (e *APIError) Error() string { return fmt.Sprintf("%d: %s", e.Status, e.Message) } + +// do sends a request and decodes a JSON answer into out, waiting out answers +// that say to retry. +func (c *Client) do(ctx context.Context, method, path string, body any, out any) error { + var payload []byte + if body != nil { + var err error + if payload, err = json.Marshal(body); err != nil { + return err + } + } + wait := c.Wait + if wait <= 0 { + wait = 15 * time.Minute + } + deadline := time.Now().Add(wait) + for { + req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.Base, "/")+path, bytes.NewReader(payload)) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+c.Key) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", "llmgw-kb") + if c.Project != "" { + req.Header.Set("X-Llmgw-Project", c.Project) + } + res, err := c.HTTP.Do(req) + if err != nil { + return err + } + data, err := io.ReadAll(io.LimitReader(res.Body, 256<<20)) + res.Body.Close() + if err != nil { + return err + } + if res.StatusCode < 300 { + if out == nil { + return nil + } + return json.Unmarshal(data, out) + } + retry := res.Header.Get("Retry-After") + if retry != "" && (res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500) && time.Now().Before(deadline) { + secs, _ := strconv.Atoi(retry) + d := time.Duration(max(secs, 1)) * time.Second + if c.Waiting != nil { + c.Waiting(errorMessage(data, res.Status), d) + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(d): + } + continue + } + return &APIError{Status: res.StatusCode, Message: errorMessage(data, res.Status)} + } +} + +// errorMessage is the message in an OpenAI-shaped error, or the body. +func errorMessage(data []byte, status string) string { + var e struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal(data, &e) == nil && e.Error.Message != "" { + return e.Error.Message + } + if s := strings.TrimSpace(string(data)); s != "" && len(s) < 500 { + return s + } + return status +} + +// Collections lists what this key may read. +func (c *Client) Collections(ctx context.Context) ([]Collection, error) { + var out struct { + Data []Collection `json:"data"` + } + err := c.do(ctx, http.MethodGet, "/v1/kb/collections", nil, &out) + return out.Data, err +} + +// Manifest reads a collection's documents, by its name or id. +func (c *Client) Manifest(ctx context.Context, collection string) (Manifest, error) { + var m Manifest + err := c.do(ctx, http.MethodGet, "/v1/kb/collections/"+url.PathEscape(collection)+"/manifest", nil, &m) + return m, err +} + +// Put sends a document; changed is false when the collection already had it. +func (c *Client) Put(ctx context.Context, collection string, d DocOut) (bool, error) { + var out struct { + Changed bool `json:"changed"` + } + err := c.do(ctx, http.MethodPut, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents", d, &out) + return out.Changed, err +} + +// Delete removes a document; one already gone is not an error. +func (c *Client) Delete(ctx context.Context, collection, name string) error { + err := c.do(ctx, http.MethodDelete, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents?name="+url.QueryEscape(name), nil, nil) + var ae *APIError + if errors.As(err, &ae) && ae.Status == http.StatusNotFound && strings.HasPrefix(ae.Message, "no document") { + return nil + } + return err +} + +// embedBatch is how many chunks go in one embeddings request. +const embedBatch = 32 + +// Embed has the gateway embed texts with a model, in batches, in order. +func (c *Client) Embed(ctx context.Context, model string, texts []string) ([][]float32, error) { + out := make([][]float32, 0, len(texts)) + for start := 0; start < len(texts); start += embedBatch { + end := min(start+embedBatch, len(texts)) + var res struct { + Data []struct { + Index int `json:"index"` + Embedding []float32 `json:"embedding"` + } `json:"data"` + } + if err := c.do(ctx, http.MethodPost, "/v1/embeddings", + map[string]any{"model": model, "input": texts[start:end], "encoding_format": "float"}, &res); err != nil { + return nil, err + } + if len(res.Data) != end-start { + return nil, fmt.Errorf("the gateway embedded %d of %d texts", len(res.Data), end-start) + } + batch := make([][]float32, end-start) + for _, d := range res.Data { + if d.Index < 0 || d.Index >= len(batch) { + return nil, fmt.Errorf("an embedding came back with index %d", d.Index) + } + batch[d.Index] = d.Embedding + } + out = append(out, batch...) + } + return out, nil +} diff --git a/internal/kb/kbsync/ledger.go b/internal/kb/kbsync/ledger.go new file mode 100644 index 0000000..5e1685d --- /dev/null +++ b/internal/kb/kbsync/ledger.go @@ -0,0 +1,222 @@ +// Package kbsync is llmgw-kb's work (BACKLOG F2): it finds a person's +// documents, keeps a ledger of what it has done with each, and sends a +// knowledge base only what changed — text and vectors, never the files. +package kbsync + +import ( + "context" + "database/sql" + "fmt" + "os" + "path/filepath" + "time" + + _ "modernc.org/sqlite" // pure Go, so the client builds for every system +) + +// Ledger is the record, on this machine, of every document sent or skipped: +// which file it came from, its size, modification time and digest, and what +// happened. A file whose size and time have not changed is not read again to +// learn that it has not changed, and --delete removes only what this +// machine sent. +type Ledger struct { + db *sql.DB +} + +// State is what happened to a document. +type State string + +// The states a document is left in. +const ( + Sent State = "sent" + Skipped State = "skipped" + Failed State = "failed" +) + +// Entry is one document in the ledger. +type Entry struct { + Gateway string + Collection string + // Name is the document's name in the collection. + Name string + // File is the path on disk it came from, absolute. + File string + Size int64 + MTime time.Time + Digest string + State State + Reason string + Chunks int + At time.Time +} + +// Run is one sync, for history. +type Run struct { + ID int64 + Gateway, Collection string + Started, Finished time.Time + Sent, Unchanged, Skipped, Failed, Delete int +} + +// DefaultLedgerPath is where the ledger is kept: the user's configuration +// directory, which, unlike a cache directory, nothing cleans up. +func DefaultLedgerPath() (string, error) { + dir, err := os.UserConfigDir() + if err != nil { + return "", err + } + return filepath.Join(dir, "llmgw-kb", "ledger.db"), nil +} + +const ledgerSchema = ` +CREATE TABLE IF NOT EXISTS documents ( + gateway TEXT NOT NULL, + collection TEXT NOT NULL, + name TEXT NOT NULL, + file TEXT NOT NULL, + size INTEGER NOT NULL, + mtime_ns INTEGER NOT NULL, + digest TEXT NOT NULL, + state TEXT NOT NULL, + reason TEXT NOT NULL DEFAULT '', + chunks INTEGER NOT NULL DEFAULT 0, + at_ms INTEGER NOT NULL, + PRIMARY KEY (gateway, collection, name) +); +CREATE INDEX IF NOT EXISTS documents_by_file ON documents (file); +CREATE TABLE IF NOT EXISTS runs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + gateway TEXT NOT NULL, + collection TEXT NOT NULL, + started_ms INTEGER NOT NULL, + ended_ms INTEGER NOT NULL DEFAULT 0, + sent INTEGER NOT NULL DEFAULT 0, + unchanged INTEGER NOT NULL DEFAULT 0, + skipped INTEGER NOT NULL DEFAULT 0, + failed INTEGER NOT NULL DEFAULT 0, + deleted INTEGER NOT NULL DEFAULT 0 +);` + +// OpenLedger opens (creating) the ledger at a path. +func OpenLedger(ctx context.Context, path string) (*Ledger, error) { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return nil, err + } + // WAL and a busy timeout: two syncs at once, from two terminals, wait + // for each other rather than failing. + db, err := sql.Open("sqlite", path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(10000)&_pragma=synchronous(NORMAL)") + if err != nil { + return nil, err + } + db.SetMaxOpenConns(1) + if _, err := db.ExecContext(ctx, ledgerSchema); err != nil { + _ = db.Close() + return nil, fmt.Errorf("opening the ledger %s: %w", path, err) + } + _ = os.Chmod(path, 0o600) + return &Ledger{db: db}, nil +} + +// Close closes the ledger. +func (l *Ledger) Close() error { return l.db.Close() } + +const entryCols = `gateway, collection, name, file, size, mtime_ns, digest, state, reason, chunks, at_ms` + +func scanEntry(sc interface{ Scan(...any) error }) (Entry, error) { + var e Entry + var mtime, at int64 + var state string + err := sc.Scan(&e.Gateway, &e.Collection, &e.Name, &e.File, &e.Size, &mtime, &e.Digest, &state, &e.Reason, &e.Chunks, &at) + e.MTime, e.At, e.State = time.Unix(0, mtime), time.UnixMilli(at), State(state) + return e, err +} + +// Put records what happened to a document. +func (l *Ledger) Put(ctx context.Context, e Entry) error { + if e.At.IsZero() { + e.At = time.Now() + } + _, err := l.db.ExecContext(ctx, `INSERT INTO documents (`+entryCols+`) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (gateway, collection, name) DO UPDATE SET file = excluded.file, size = excluded.size, + mtime_ns = excluded.mtime_ns, digest = excluded.digest, state = excluded.state, + reason = excluded.reason, chunks = excluded.chunks, at_ms = excluded.at_ms`, + e.Gateway, e.Collection, e.Name, e.File, e.Size, e.MTime.UnixNano(), e.Digest, string(e.State), + e.Reason, e.Chunks, e.At.UnixMilli()) + return err +} + +// Remove forgets a document. +func (l *Ledger) Remove(ctx context.Context, gateway, collection, name string) error { + _, err := l.db.ExecContext(ctx, `DELETE FROM documents WHERE gateway = ? AND collection = ? AND name = ?`, + gateway, collection, name) + return err +} + +// Entries lists a collection's documents as this machine left them. +func (l *Ledger) Entries(ctx context.Context, gateway, collection string) ([]Entry, error) { + rows, err := l.db.QueryContext(ctx, `SELECT `+entryCols+` FROM documents WHERE gateway = ? AND collection = ? ORDER BY name`, + gateway, collection) + if err != nil { + return nil, err + } + defer rows.Close() + var out []Entry + for rows.Next() { + e, err := scanEntry(rows) + if err != nil { + return nil, err + } + out = append(out, e) + } + return out, rows.Err() +} + +// Known is the digest the ledger has for a file of this size and time, from +// any collection: a file is hashed once however many collections it goes to. +func (l *Ledger) Known(ctx context.Context, file string, size int64, mtime time.Time) (string, bool) { + var d string + err := l.db.QueryRowContext(ctx, `SELECT digest FROM documents WHERE file = ? AND size = ? AND mtime_ns = ? LIMIT 1`, + file, size, mtime.UnixNano()).Scan(&d) + return d, err == nil +} + +// StartRun records that a sync began. +func (l *Ledger) StartRun(ctx context.Context, gateway, collection string) (int64, error) { + res, err := l.db.ExecContext(ctx, `INSERT INTO runs (gateway, collection, started_ms) VALUES (?, ?, ?)`, + gateway, collection, time.Now().UnixMilli()) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +// EndRun records how a sync ended. +func (l *Ledger) EndRun(ctx context.Context, id int64, r Run) error { + _, err := l.db.ExecContext(ctx, `UPDATE runs SET ended_ms = ?, sent = ?, unchanged = ?, skipped = ?, failed = ?, deleted = ? + WHERE id = ?`, time.Now().UnixMilli(), r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete, id) + return err +} + +// Runs lists the latest syncs, newest first. +func (l *Ledger) Runs(ctx context.Context, limit int) ([]Run, error) { + rows, err := l.db.QueryContext(ctx, `SELECT id, gateway, collection, started_ms, ended_ms, sent, unchanged, skipped, failed, deleted + FROM runs ORDER BY id DESC LIMIT ?`, limit) + if err != nil { + return nil, err + } + defer rows.Close() + var out []Run + for rows.Next() { + var r Run + var s, e int64 + if err := rows.Scan(&r.ID, &r.Gateway, &r.Collection, &s, &e, &r.Sent, &r.Unchanged, &r.Skipped, &r.Failed, &r.Delete); err != nil { + return nil, err + } + r.Started = time.UnixMilli(s) + if e > 0 { + r.Finished = time.UnixMilli(e) + } + out = append(out, r) + } + return out, rows.Err() +} diff --git a/internal/kb/kbsync/sync.go b/internal/kb/kbsync/sync.go new file mode 100644 index 0000000..45bdb36 --- /dev/null +++ b/internal/kb/kbsync/sync.go @@ -0,0 +1,574 @@ +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") diff --git a/internal/kb/kbsync/sync_test.go b/internal/kb/kbsync/sync_test.go new file mode 100644 index 0000000..b8cec47 --- /dev/null +++ b/internal/kb/kbsync/sync_test.go @@ -0,0 +1,272 @@ +package kbsync + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +// fakeGateway holds one collection in memory and embeds by counting letters. +type fakeGateway struct { + mu sync.Mutex + writable bool + docs map[string]RemoteDoc + puts int + embedded int + busyFirst bool +} + +func (g *fakeGateway) serve(t *testing.T) *Client { + t.Helper() + mux := http.NewServeMux() + coll := Collection{ID: "kb_x", Name: "ik", EmbedModel: "emb", Dims: 3} + mux.HandleFunc("GET /v1/kb/collections/{id}/manifest", func(w http.ResponseWriter, r *http.Request) { + g.mu.Lock() + defer g.mu.Unlock() + if r.PathValue("id") != "ik" && r.PathValue("id") != "kb_x" { + http.Error(w, `{"error":{"message":"no collection"}}`, http.StatusNotFound) + return + } + m := Manifest{Collection: coll} + m.Collection.Writable = g.writable + for _, d := range g.docs { + m.Documents = append(m.Documents, d) + } + _ = json.NewEncoder(w).Encode(m) + }) + mux.HandleFunc("PUT /v1/kb/collections/{id}/documents", func(w http.ResponseWriter, r *http.Request) { + var d DocOut + _ = json.NewDecoder(r.Body).Decode(&d) + g.mu.Lock() + defer g.mu.Unlock() + g.puts++ + g.docs[d.Name] = RemoteDoc{Name: d.Name, Digest: d.Digest, Chunks: len(d.Chunks)} + _, _ = w.Write([]byte(`{"changed":true}`)) + }) + mux.HandleFunc("DELETE /v1/kb/collections/{id}/documents", func(w http.ResponseWriter, r *http.Request) { + g.mu.Lock() + defer g.mu.Unlock() + delete(g.docs, r.URL.Query().Get("name")) + _, _ = w.Write([]byte(`{}`)) + }) + mux.HandleFunc("POST /v1/embeddings", func(w http.ResponseWriter, r *http.Request) { + g.mu.Lock() + busy := g.busyFirst + g.busyFirst = false + g.mu.Unlock() + if busy { + w.Header().Set("Retry-After", "1") + http.Error(w, `{"error":{"message":"the model is loading"}}`, http.StatusServiceUnavailable) + return + } + var in struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&in) + type row struct { + Index int `json:"index"` + Embedding []float32 `json:"embedding"` + } + var out struct { + Data []row `json:"data"` + } + for i, s := range in.Input { + out.Data = append(out.Data, row{i, []float32{float32(len(s)), 1, 0}}) + } + g.mu.Lock() + g.embedded += len(in.Input) + g.mu.Unlock() + _ = json.NewEncoder(w).Encode(out) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + h, _ := NewHTTP("") + return &Client{Base: srv.URL, Key: "k", HTTP: h} +} + +func setup(t *testing.T) (*fakeGateway, *Client, *Ledger, string) { + t.Helper() + g := &fakeGateway{writable: true, docs: map[string]RemoteDoc{}} + c := g.serve(t) + l, err := OpenLedger(context.Background(), filepath.Join(t.TempDir(), "ledger.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = l.Close() }) + dir := t.TempDir() + write(t, dir, "politika/izin.md", "# İzin\n\nYirmi gün.\n") + write(t, dir, "politika/masraf.txt", "Masraflar finansa.\n") + write(t, dir, "eski.doc", "\xd0\xcf\x11\xe0\x00binary") + write(t, dir, "resim.bin", "\x00\x01") + write(t, dir, ".git/config", "[core]") + write(t, dir, "politika/~$izin.docx", "lock") + return g, c, l, dir +} + +func write(t *testing.T, dir, name, s string) { + t.Helper() + p := filepath.Join(dir, filepath.FromSlash(name)) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(s), 0o644); err != nil { + t.Fatal(err) + } +} + +func kinds(events []Event) map[string][]string { + out := map[string][]string{} + for _, e := range events { + out[e.Kind] = append(out[e.Kind], e.Name) + } + return out +} + +func run(t *testing.T, c *Client, l *Ledger, o Options) (Report, map[string][]string) { + t.Helper() + var mu sync.Mutex + var ev []Event + rep, err := Sync(context.Background(), c, l, o, func(e Event) { + mu.Lock() + ev = append(ev, e) + mu.Unlock() + }) + if err != nil { + t.Fatal(err) + } + return rep, kinds(ev) +} + +func TestASecondRunSendsNothingAndReadsNothing(t *testing.T) { + g, c, l, dir := setup(t) + rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}}) + if rep.Sent != 2 || len(g.docs) != 2 { + t.Fatalf("first run %+v %v", rep, k) + } + if _, ok := g.docs["politika/izin.md"]; !ok { + t.Errorf("documents %v", g.docs) + } + // The legacy Word file is skipped with the reason, and no hidden, + // lock or unknown file is touched. + if len(k["skipped"]) != 1 || k["skipped"][0] != "eski.doc" { + t.Errorf("skipped %v", k["skipped"]) + } + + puts, embedded := g.puts, g.embedded + rep, k = run(t, c, l, Options{Collection: "ik", Roots: []string{dir}}) + if rep.Sent != 0 || rep.Unchanged != 2 || g.puts != puts || g.embedded != embedded { + t.Errorf("second run %+v %v", rep, k) + } + if rep.Skipped != 1 { + t.Errorf("the skipped file was not reported as before: %+v", rep) + } + es, _ := l.Entries(context.Background(), c.Base, "kb_x") + if len(es) != 3 { + t.Errorf("ledger %+v", es) + } + runs, _ := l.Runs(context.Background(), 10) + if len(runs) != 2 || runs[0].Unchanged != 2 { + t.Errorf("runs %+v", runs) + } +} + +func TestAChangedFileIsSentAgain(t *testing.T) { + g, c, l, dir := setup(t) + run(t, c, l, Options{Collection: "ik", Roots: []string{dir}}) + old := g.docs["politika/izin.md"].Digest + write(t, dir, "politika/izin.md", "# İzin\n\nOtuz gün.\n") + // A changed time is what makes it look again; some filesystems keep + // seconds only. + later := time.Now().Add(2 * time.Second) + _ = os.Chtimes(filepath.Join(dir, "politika", "izin.md"), later, later) + rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}}) + if rep.Sent != 1 || k["sent"][0] != "politika/izin.md" || g.docs["politika/izin.md"].Digest == old { + t.Errorf("after a change %+v %v", rep, k) + } +} + +func TestDeleteRemovesOnlyWhatThisMachineSentUnderTheseRoots(t *testing.T) { + g, c, l, dir := setup(t) + run(t, c, l, Options{Collection: "ik", Roots: []string{dir}}) + // Somebody else's document, sent from another machine. + g.docs["baska/makine.md"] = RemoteDoc{Name: "baska/makine.md", Digest: "x"} + // A file of this machine from another folder. + other := t.TempDir() + write(t, other, "not.md", "başka klasör") + run(t, c, l, Options{Collection: "ik", Roots: []string{other}}) + + if err := os.Remove(filepath.Join(dir, "politika", "masraf.txt")); err != nil { + t.Fatal(err) + } + rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, Delete: true}) + if rep.Deleted != 1 || k["deleted"][0] != "politika/masraf.txt" { + t.Errorf("deleted %+v %v", rep, k) + } + for _, keep := range []string{"baska/makine.md", "not.md", "politika/izin.md"} { + if _, ok := g.docs[keep]; !ok { + t.Errorf("%s was deleted", keep) + } + } +} + +func TestADryRunSendsAndRecordsNothing(t *testing.T) { + g, c, l, dir := setup(t) + rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, DryRun: true}) + if rep.Sent != 2 || len(k["would-send"]) != 2 || rep.Bytes == 0 || rep.Chunks != 2 { + t.Errorf("dry run %+v %v", rep, k) + } + if g.puts != 0 || g.embedded != 0 { + t.Errorf("a dry run sent %d documents, embedded %d chunks", g.puts, g.embedded) + } + if es, _ := l.Entries(context.Background(), c.Base, "kb_x"); len(es) != 0 { + t.Errorf("a dry run wrote the ledger: %+v", es) + } +} + +func TestAReadOnlyCollectionIsRefusedAndALoadingModelWaitedFor(t *testing.T) { + g, c, l, dir := setup(t) + g.writable = false + if _, err := Sync(context.Background(), c, l, Options{Collection: "ik", Roots: []string{dir}}, nil); err == nil || + !strings.Contains(err.Error(), "not add") { + t.Errorf("read-only gave %v", err) + } + g.writable, g.busyFirst = true, true + var waited []string + c.Waiting = func(why string, _ time.Duration) { waited = append(waited, why) } + rep, _ := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, Workers: 1}) + if rep.Sent != 2 || len(waited) != 1 || !strings.Contains(waited[0], "loading") { + t.Errorf("after a 503 %+v, waited %v", rep, waited) + } + if _, err := Sync(context.Background(), c, l, Options{Collection: "yok", Roots: []string{dir}}, nil); err == nil { + t.Error("a collection that is not there was synced") + } +} + +func TestFiltersAndNames(t *testing.T) { + _, c, l, dir := setup(t) + write(t, dir, "rapor/2026/ocak.md", "ocak") + write(t, dir, "rapor/taslak.md", "taslak") + rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, DryRun: true, + Include: []string{"rapor/**/*.md"}, Exclude: []string{"taslak*"}}) + if rep.Sent != 1 || k["would-send"][0] != "rapor/2026/ocak.md" { + t.Errorf("filtered %+v %v", rep, k) + } + // Two roots keep their names apart. + other := t.TempDir() + write(t, other, "politika/izin.md", "aynı ad") + _, k = run(t, c, l, Options{Collection: "ik", Roots: []string{dir, other}, DryRun: true, Include: []string{"izin.md"}}) + if len(k["would-send"]) != 2 || k["would-send"][0] == k["would-send"][1] { + t.Errorf("two roots: %v", k["would-send"]) + } + if _, err := Sync(context.Background(), c, l, Options{Collection: "ik", Roots: []string{t.TempDir()}}, nil); !errors.Is(err, ErrNothing) { + t.Errorf("an empty folder gave %v", err) + } +} diff --git a/internal/router/kb_e2e_test.go b/internal/router/kb_e2e_test.go new file mode 100644 index 0000000..df3c0ee --- /dev/null +++ b/internal/router/kb_e2e_test.go @@ -0,0 +1,68 @@ +package router + +import ( + "context" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// The whole way: llmgw-kb's sync against this gateway, embeddings served by +// an engine through the ordinary path, chunks in PostgreSQL, found again by +// a search. +func TestASyncFromAMachineIsSearchable(t *testing.T) { + rt := kbRouter(t) + ctx := t.Context() + s := rt.store + sfx := strings.ToLower(store.NewID("")) + org, _ := s.CreateOrg(ctx, "e2e-"+sfx) + team, _ := s.CreateTeam(ctx, org, "t") + proj, _ := s.CreateProject(ctx, team, "p") + coll, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj, + Name: "belgeler", EmbedModel: "qwen2.5-14b-instruct", Dims: 3}) + if err != nil { + t.Fatal(err) + } + srv := httptest.NewServer(rt.Handler()) + t.Cleanup(srv.Close) + h, _ := kbsync.NewHTTP("") + c := &kbsync.Client{Base: srv.URL, Key: kbKey(t, s, proj, ""), HTTP: h} + + dir := t.TempDir() + for name, body := range map[string]string{ + "izin.md": "# Yıllık izin\n\nÇalışanlar yılda yirmi gün ücretli izin kullanır.\n", + "masraf.md": "# Masraflar\n\nMasraflar finans ekibine iletilir.\n", + } { + if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o644); err != nil { + t.Fatal(err) + } + } + l, err := kbsync.OpenLedger(ctx, filepath.Join(t.TempDir(), "ledger.db")) + if err != nil { + t.Fatal(err) + } + defer l.Close() + rep, err := kbsync.Sync(ctx, c, l, kbsync.Options{Collection: "belgeler", Roots: []string{dir}}, nil) + if err != nil || rep.Sent != 2 { + t.Fatalf("sync: %+v %v", rep, err) + } + docs, _ := s.Documents(ctx, coll.ID) + if len(docs) != 2 { + t.Fatalf("documents %+v", docs) + } + hits, err := s.SearchChunks(context.Background(), store.KBSearch{OrgID: org, Collections: []store.KBCollection{coll}, + Text: "kaç gün izin", K: 1}) + if err != nil || len(hits) != 1 || hits[0].Document != "izin.md" || !strings.Contains(hits[0].Body, "Yıllık izin\n\n") { + t.Errorf("search: %+v %v", hits, err) + } + // Again: nothing sent. + rep, err = kbsync.Sync(ctx, c, l, kbsync.Options{Collection: "belgeler", Roots: []string{dir}}, nil) + if err != nil || rep.Sent != 0 || rep.Unchanged != 2 { + t.Errorf("second sync: %+v %v", rep, err) + } +} diff --git a/internal/router/router_test.go b/internal/router/router_test.go index 9103f2f..8bb42a1 100644 --- a/internal/router/router_test.go +++ b/internal/router/router_test.go @@ -66,6 +66,20 @@ func newFakeEngine(t *testing.T, model string) *fakeEngine { w.Header().Set("Content-Type", "application/json") fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model) }) + // Three dimensions: the input's length in bytes, its number of words, + // and 1. + mux.HandleFunc("POST /v1/embeddings", func(w http.ResponseWriter, r *http.Request) { + var in struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&in) + var rows []string + for i, s := range in.Input { + rows = append(rows, fmt.Sprintf(`{"index":%d,"embedding":[%d,%d,1]}`, i, len(s), len(strings.Fields(s)))) + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"model":%q,"data":[%s],"usage":{"prompt_tokens":%d}}`, model, strings.Join(rows, ","), len(in.Input)) + }) f.srv = httptest.NewServer(mux) t.Cleanup(f.srv.Close) return f