llmgw/internal/cli/tui_kb.go

package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"

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

// Knowledge bases, as the console shows them: each collection with its
// owner, model, documents and grants; Enter for what can be done to one, A
// to make one. People fill them from their own machines with llmgw-kb.

type kbState struct {
	view adminapi.KBView
	err  error
	// loaded is set once an answer came, so "none yet" is not shown while
	// asking.
	loaded bool
}

type consoleKBMsg struct {
	view adminapi.KBView
	err  error
}

func (m consoleModel) fetchKB() tea.Cmd {
	c := m.client
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		var msg consoleKBMsg
		msg.err = c.get(ctx, "/api/v1/kb", &msg.view)
		return msg
	}
}

func (m consoleModel) kbPanel(pw int) (string, string, string) {
	keys := "→ or Enter open · ↑↓ choose · Q leave"
	if m.inPanel {
		keys = "↑↓ collection · Enter actions · A new collection · ← back · Q leave"
	}
	k := m.kb
	var b strings.Builder
	b.WriteString("\n  " + sDim.Render("Documents requests are answered from; people add them with llmgw-kb.") + "\n\n")
	switch {
	case k.err != nil:
		b.WriteString("  " + sBad.Render(k.err.Error()))
		return "Knowledge", b.String(), keys
	case !k.loaded:
		b.WriteString("  " + sDim.Render("asking the control plane…"))
		return "Knowledge", b.String(), keys
	case !k.view.Available:
		b.WriteString("  " + sBad.Render("Off: ") + k.view.Reason)
		return "Knowledge", b.String(), keys
	case len(k.view.Collections) == 0:
		b.WriteString("  " + sDim.Render("No collections yet. A to make one."))
		return "Knowledge", b.String(), keys
	}
	var rows [][]string
	for _, c := range k.view.Collections {
		var grants []string
		for _, g := range c.Grants {
			grants = append(grants, or(g.Name, g.Target)+" "+g.Perm)
		}
		rows = append(rows, []string{or(c.OrgName, c.Org) + "/" + c.Name, or(c.OwnerName, c.Owner),
			fmt.Sprintf("%s (%d)", c.EmbedModel, c.Dims), fmt.Sprint(c.Documents), fmt.Sprint(c.Chunks), dash(strings.Join(grants, ", "))})
	}
	b.WriteString(table([]string{"COLLECTION", "OWNER", "MODEL", "DOCUMENTS", "CHUNKS", "GRANTED TO"}, rows, m.row, pw, m.inPanel))
	b.WriteString("\n\n  " + sDim.Render("On a person's machine: llmgw-kb sync COLLECTION FOLDER"))
	return "Knowledge", b.String(), keys
}

// kbRef is how the admin API is told which collection: by its ID.
func kbRef(c adminapi.KBCollectionView) string { return url.PathEscape(c.ID) }

// kbCreateForm makes a collection.
func (m *consoleModel) kbCreateForm() {
	model := ""
	for _, md := range m.data.models {
		n := strings.ToLower(md.Name)
		if strings.Contains(n, "embed") || strings.Contains(n, "bge") || strings.Contains(n, "e5") {
			model = md.Name
			break
		}
	}
	m.openForm(&formDialog{title: "New knowledge-base collection",
		note: "Owned by part of the tree: it and everything below it may read and add to it.\n" +
			"The model must embed (capabilities.embeddings) and be able to serve now.",
		fields: []formField{
			choiceF("owner", []string{"project", "team", "org", "user"}, "team"),
			textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
			textF("org", "", "for a person's collection only"),
			textF("name", "", "how people name it to llmgw-kb"),
			textF("model", model, "the catalog model that embeds it"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			for i, what := range map[int]string{1: "which " + v[0], 3: "a name", 4: "a model"} {
				if strings.TrimSpace(v[i]) == "" {
					return nil, fmt.Errorf("%s?", what)
				}
			}
			req := adminapi.KBCreateRequest{Owner: v[0] + ":" + v[1], Org: v[2], Name: v[3], Model: v[4]}
			return m.call("New collection "+v[3], http.MethodPost, "/api/v1/kb/collections", req, func(raw json.RawMessage) string {
				var c adminapi.KBCollectionView
				_ = json.Unmarshal(raw, &c)
				return fmt.Sprintf("Collection %s, embedded by %s in %d dimensions.\n\nPeople with access fill it with:\n  llmgw-kb sync %s FOLDER",
					c.Name, c.EmbedModel, c.Dims, c.Name)
			}), nil
		}})
}

// kbMenu is what can be done to one collection.
func (m consoleModel) kbMenu(c adminapi.KBCollectionView) *menuDialog {
	name := or(c.OrgName, c.Org) + "/" + c.Name
	items := []menuItem{
		{"Documents", func(m *consoleModel) tea.Cmd { return m.kbDocuments(c) }},
		{"Grant…", func(m *consoleModel) tea.Cmd {
			m.openForm(&formDialog{title: "Grant " + name,
				note: "It and everything below it may then read, or with write add to, this collection.\nNever outside " + or(c.OrgName, c.Org) + ".",
				fields: []formField{
					choiceF("to", []string{"project", "team", "org", "user", "key"}, "project"),
					textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
					choiceF("may", []string{"read", "write"}, "read"),
				},
				submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
					if strings.TrimSpace(v[1]) == "" {
						return nil, fmt.Errorf("which %s?", v[0])
					}
					req := adminapi.KBGrantRequest{Collection: c.ID, Target: v[0] + ":" + strings.TrimSpace(v[1]), Perm: v[2]}
					return m.call("Grant "+c.Name, http.MethodPost, "/api/v1/kb/grants", req, nil), nil
				}})
			return nil
		}},
	}
	for _, g := range c.Grants {
		items = append(items, menuItem{"Revoke " + or(g.Name, g.Target) + " (" + g.Perm + ")…", func(m *consoleModel) tea.Cmd {
			m.confirmD = &confirmDialog{title: "Revoke a grant on " + name,
				text: or(g.Name, g.Target) + " may then no longer " + g.Perm + " it, unless it owns it.",
				yes: func(m *consoleModel) tea.Cmd {
					return m.call("Revoke", http.MethodDelete, "/api/v1/kb/grants",
						adminapi.KBGrantRequest{Collection: c.ID, Target: g.Target}, nil)
				}}
			return nil
		}})
	}
	items = append(items,
		menuItem{"Try a search…", func(m *consoleModel) tea.Cmd {
			m.openForm(&formDialog{title: "Search " + name,
				note:   "As an administrator: no tenancy applies. Shows what a question would find.",
				fields: []formField{textF("question", "", ""), textF("results", "5", "how many chunks")},
				submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
					if v[0] == "" {
						return nil, fmt.Errorf("what question?")
					}
					k, _ := strconv.Atoi(v[1])
					req := adminapi.KBSearchRequest{Collection: c.ID, Query: v[0], K: k}
					return m.call("Search "+c.Name, http.MethodPost, "/api/v1/kb/search", req, func(raw json.RawMessage) string {
						var hits []adminapi.KBHitView
						_ = json.Unmarshal(raw, &hits)
						if len(hits) == 0 {
							return "Nothing found."
						}
						var b strings.Builder
						for i, h := range hits {
							text := h.Text
							if r := []rune(text); len(r) > 400 {
								text = string(r[:400]) + "…"
							}
							fmt.Fprintf(&b, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(text), "   "))
						}
						return strings.TrimRight(b.String(), "\n")
					}), nil
				}})
			return nil
		}},
		menuItem{"Delete…", func(m *consoleModel) tea.Cmd {
			m.confirmD = &confirmDialog{title: "Delete " + name,
				text: fmt.Sprintf("Its %d documents and %d chunks, and every grant, go with it. People's own files are not touched.",
					c.Documents, c.Chunks),
				yes: func(m *consoleModel) tea.Cmd {
					return m.call("Delete "+c.Name, http.MethodDelete, "/api/v1/kb/collections/"+kbRef(c), nil, nil)
				}}
			return nil
		}},
	)
	return &menuDialog{title: name, items: items}
}

// kbDocuments shows a collection's documents in a dialog.
func (m *consoleModel) kbDocuments(c adminapi.KBCollectionView) tea.Cmd {
	return m.call("Documents of "+c.Name, http.MethodGet, "/api/v1/kb/collections/"+kbRef(c)+"/documents", nil,
		func(raw json.RawMessage) string {
			var docs []adminapi.KBDocumentView
			_ = json.Unmarshal(raw, &docs)
			if len(docs) == 0 {
				return "None yet. On a person's machine: llmgw-kb sync " + c.Name + " FOLDER"
			}
			var b strings.Builder
			for i, d := range docs {
				if i == 40 {
					fmt.Fprintf(&b, "… and %d more (llmgw ctl kb docs %s)", len(docs)-40, c.ID)
					break
				}
				fmt.Fprintf(&b, "%-48s %5d chunks  %s  %s\n", d.Name, d.Chunks, d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
			}
			return strings.TrimRight(b.String(), "\n")
		})
}