llmgw/cmd/llmgw-kb/main.go

// 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, or(x.OwnerName, 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.BoolVar(&o.EmbedHere, "embed-here", false,
		"embed on this machine and send the vectors, even to a gateway that embeds text it is sent")
	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()
}