llmgw/internal/cli/kb.go

package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)

func newKBCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "kb",
		Short: "Knowledge bases: collections of documents requests are answered from",
		Long: "A collection belongs to an organization, team, project or person, and is\n" +
			"embedded by one model of the catalog. People put documents into it from\n" +
			"their own machines with llmgw-kb (its own download, for Windows, macOS and\n" +
			"Linux), which sends text and vectors, never the files.\n\n" +
			"Who may read or add to a collection: its owner and everything below it,\n" +
			"and whoever it is granted to (grant), never outside its organization.\n" +
			"Targets are written org:Acme, team:Acme/hr, project:Acme/hr/docs (or with\n" +
			"IDs), key:KEY-ID or user:USER-ID. A collection is named by its ID, its name\n" +
			"when that is unique, or ORG/NAME.\n\n" +
			"Needs the PostgreSQL store with pgvector; `llmgw ctl kb list` says why not.",
	}
	c.AddCommand(&cobra.Command{
		Use:   "list",
		Short: "Every collection: owner, model, documents and grants",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var v adminapi.KBView
			if err := client().get(cmd.Context(), "/api/v1/kb", &v); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
			}
			out := cmd.OutOrStdout()
			if !v.Available {
				fmt.Fprintln(out, "Knowledge bases are off: "+v.Reason)
				return nil
			}
			if len(v.Collections) == 0 {
				fmt.Fprintln(out, "No collections yet (llmgw ctl kb create).")
				return nil
			}
			w := tab()
			fmt.Fprintln(w, "COLLECTION\tOWNER\tMODEL\tDOCUMENTS\tCHUNKS\tGRANTED TO")
			for _, c := range v.Collections {
				var grants []string
				for _, g := range c.Grants {
					grants = append(grants, or(g.Name, g.Target)+" ("+g.Perm+")")
				}
				fmt.Fprintf(w, "%s/%s\t%s\t%s (%d)\t%d\t%d\t%s\n", or(c.OrgName, c.Org), c.Name, or(c.OwnerName, c.Owner),
					c.EmbedModel, c.Dims, c.Documents, c.Chunks, strings.Join(grants, ", "))
			}
			if err := w.Flush(); err != nil {
				return err
			}
			if len(v.Defaults) > 0 {
				fmt.Fprintln(out, "\nAnswered from by default, when a request names no collection:")
				w = tab()
				for _, d := range v.Defaults {
					fmt.Fprintf(w, "  %s\t%s\n", or(d.Name, d.Target), strings.Join(d.Collections, ", "))
				}
				return w.Flush()
			}
			return nil
		},
	})

	var org string
	var dims int
	create := &cobra.Command{
		Use:   "create OWNER NAME MODEL",
		Short: "Make a collection owned by an organization, team, project or person",
		Long: "MODEL is a catalog model with capabilities.embeddings. Its dimension is asked\n" +
			"of it, so it has to be able to serve now; --dims gives it instead.",
		Example: "  llmgw ctl kb create team:Acme/hr policies bge-m3\n" +
			"  llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme",
		Args: cobra.ExactArgs(3),
		RunE: func(cmd *cobra.Command, args []string) error {
			body, _ := json.Marshal(adminapi.KBCreateRequest{Owner: args[0], Name: args[1], Model: args[2], Org: org, Dims: dims})
			var v adminapi.KBCollectionView
			if err := client().do(cmd.Context(), "POST", "/api/v1/kb/collections", json.RawMessage(body), &v); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "collection %s (%s), embedded by %s in %d dimensions\n", v.Name, v.ID, v.EmbedModel, v.Dims)
			fmt.Fprintf(cmd.OutOrStdout(), "people with access fill it with: llmgw-kb sync %s FOLDER\n", v.Name)
			return nil
		},
	}
	create.Flags().StringVar(&org, "org", "", "the organization of a person's collection")
	create.Flags().IntVar(&dims, "dims", 0, "the model's embedding dimension, instead of asking it")
	c.AddCommand(create)

	c.AddCommand(&cobra.Command{
		Use:   "delete COLLECTION",
		Short: "Remove a collection, its documents and its grants",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := client().do(cmd.Context(), "DELETE", "/api/v1/kb/collections/"+url.PathEscape(args[0]), nil, nil); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "collection %s removed\n", args[0])
			return nil
		},
	})

	for _, verb := range []struct {
		use, short, method string
		perm               bool
	}{
		{"grant", "Let a scope and everything below it read a collection, or with --write add to it", "POST", true},
		{"revoke", "Take a grant away", "DELETE", false},
	} {
		var write bool
		cmd := &cobra.Command{
			Use:     verb.use + " COLLECTION TARGET",
			Short:   verb.short,
			Example: "  llmgw ctl kb " + verb.use + " Acme/policies project:Acme/sales/crm",
			Args:    cobra.ExactArgs(2),
			RunE: func(cmd *cobra.Command, args []string) error {
				req := adminapi.KBGrantRequest{Collection: args[0], Target: args[1], Perm: "read"}
				if write {
					req.Perm = "write"
				}
				body, _ := json.Marshal(req)
				if err := client().do(cmd.Context(), verb.method, "/api/v1/kb/grants", json.RawMessage(body), nil); err != nil {
					return err
				}
				if verb.perm {
					fmt.Fprintf(cmd.OutOrStdout(), "%s may %s %s\n", args[1], req.Perm, args[0])
				} else {
					fmt.Fprintf(cmd.OutOrStdout(), "%s no longer granted %s\n", args[1], args[0])
				}
				return nil
			},
		}
		if verb.perm {
			cmd.Flags().BoolVar(&write, "write", false, "add and remove documents too")
		}
		c.AddCommand(cmd)
	}

	c.AddCommand(&cobra.Command{
		Use:   "docs COLLECTION",
		Short: "A collection's documents",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var docs []adminapi.KBDocumentView
			if err := client().get(cmd.Context(), "/api/v1/kb/collections/"+url.PathEscape(args[0])+"/documents", &docs); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(docs)
			}
			w := tab()
			fmt.Fprintln(w, "DOCUMENT\tCHUNKS\tSIZE\tADDED\tBY")
			for _, d := range docs {
				fmt.Fprintf(w, "%s\t%d\t%s\t%s\t%s\n", d.Name, d.Chunks, humanBytes(d.Bytes), d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
			}
			return w.Flush()
		},
	})
	c.AddCommand(&cobra.Command{
		Use:   "rm COLLECTION DOCUMENT",
		Short: "Remove one document from a collection",
		Args:  cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			path := "/api/v1/kb/collections/" + url.PathEscape(args[0]) + "/documents?name=" + url.QueryEscape(args[1])
			if err := client().do(cmd.Context(), "DELETE", path, nil, nil); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed from %s\n", args[1], args[0])
			return nil
		},
	})

	var clear bool
	def := &cobra.Command{
		Use:   "default TARGET [COLLECTION...]",
		Short: "Answer a part of the tree's chat requests from collections, when they name none",
		Long: "Requests from TARGET and everything below it are then answered from these\n" +
			"collections without naming them (X-Llmgw-Knowledge). The narrowest default\n" +
			"wins: a key's over its project's, a project's over its team's. A request's own\n" +
			"header overrides it, and X-Llmgw-Knowledge: none turns it off. A collection a\n" +
			"caller may not read is left out for them, and a search that fails answers\n" +
			"without passages rather than refusing. In force within 30 seconds.",
		Example: "  llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms\n" +
			"  llmgw ctl kb default team:Acme/hr --clear",
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if clear != (len(args) == 1) {
				return fmt.Errorf("name collections, or --clear")
			}
			body, _ := json.Marshal(adminapi.KBDefaultRequest{Target: args[0], Collections: args[1:]})
			if err := client().do(cmd.Context(), "POST", "/api/v1/kb/defaults", json.RawMessage(body), nil); err != nil {
				return err
			}
			if clear {
				fmt.Fprintf(cmd.OutOrStdout(), "%s answers from no collection by default\n", args[0])
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s answers from %s by default\n", args[0], strings.Join(args[1:], ", "))
			}
			return nil
		},
	}
	def.Flags().BoolVar(&clear, "clear", false, "no default any more")
	c.AddCommand(def)

	var dsModel, dsOut string
	var dsPer, dsMax int
	dataset := &cobra.Command{
		Use:   "dataset COLLECTION --model MODEL",
		Short: "Write training data from a collection: questions and answers, for an adapter",
		Long: "A chat model of the cluster reads each chunk of the collection and writes\n" +
			"questions a person might ask of it, with the answers it holds, one JSON object\n" +
			"a line ({\"q\", \"a\", \"doc\", \"ord\"}). Read it through and correct it before\n" +
			"training on it (llmgw ctl train): what the model got wrong, an adapter learns.\n\n" +
			"An adapter learns how answers sound — the organisation's words, format and\n" +
			"tone — far better than what they say; the facts stay the knowledge base's.",
		Example: "  llmgw ctl kb dataset Acme/policies --model qwen2.5-coder-14b-instruct --out policies.jsonl",
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			w := cmd.OutOrStdout()
			if dsOut != "" {
				f, err := os.Create(dsOut)
				if err != nil {
					return err
				}
				defer f.Close()
				w = f
			}
			lw := &datasetWriter{pairs: w, problems: cmd.ErrOrStderr()}
			body, _ := json.Marshal(adminapi.KBDatasetRequest{Model: dsModel, PerChunk: dsPer, MaxChunks: dsMax})
			err := client().patient(6*time.Hour).do(cmd.Context(), "POST",
				"/api/v1/kb/collections/"+url.PathEscape(args[0])+"/dataset", json.RawMessage(body), lw)
			lw.flush()
			fmt.Fprintf(cmd.ErrOrStderr(), "%d pairs from %d chunks; %d chunks gave none\n", lw.n, len(lw.chunks), lw.failed)
			return err
		},
	}
	dataset.Flags().StringVar(&dsModel, "model", "", "the chat model that writes the questions and answers")
	dataset.Flags().StringVar(&dsOut, "out", "", "the file to write (default: standard output)")
	dataset.Flags().IntVar(&dsPer, "per-chunk", 8, "pairs from each chunk")
	dataset.Flags().IntVar(&dsMax, "max-chunks", 500, "the most chunks read")
	_ = dataset.MarkFlagRequired("model")
	c.AddCommand(dataset)

	var k int
	search := &cobra.Command{
		Use:   "search COLLECTION QUERY...",
		Short: "Try a search of a collection, as an administrator (no tenancy applies)",
		Args:  cobra.MinimumNArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			body, _ := json.Marshal(adminapi.KBSearchRequest{Collection: args[0], Query: strings.Join(args[1:], " "), K: k})
			var hits []adminapi.KBHitView
			if err := client().do(cmd.Context(), "POST", "/api/v1/kb/search", json.RawMessage(body), &hits); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(hits)
			}
			out := cmd.OutOrStdout()
			if len(hits) == 0 {
				fmt.Fprintln(out, "Nothing found.")
			}
			for i, h := range hits {
				fmt.Fprintf(out, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(h.Text), "   "))
			}
			return nil
		},
	}
	search.Flags().IntVarP(&k, "k", "k", 5, "how many chunks")
	c.AddCommand(search)
	return c
}

// datasetWriter splits the dataset stream: pairs to the file, problems to
// the terminal, counted as they pass.
type datasetWriter struct {
	pairs, problems io.Writer
	buf             []byte
	n, failed       int
	chunks          map[string]bool
}

func (d *datasetWriter) Write(p []byte) (int, error) {
	d.buf = append(d.buf, p...)
	for {
		i := bytes.IndexByte(d.buf, '\n')
		if i < 0 {
			return len(p), nil
		}
		d.line(d.buf[:i])
		d.buf = d.buf[i+1:]
	}
}

func (d *datasetWriter) flush() {
	if len(bytes.TrimSpace(d.buf)) > 0 {
		d.line(d.buf)
	}
	d.buf = nil
}

func (d *datasetWriter) line(b []byte) {
	var l adminapi.KBDatasetLine
	if json.Unmarshal(b, &l) != nil {
		return
	}
	if d.chunks == nil {
		d.chunks = map[string]bool{}
	}
	d.chunks[fmt.Sprintf("%s#%d", l.Doc, l.Ord)] = true
	if l.Error != "" {
		d.failed++
		fmt.Fprintf(d.problems, "  %s #%d: %s\n", l.Doc, l.Ord, l.Error)
		return
	}
	d.n++
	out, _ := json.Marshal(l)
	fmt.Fprintln(d.pairs, string(out))
}