llmgw/internal/cli/kb.go

package cli

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strings"

	"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, ", "))
			}
			return w.Flush()
		},
	})

	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 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
}