llmgw/internal/cli/tui_mirror.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

// The mirror, as the control plane reaches it: its management API, with the
// address and certificate server.yaml gives under mirror.

type consoleMirrorMsg struct {
	res mirror.ModelsResponse
	err error
}

func (m consoleModel) mirrorClient() (*mirror.Client, error) {
	cfg, err := config.Load(m.file)
	if err != nil {
		return nil, err
	}
	if cfg.Mirror.API == "" {
		return nil, errors.New("no mirror: mirror.api in server.yaml is empty")
	}
	return mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, 10*time.Second)
}

func (m consoleModel) fetchMirror() tea.Cmd {
	return func() tea.Msg {
		c, err := m.mirrorClient()
		if err != nil {
			return consoleMirrorMsg{err: err}
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		res, err := c.List(ctx)
		return consoleMirrorMsg{res, err}
	}
}

// mirrorCall runs one request to the mirror in the background and reports it.
func (m *consoleModel) mirrorCall(title string, do func(context.Context, *mirror.Client) (string, error)) tea.Cmd {
	m.busy = title
	mm := *m
	return func() tea.Msg {
		c, err := mm.mirrorClient()
		if err != nil {
			return actionDoneMsg{title: title, err: err}
		}
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		text, err := do(ctx, c)
		return actionDoneMsg{title: title, text: text, err: err}
	}
}

// mirrorGroup is one version of one model on the mirror: a file, or the set
// of files a model of several is fetched as.
type mirrorGroup struct {
	model, version string
	files          []mirror.Status
}

func (g mirrorGroup) bytes() (done, total int64) {
	for _, f := range g.files {
		done += f.BytesDone
		total += max(f.BytesTotal, f.BytesDone)
	}
	return done, total
}

// source is what a model's definition names for it: the file, or the
// directory of a set.
func (g mirrorGroup) source(base string) string {
	if len(g.files) == 1 {
		return mirror.Entry{Model: g.model, Version: g.version, File: g.files[0].File}.URL(base)
	}
	return strings.TrimSuffix(mirror.Entry{Model: g.model, Version: g.version}.URL(base), "/") + "/"
}

func (m consoleModel) mirrorGroups() []mirrorGroup {
	var out []mirrorGroup
	at := map[string]int{}
	for _, st := range m.mirror.res.Models {
		k := st.Model + "/" + st.Version
		i, ok := at[k]
		if !ok {
			i = len(out)
			at[k] = i
			out = append(out, mirrorGroup{model: st.Model, version: st.Version})
		}
		out[i].files = append(out[i].files, st)
	}
	return out
}

// usedBy names the catalog's models that run from a version's files.
func (m consoleModel) usedBy(base string, g mirrorGroup) []string {
	dir := strings.TrimSuffix(mirror.Entry{Model: g.model, Version: g.version}.URL(base), "/") + "/"
	var names []string
	for _, md := range m.data.models {
		if strings.HasPrefix(md.Source, dir) {
			names = append(names, md.Name)
		}
	}
	return names
}

func (m consoleModel) mirrorMenu(g mirrorGroup) *menuDialog {
	what := g.model + "/" + g.version
	used := m.usedBy(m.mirror.res.Base, g)
	items := []menuItem{{"Its files, and where nodes fetch them", func(m *consoleModel) tea.Cmd {
		var b strings.Builder
		for _, f := range g.files {
			sum := ""
			if f.SHA256 != "" {
				sum = sDim.Render("  sha256 " + f.SHA256[:min(12, len(f.SHA256))])
			}
			fmt.Fprintf(&b, "%-50s %9s%s\n", truncate(f.File, 50), bytesShort(max(f.BytesTotal, f.BytesDone)), sum)
			if f.Source != "" {
				b.WriteString(sDim.Render("  from "+truncate(f.Source, 90)) + "\n")
			}
		}
		m.dialog, m.dialogTitle = "error", what
		m.dialogText = b.String() + "\nA model's source: " + g.source(m.mirror.res.Base) + "\n\n" + sDim.Render("any key")
		return nil
	}}}
	items = append(items, menuItem{"Remove from the mirror…", func(m *consoleModel) tea.Cmd {
		text := "Every file of this version is deleted from the mirror."
		if len(used) > 0 {
			text += "\n\n" + sWarn.Render("! ") + "The catalog uses it: " + strings.Join(used, ", ") +
				".\nNodes that do not hold it already cannot start them."
		}
		m.confirmD = &confirmDialog{title: "Remove " + what, text: text,
			yes: func(m *consoleModel) tea.Cmd {
				return m.mirrorCall("Remove "+what, func(ctx context.Context, c *mirror.Client) (string, error) {
					removed, err := c.Remove(ctx, g.model, g.version)
					return fmt.Sprintf("Removed %d file(s).", len(removed)), err
				})
			}}
		return nil
	}})
	return &menuDialog{title: what, items: items}
}

func (m *consoleModel) mirrorFetchForm() {
	m.openForm(&formDialog{title: "Download into the mirror",
		note: "The mirror downloads it once; nodes then fetch it from the mirror.",
		fields: []formField{
			textF("from", "", "an https:// URL, like a Hugging Face resolve link"),
			textF("model", "", "one directory level, like qwen2.5-7b"),
			textF("version", time.Now().Format("2006-01-02"), "one directory level; a date works well"),
			textF("file", "", "optional: taken from the URL"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			if !strings.HasPrefix(v[0], "http://") && !strings.HasPrefix(v[0], "https://") {
				return nil, fmt.Errorf("from is a URL")
			}
			if v[1] == "" || v[2] == "" {
				return nil, fmt.Errorf("a model and a version")
			}
			req := mirror.FetchRequest{Source: v[0], Model: v[1], Version: v[2], File: v[3]}
			return m.mirrorCall("Download "+v[1], func(ctx context.Context, c *mirror.Client) (string, error) {
				r, err := c.Fetch(ctx, req)
				if err != nil {
					return "", err
				}
				return "Downloading. Once it is here, a model's source is\n" + r.URL, nil
			}), nil
		}})
}

func (m consoleModel) mirrorPanel(pw int) (string, string, string) {
	keys := "→ or Enter open · ↑↓ choose · Q leave"
	if m.inPanel {
		keys = "↑↓ row · Enter actions · A download one · ← back · Q leave"
	}
	if m.mirror.err != nil {
		return "Mirror", "\n  " + sBad.Render(m.mirror.err.Error()), keys
	}
	head := "\n  " + sDim.Render("Nodes fetch from "+m.mirror.res.Base) + "\n\n"
	var rows [][]string
	for _, g := range m.mirrorGroups() {
		done, total := g.bytes()
		state := "present"
		for _, f := range g.files {
			switch f.State {
			case mirror.Fetching:
				state = sWarn.Render("fetching")
				if total > 0 {
					state = sWarn.Render(fmt.Sprintf("fetching %d%%", done*100/total))
				}
			case mirror.Failed:
				state = sBad.Render("failed")
			}
			if f.State != mirror.Present {
				break
			}
		}
		what := g.files[0].File
		if len(g.files) > 1 {
			what = fmt.Sprintf("%d files", len(g.files))
		}
		used := strings.Join(m.usedBy(m.mirror.res.Base, g), ", ")
		if used == "" {
			used = sDim.Render("unused")
		}
		rows = append(rows, []string{g.model, g.version, what, state, bytesShort(total), used})
	}
	if len(rows) == 0 {
		return "Mirror", head + "  " + sDim.Render("Empty. A downloads one."), keys
	}
	return "Mirror", head + table([]string{"MODEL", "VERSION", "FILE", "STATE", "SIZE", "USED BY"}, rows, m.row, pw, m.inPanel), keys
}

// addModelForm adds a model from a definition file on this machine: the
// same YAML a model takes in server.yaml.
func (m *consoleModel) addModelForm() {
	server := "--server=" + m.client.base
	m.openForm(&formDialog{title: "Add a model",
		note: "From a definition file on this machine: the YAML a model takes in\nserver.yaml's models. Serve a model, under Do, needs none.",
		fields: []formField{
			textF("file", "", "like /root/qwen.yaml"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			if _, err := os.Stat(v[0]); err != nil {
				return nil, err
			}
			return m.runCLI("Add a model", "ctl", server, "model", "add", "-f", v[0]), nil
		}})
}