llmgw/internal/cli/tui_hf.go

package cli

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	tea "github.com/charmbracelet/bubbletea"

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

// 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 {
	plan modelserve.Plan
	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 := loadEffectiveConfig(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}
	}
}

// fetchRepo asks the control plane for a repository's plan: its files, and
// how each could run on this cluster.
func (m *consoleModel) fetchRepo(id string) tea.Cmd {
	m.busy = "Planning " + id
	c := m.client
	format := "gguf"
	if m.hf.safetensors {
		format = "safetensors"
	}
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		var p modelserve.Plan
		err := c.get(ctx, "/api/v1/plan?"+url.Values{"repo": {id}, "format": {format}}.Encode(), &p)
		return consoleRepoMsg{p, err}
	}
}

// optionLabel is one way to run a file, with what it needs and whether the
// hardware is free.
func optionLabel(o modelplan.Option) string {
	now := sDim.Render("busy now")
	if o.FreeNow {
		now = sOK.Render("free")
	}
	return fmt.Sprintf("%s · %s · %s", o.Label(), o.MemoryText(), now)
}

// repoMenu offers a repository's files, each with the best way it runs here;
// the plan's choice first.
func (m consoleModel) repoMenu(p modelserve.Plan) *menuDialog {
	title := p.Repo
	if p.Gated {
		title += " · gated"
	}
	var items []menuItem
	if !p.Servable {
		return &menuDialog{title: title, items: []menuItem{{"The nodes' engine cannot run it: " + p.Reason,
			func(*consoleModel) tea.Cmd { return nil }}}}
	}
	rec, hasRec := p.Recommended()
	arts := append([]modelserve.PlannedArtifact{}, p.Artifacts...)
	// The recommended file first, then the rest as the repository lists them.
	for i, a := range arts {
		if hasRec && a.Name == rec.Artifact && i > 0 {
			arts = append([]modelserve.PlannedArtifact{a}, append(arts[:i:i], arts[i+1:]...)...)
			break
		}
	}
	for _, a := range arts {
		name := a.Name
		if a.Format == "safetensors" {
			name = "original weights"
		}
		// The quantisation first: long names end in it, and it is what
		// tells one file of a repository from the next.
		q := a.Quant
		if a.Shards > 1 && a.Format == "gguf" {
			q += fmt.Sprintf(" ×%d", a.Shards)
		}
		label := fmt.Sprintf("%-10s %-30s %8s  ", q, truncate(name, 30), bytesShort(a.Bytes))
		switch {
		case a.Best == nil:
			label += sBad.Render("fits nowhere")
		default:
			label += optionLabel(*a.Best)
		}
		if hasRec && a.Name == rec.Artifact {
			label = sOK.Render("★ ") + label
		} else {
			label = "  " + label
		}
		a := a
		items = append(items, menuItem{label, func(m *consoleModel) tea.Cmd {
			if a.Best == nil {
				m.dialog, m.dialogTitle = "error", a.Name
				m.dialogText = fmt.Sprintf("No card holds %s with a %d-token context, nor does any node's memory.\n"+
					"A smaller quantisation, or more cards free.\n\n%s", a.Name, p.Context, sDim.Render("any key"))
				return nil
			}
			m.serveForm(p, a)
			return nil
		}})
	}
	return &menuDialog{title: title, items: items}
}

func (m *consoleModel) serveForm(p modelserve.Plan, a modelserve.PlannedArtifact) {
	var opts []modelplan.Option
	for _, o := range p.Options {
		if o.Artifact == a.Name {
			opts = append(opts, o)
		}
	}
	var where []string
	for _, o := range opts {
		where = append(where, optionLabel(o))
	}
	note := fmt.Sprintf("%s, %s", a.Name, bytesShort(a.Bytes))
	if a.Format == "safetensors" {
		note = fmt.Sprintf("Original weights, %s in %d files with their config and tokenizer, for vLLM", bytesShort(a.Bytes), a.Shards)
	}
	if p.Arch != "" {
		note += ", " + p.Arch
	}
	if p.License != "" {
		note += ", licence " + p.License
	}
	if p.Note != "" {
		note += "\n" + sWarn.Render("! ") + p.Note
	}
	note += fmt.Sprintf("\nMemory for a %d-token context, each way it runs:", p.Context)
	for _, o := range opts {
		note += "\n  " + o.Label() + ": " + o.MemoryText()
	}
	if p.KVBytesPerToken > 0 {
		note += fmt.Sprintf("\nThe context costs %.2f GB per 1000 tokens, the whole model; a longer one is checked\n"+
			"before it is served.", modelplan.GB(p.KVBytesPerToken*1000))
	}
	note += "\nThe mirror fetches it and checks it, it joins the catalog, and a replica is started\nand asked to say hello."
	ctxDefault := "4096"
	if a.Format == "safetensors" {
		ctxDefault = strconv.Itoa(p.Context)
	}
	fields := []formField{
		textF("name", hf.ModelName(p.Repo), "what clients ask for"),
		textF("replicas", "1", "0 runs none until it is first asked for"),
		textF("context length", ctxDefault, "what one request may use"),
		textF("description", "", "optional: shown to clients in /v1/models"),
		choiceF("runs on", where, where[0]),
		choiceF("when that is busy", []string{"the next best way", "wait for it"}, "the next best way"),
	}
	fields[4].hint = "every way this file fits, best first"
	proj := -1
	if a.Format == "gguf" && len(p.Projectors) > 0 {
		fields = append(fields, choiceF("images", append(append([]string{}, p.Projectors...), "none: text only"), p.Projectors[0]))
		fields[len(fields)-1].hint = "the vision projector, fetched and loaded beside the model"
		proj = len(fields) - 1
	}
	m.openForm(&formDialog{title: "Serve " + p.Repo, 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 c, err := strconv.Atoi(v[2]); err != nil || c <= 0 {
				return nil, fmt.Errorf("the context length is a whole number")
			}
			ctxLen, _ := strconv.Atoi(v[2])
			req := modelserve.Request{Repo: p.Repo, Name: v[0], Replicas: n, ContextLength: ctxLen,
				Description: v[3], Format: a.Format}
			if a.Format == "gguf" {
				req.File = a.Name
			}
			o := opts[slices.Index(where, v[4])]
			// Asked for a context other than planned: does it still fit?
			if sized := o.WithContext(p.KVBytesPerToken, ctxLen); !sized.Fits() {
				return nil, fmt.Errorf("with a %d-token context it needs %s: a shorter context, or another way", ctxLen, sized.MemoryText())
			}
			if o.CPU {
				req.CPU = true
			} else {
				req.GPUClass, req.GPUsPerNode, req.Nodes = o.Class, o.GPUsPerNode, max(o.Nodes, 1)
			}
			req.Fallbacks = strings.HasPrefix(v[5], "the next") && len(opts) > 1
			if proj >= 0 && !strings.HasPrefix(v[proj], "none") {
				req.Projector = v[proj]
			}
			// Started on the control plane, which runs it whatever happens
			// to this screen: the console is free again at once, and Jobs
			// shows how far it has got.
			return m.call("Serve "+p.Repo, http.MethodPost, "/api/v1/models/serve", req, startedSays), 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.Minute:
		return "just now"
	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))
}