llmgw/internal/cli/tui_hf.go

package cli

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	tea "github.com/charmbracelet/bubbletea"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
)

// Serving from Hugging Face: the registry's GGUF models, trending or
// searched, then a file of the one chosen, then `llmgw ctl model serve` with
// it — which fetches it into the mirror, checks it, adds it and starts it.

var hfSorts = []struct {
	label string
	sort  hf.Sort
}{{"trending", hf.SortTrending}, {"most downloaded", hf.SortDownloads}, {"most liked", hf.SortLikes}, {"newest", hf.SortCreated}}

type hfState struct {
	query string
	sort  int
	// safetensors lists original weights, for vLLM on GPUs, rather than
	// GGUF files for llama.cpp.
	safetensors bool
	models      []hf.Model
	err         error
}

type consoleHFMsg struct {
	query  string
	sort   int
	st     bool
	models []hf.Model
	err    error
}

type consoleRepoMsg struct {
	repo hf.Repo
	err  error
}

// hfClients keeps one client per registry, for the client's few minutes of
// remembered listings.
var hfClients sync.Map

func (m consoleModel) registry() *hf.Client {
	endpoint := ""
	if cfg, err := config.Load(m.file); err == nil {
		endpoint = cfg.Mirror.Registry
	}
	c, _ := hfClients.LoadOrStore(endpoint, hf.New(endpoint, 20*time.Second))
	return c.(*hf.Client)
}

func (m consoleModel) fetchHF() tea.Cmd {
	c, q, sort, st := m.registry(), m.hf.query, m.hf.sort, m.hf.safetensors
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
		defer cancel()
		hq := hf.Query{Search: q, Sort: hfSorts[sort].sort, GGUF: !st, Limit: 50}
		if st {
			hq.Filters = []string{"safetensors"}
		}
		// Browsing: models that chat. A search finds the rest too, an
		// embedding model included.
		if q == "" {
			hq.Task = "text-generation"
		}
		got, err := c.Explore(ctx, hq)
		return consoleHFMsg{query: q, sort: sort, st: st, models: got, err: err}
	}
}

func (m *consoleModel) fetchRepo(id string) tea.Cmd {
	m.busy = "Reading " + id
	c := m.registry()
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
		defer cancel()
		r, err := c.Repo(ctx, id)
		return consoleRepoMsg{r, err}
	}
}

// repoMenu offers a repository's GGUF files, smallest first as listed.
func (m consoleModel) repoMenu(r hf.Repo) *menuDialog {
	var items []menuItem
	if m.hf.safetensors {
		title := r.ID
		set, err := modelserve.SafetensorsSet(r)
		if err != nil {
			return &menuDialog{title: title, items: []menuItem{{err.Error(), func(*consoleModel) tea.Cmd { return nil }}}}
		}
		return &menuDialog{title: title, items: []menuItem{{
			fmt.Sprintf("Original weights for vLLM: %s in %d files", bytesShort(set.Bytes), len(set.Files)),
			func(m *consoleModel) tea.Cmd { m.serveVLLMForm(r, set); return nil },
		}}}
	}
	for _, set := range hf.Group(r.Files) {
		if !strings.HasSuffix(strings.ToLower(set.Files[0].Path), ".gguf") {
			continue
		}
		low := strings.ToLower(set.Name)
		// Projectors and imatrix data are not a model to serve on their own.
		if strings.Contains(low, "mmproj") || strings.Contains(low, "imatrix") {
			continue
		}
		label := fmt.Sprintf("%-44s %8s", truncate(set.Name, 44), bytesShort(set.Bytes))
		if set.Shards > 1 {
			label += fmt.Sprintf("  %d parts", set.Shards)
		}
		set := set
		items = append(items, menuItem{label, func(m *consoleModel) tea.Cmd {
			m.serveForm(r, set)
			return nil
		}})
	}
	title := r.ID
	if len(items) == 0 {
		items = []menuItem{{"No GGUF file here to serve", func(*consoleModel) tea.Cmd { return nil }}}
	}
	if r.Gated.Is() {
		title += " · gated"
	}
	return &menuDialog{title: title, items: items}
}

func (m *consoleModel) serveForm(r hf.Repo, set hf.FileSet) {
	server := "--server=" + m.client.base
	note := fmt.Sprintf("%s, %s", set.Name, bytesShort(set.Bytes))
	if r.Architecture != "" {
		note += ", " + r.Architecture
	}
	if r.License != "" {
		note += ", licence " + r.License
	}
	note += "\nThe mirror fetches it and checks it, it joins the catalog, and a\nreplica is started and asked to say hello. The output is shown as it goes."
	fields := []formField{
		textF("name", hf.ModelName(r.ID), "what clients ask for"),
		textF("replicas", "1", "0 runs none until it is first asked for"),
		textF("context length", "4096", "what one request may use"),
		textF("description", "", "optional: shown to clients in /v1/models"),
	}
	// A vision model: its projector, loaded beside it, is what reads images.
	if proj := modelserve.Projectors(r); len(proj) > 0 {
		fields = append(fields, choiceF("images", append(proj, "none: text only"), proj[0]))
		fields[len(fields)-1].hint = "the vision projector, fetched and loaded beside the model"
	}
	m.openForm(&formDialog{title: "Serve " + r.ID, note: note,
		fields: fields,
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			n, err := strconv.Atoi(v[1])
			if err != nil || n < 0 {
				return nil, fmt.Errorf("replicas is a whole number")
			}
			if _, err := strconv.Atoi(v[2]); err != nil {
				return nil, fmt.Errorf("the context length is a whole number")
			}
			args := []string{"ctl", server, "model", "serve", r.ID, "--file=" + set.Name,
				"--name=" + v[0], "--replicas=" + v[1], "--context-length=" + v[2]}
			if n == 0 {
				args = append(args, "--on-demand")
			}
			if v[3] != "" {
				args = append(args, "--description="+v[3])
			}
			if len(v) > 4 && !strings.HasPrefix(v[4], "none") {
				args = append(args, "--projector="+v[4])
			}
			return m.runAfter(args...), nil
		}})
}

func (m *consoleModel) serveVLLMForm(r hf.Repo, set hf.FileSet) {
	server := "--server=" + m.client.base
	var classes []string
	for _, g := range m.data.sum.GPUs {
		classes = append(classes, g.Class)
	}
	if len(classes) == 0 {
		m.dialog, m.dialogTitle = "error", "No GPUs"
		m.dialogText = "Original weights run on vLLM, on GPUs, and this fleet reports none.\n" +
			"F goes back to GGUF models, which run on the CPU.\n\n" + sDim.Render("any key")
		return
	}
	m.openForm(&formDialog{title: "Serve " + r.ID,
		note: fmt.Sprintf("%s in %d files: the weights, config.json and the tokenizer.\n"+
			"Fetched into the mirror and checked, added, and run by vLLM on the GPUs chosen.", bytesShort(set.Bytes), len(set.Files)),
		fields: []formField{
			textF("name", hf.ModelName(r.ID), "what clients ask for"),
			textF("replicas", "1", "0 runs none until it is first asked for"),
			textF("context length", "8192", "what one request may use; vLLM's --max-model-len"),
			textF("description", "", "optional: shown to clients in /v1/models"),
			choiceF("GPU class", classes, classes[0]),
			textF("GPUs per replica", "1", "tensor parallel: the cards one replica spans"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			for i, what := range map[int]string{1: "replicas", 2: "the context length", 5: "GPUs per replica"} {
				if n, err := strconv.Atoi(v[i]); err != nil || n < 0 || (i == 5 && n < 1) {
					return nil, fmt.Errorf("%s is a whole number", what)
				}
			}
			args := []string{"ctl", server, "model", "serve", r.ID, "--format=safetensors", "--gpu-class=" + v[4],
				"--tp=" + v[5], "--name=" + v[0], "--replicas=" + v[1], "--context-length=" + v[2]}
			if v[1] == "0" {
				args = append(args, "--on-demand")
			}
			if v[3] != "" {
				args = append(args, "--description="+v[3])
			}
			return m.runAfter(args...), nil
		}})
}

func (m *consoleModel) hfSearchDialog() {
	m.openForm(&formDialog{title: "Search Hugging Face",
		note:   "GGUF models, which is what the nodes run.",
		fields: []formField{textF("search", m.hf.query, "like qwen2.5 coder; empty lists them all")},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			m.hf.query, m.hf.models, m.hf.err, m.row = v[0], nil, nil, 0
			return m.fetchHF(), nil
		}})
}

func (m consoleModel) hfPanel(pw int) (string, string, string) {
	keys := "→ or Enter open · ↑↓ choose · Q leave"
	if m.inPanel {
		keys = "↑↓ row · Enter choose · / search · S order · F GGUF or original weights · ← back · Q leave"
	}
	what := hfSorts[m.hf.sort].label
	if m.hf.query != "" {
		what = "“" + m.hf.query + "”, " + what
	}
	if m.hf.query == "" {
		what = "text generation, " + what
	}
	kind := "GGUF models, for llama.cpp"
	if m.hf.safetensors {
		kind = "Original weights, for vLLM on GPUs"
	}
	head := "\n  " + kind + " on Hugging Face: " + what + "\n\n"
	switch {
	case m.hf.err != nil:
		return "Serve from Hugging Face", head + "  " + sBad.Render(m.hf.err.Error()) +
			"\n\n  " + sDim.Render("mirror.registry in server.yaml names another registry, for a site without the internet."), keys
	case m.hf.models == nil:
		return "Serve from Hugging Face", head + "  " + m.spin.View() + " asking the registry…", keys
	}
	var rows [][]string
	for _, x := range m.hf.models {
		params := ""
		if p := x.Params(); p > 0 {
			params = fmt.Sprintf("%.1fB", float64(p)/1e9)
		}
		gated := ""
		if x.Gated.Is() {
			gated = sWarn.Render("gated")
		}
		rows = append(rows, []string{x.ID, params, x.Architecture(), compactNum(x.Downloads), compactNum(x.Likes),
			ageShort(x.LastModified), gated})
	}
	if len(rows) == 0 {
		return "Serve from Hugging Face", head + "  " + sDim.Render("Nothing found. / searches again."), keys
	}
	return "Serve from Hugging Face", head + table([]string{"REPOSITORY", "SIZE", "ARCH", "DOWNLOADS", "LIKES", "UPDATED", ""},
		rows, m.row, pw, m.inPanel), keys
}

// ageShort is how long ago, in the largest unit that says it: 5m, 3h, 12d, 7mo, 2y.
func ageShort(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	d := time.Since(t)
	switch {
	case d < time.Hour:
		return fmt.Sprintf("%dm ago", int(d.Minutes()))
	case d < 48*time.Hour:
		return fmt.Sprintf("%dh ago", int(d.Hours()))
	case d < 60*24*time.Hour:
		return fmt.Sprintf("%dd ago", int(d.Hours()/24))
	case d < 2*365*24*time.Hour:
		return fmt.Sprintf("%dmo ago", int(d.Hours()/24/30))
	}
	return fmt.Sprintf("%dy ago", int(d.Hours()/24/365))
}