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

// usedBy names the catalog's models whose source is this file.
func (m consoleModel) usedBy(base string, s mirror.Status) []string {
	u := mirror.Entry{Model: s.Model, Version: s.Version, File: s.File}.URL(base)
	dir := strings.TrimSuffix(mirror.Entry{Model: s.Model, Version: s.Version}.URL(base), "/") + "/"
	var names []string
	for _, md := range m.data.models {
		if md.Source == u || strings.HasPrefix(md.Source, dir) {
			names = append(names, md.Name)
		}
	}
	return names
}

func (m consoleModel) mirrorMenu(s mirror.Status) *menuDialog {
	what := s.Model + "/" + s.Version
	used := m.usedBy(m.mirror.res.Base, s)
	items := []menuItem{{"Where nodes fetch it", func(m *consoleModel) tea.Cmd {
		m.dialog, m.dialogTitle = "error", what
		m.dialogText = mirror.Entry{Model: s.Model, Version: s.Version, File: s.File}.URL(m.mirror.res.Base) +
			"\n\n" + sDim.Render("the source to give a model definition · 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, s.Model, s.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 _, s := range m.mirror.res.Models {
		state := string(s.State)
		switch s.State {
		case mirror.Fetching:
			if p, ok := s.Percent(); ok {
				state = sWarn.Render(fmt.Sprintf("fetching %d%%", p))
			} else {
				state = sWarn.Render("fetching " + bytesShort(s.BytesDone))
			}
		case mirror.Failed:
			state = sBad.Render("failed")
		}
		size := bytesShort(s.BytesDone)
		used := strings.Join(m.usedBy(m.mirror.res.Base, s), ", ")
		if used == "" {
			used = sDim.Render("unused")
		}
		rows = append(rows, []string{s.Model, s.Version, s.File, state, size, 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
		}})
}