llmgw-kb: sync a person's documents into a knowledge base from their own machine

Its own binary for Windows, macOS and Linux, carrying no server code. It reads
the files there, cuts them, has the gateway embed them with the person's key
and sends text and vectors for what changed only: a SQLite ledger on the
machine remembers each document's file, size, time, digest and what happened,
so a rerun reads nothing unchanged, a file skipped stays skipped until
--retry, and --delete removes only what this machine sent under the paths
given. --dry-run says what would leave the machine. Commands: login,
collections, sync, status, history, formats.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-10-08 13:43:47 +02:00
1 parent d9e030e0b6
commit ba5df90602
9 files changed
+1887

No files matched your search

+467
View File
@@ -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()
}
+24
View File
@@ -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)
}
}
}
}
+1
View File
@@ -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
+245
View File
@@ -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
}
+222
View File
@@ -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()
}
+574
View File
@@ -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")
+272
View File
@@ -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)
}
}
+68
View File
@@ -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)
}
}
+14
View File
@@ -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