llmgw/internal/control/adminapi/presets.go

package adminapi

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// PresetsView is the models that come with llmgw, and whether the leader adds
// them by itself.
type PresetsView struct {
	// Auto is whether a preset is added when a node fits it and the catalog
	// has no model that makes what it makes (presets.auto).
	Auto bool `json:"auto"`
	// ComfyUI is engines.comfyui as it stands: the version the nodes
	// install, "" for none — in which case adding a preset sets it, and the
	// control planes start again on that.
	ComfyUI string       `json:"comfyui"`
	Presets []PresetView `json:"presets"`
}

// PresetView is one preset as it stands in this cluster.
type PresetView struct {
	Name        string          `json:"name"`
	Makes       string          `json:"makes"`
	Description string          `json:"description"`
	Card        string          `json:"card"`
	Licence     catalog.License `json:"licence"`
	// Fits is whether some node has a card that meets the preset's need;
	// FitsClasses names those classes.
	Fits        bool     `json:"fits"`
	FitsClasses []string `json:"fits_classes,omitempty"`
	InCatalog   bool     `json:"in_catalog"`
	// AddedBy is "auto" when the leader added it, who did otherwise, ""
	// when it was not added as a preset; AddedAt is when.
	AddedBy string    `json:"added_by,omitempty"`
	AddedAt time.Time `json:"added_at,omitzero"`
	// Declined is whether somebody removed it after it was added, so it is
	// not added again by itself; adding it again clears that.
	Declined bool `json:"declined"`
	// Blocked is why the leader does not add it by itself — presets.auto
	// off, no ComfyUI, no card that fits, declined, another model already
	// making the same — empty when it would or has.
	Blocked string `json:"blocked,omitempty"`
}

// PresetsAutoRequest turns the automatic adding of presets on or off.
type PresetsAutoRequest struct {
	On bool `json:"on"`
}

// errPresetsUnavailable is a control plane without the store presets are
// kept in.
var errPresetsUnavailable = errors.New("this control plane has no store, so presets cannot be added here")

// PresetStore keeps what became of each preset, and the audit log the
// leader writes its own additions to.
type PresetStore interface {
	PresetStates(ctx context.Context) (map[string]store.PresetState, error)
	PresetAdded(ctx context.Context, name, by string, at time.Time) error
	PresetDeclined(ctx context.Context, name, by string, at time.Time) error
	WriteAudit(ctx context.Context, e store.AuditEvent) error
}

// PresetConfig is what the presets read of the settings in force.
type PresetConfig struct {
	// Auto is presets.auto.
	Auto bool
	// ComfyUI is engines.comfyui: the version the nodes install, empty for
	// none.
	ComfyUI string
}

// SetPresets enables the presets: where what became of them is kept, and
// the settings in force.
func (a *API) SetPresets(st PresetStore, cfg func() PresetConfig) {
	a.presets, a.presetCfg = st, cfg
}

// Presets is the presets as they stand in this cluster.
func (a *API) Presets(ctx context.Context) (PresetsView, error) {
	if a.presets == nil || a.presetCfg == nil {
		return PresetsView{}, errPresetsUnavailable
	}
	states, err := a.presets.PresetStates(ctx)
	if err != nil {
		return PresetsView{}, err
	}
	cfg := a.presetCfg()
	cat := a.ctl.Catalog()
	kinds := a.gpuKinds()
	out := PresetsView{Auto: cfg.Auto, ComfyUI: cfg.ComfyUI}
	for _, p := range catalog.Presets() {
		v := PresetView{Name: p.Model.Name, Makes: p.Makes, Description: p.Model.Description,
			Card: presetCard(p.Model), FitsClasses: presetFits(p.Model, kinds)}
		if p.Model.License != nil {
			v.Licence = *p.Model.License
		}
		v.Fits = len(v.FitsClasses) > 0
		_, v.InCatalog = cat.Resolve(p.Model.Name)
		if st, ok := states[p.Model.Name]; ok {
			v.AddedBy, v.AddedAt, v.Declined = st.AddedBy, st.AddedAt, st.Declined
		}
		v.Blocked = presetBlocked(p, v, cfg, cat)
		out.Presets = append(out.Presets, v)
	}
	return out, nil
}

// AddPreset adds the preset of that name, and clears its decline. A cluster
// that installs no ComfyUI yet is set to (engines.comfyui), as `llmgw ctl
// model add` does; that change starts the control planes again.
func (a *API) AddPreset(ctx context.Context, name, by string) error {
	if a.presets == nil || a.presetCfg == nil || a.models == nil {
		return errPresetsUnavailable
	}
	p, ok := catalog.PresetNamed(name)
	if !ok {
		return fmt.Errorf("%w: %q (there are %s)", errNoSuchPreset, name, strings.Join(catalog.PresetNames(), ", "))
	}
	if _, in := a.ctl.Catalog().Resolve(name); in {
		return fmt.Errorf("%s is already in the catalog", name)
	}
	if a.presetCfg().ComfyUI == "" {
		if err := a.setComfyUI(ctx, by); err != nil {
			return err
		}
	}
	if err := a.models.Add(ctx, p.Model, by); err != nil {
		return err
	}
	return a.presets.PresetAdded(ctx, name, by, time.Now())
}

// SetAutoPresets turns presets.auto on or off.
func (a *API) SetAutoPresets(ctx context.Context, on bool, by string) error {
	if a.presets == nil {
		return errPresetsUnavailable
	}
	note := "presets added by themselves: " + map[bool]string{true: "on", false: "off"}[on]
	return a.editSettings(ctx, []string{"presets", "auto"}, strconv.FormatBool(on), by, note)
}

// PresetsPass adds, as "auto", each preset a node fits whose kind of model
// the catalog has none of: run by the leader, now and then. It adds nothing
// with presets.auto off, where the cluster installs no ComfyUI, or a preset
// somebody declined. It says what it added.
func (a *API) PresetsPass(ctx context.Context) ([]string, error) {
	v, err := a.Presets(ctx)
	if err != nil || !v.Auto {
		return nil, err
	}
	var added []string
	for _, pv := range v.Presets {
		if pv.Blocked != "" || pv.InCatalog {
			continue
		}
		p, _ := catalog.PresetNamed(pv.Name)
		if err := a.models.Add(ctx, p.Model, store.PresetAdder); err != nil {
			return added, fmt.Errorf("adding %s: %w", pv.Name, err)
		}
		now := time.Now()
		if err := a.presets.PresetAdded(ctx, pv.Name, store.PresetAdder, now); err != nil {
			return added, err
		}
		_ = a.presets.WriteAudit(ctx, store.AuditEvent{At: now, Actor: store.PresetAdder, Source: "presets",
			Action: "add preset", Target: pv.Name, Status: http.StatusOK,
			Detail: fmt.Sprintf("%s, on demand: a node has %s (%s) and the catalog made no %s",
				pv.Name, strings.Join(pv.FitsClasses, ", "), pv.Card, pv.Makes)})
		added = append(added, pv.Name)
	}
	return added, nil
}

// presetDeclined records a preset's name removed from the catalog, so the
// leader does not put it back.
func (a *API) presetDeclined(ctx context.Context, name, by string) {
	if a.presets == nil {
		return
	}
	if _, ok := catalog.PresetNamed(name); ok {
		_ = a.presets.PresetDeclined(ctx, name, by, time.Now())
	}
}

// presetBlocked is why the leader would not add a preset by itself, empty
// when it would (or has).
func presetBlocked(p catalog.Preset, v PresetView, cfg PresetConfig, cat *catalog.Catalog) string {
	switch {
	case v.InCatalog:
		return ""
	case v.Declined:
		return "removed from the catalog after it was added; add it again to undo"
	case !cfg.Auto:
		return "presets.auto is off"
	case cfg.ComfyUI == "":
		return "the cluster installs no ComfyUI (engines.comfyui); adding it by hand sets that"
	case !v.Fits:
		return "no node has a card it needs (" + v.Card + ")"
	}
	for _, m := range cat.Models() {
		if p.Makes == catalog.OutputPictures && m.MakesImages() || p.Makes == catalog.OutputMeshes && m.MakesMeshes() {
			return "the catalog has " + m.Name + ", which makes " + p.Makes + " already"
		}
	}
	return ""
}

// presetCard is what a preset needs of a card, as people read it.
func presetCard(m catalog.Model) string {
	var needs []string
	for _, pr := range m.Profiles {
		switch {
		case pr.GPU != nil:
			needs = append(needs, pr.GPU.String())
		case pr.GPUClass != "":
			needs = append(needs, pr.GPUClass)
		}
	}
	return strings.Join(needs, " or ")
}

// presetFits are the classes of card in the cluster one of the preset's
// profiles can run on.
func presetFits(m catalog.Model, kinds []catalog.GPUKind) []string {
	var out []string
	for _, k := range kinds {
		for _, pr := range m.Profiles {
			if pr.GPU != nil && pr.GPU.Fits(k) || pr.GPU == nil && pr.GPUClass == k.Class {
				out = append(out, k.Class)
				break
			}
		}
	}
	return out
}

// gpuKinds are the kinds of card the cluster's nodes have.
func (a *API) gpuKinds() []catalog.GPUKind {
	var gpus []discovery.GPU
	for _, n := range a.ctl.Nodes() {
		gpus = append(gpus, n.GPUs...)
	}
	return catalog.GPUKinds(gpus)
}

// setComfyUI has the nodes install the ComfyUI this release knows.
func (a *API) setComfyUI(ctx context.Context, by string) error {
	if a.settings != nil {
		cur, err := a.settings.Store.ClusterSettings(ctx)
		if err != nil {
			return err
		}
		if doc, err := settings.FromBytes("settings", []byte(cur.Body)); err == nil {
			if n := doc.Node([]string{"engines", catalog.EngineComfyUI}); n != nil && n.Value == "" {
				return errors.New("engines.comfyui is set to none, so the nodes install no ComfyUI: " +
					"set a version (llmgw ctl config set engines.comfyui " + provision.ComfyUIVersion() + ")")
			}
		}
	}
	return a.editSettings(ctx, []string{"engines", catalog.EngineComfyUI}, provision.ComfyUIVersion(), by,
		"ComfyUI for the presets")
}

// editSettings sets one key of the cluster's settings, as a new version.
func (a *API) editSettings(ctx context.Context, path []string, value, by, note string) error {
	if a.settings == nil {
		return errors.New("this control plane keeps no cluster settings")
	}
	cur, err := a.settings.Store.ClusterSettings(ctx)
	if err != nil {
		return err
	}
	doc, err := settings.FromBytes("settings", []byte(cur.Body))
	if err != nil {
		return err
	}
	if n := doc.Node(path); n != nil && n.Value == value {
		return nil
	}
	if err := doc.SetYAML(path, value); err != nil {
		return err
	}
	b, err := doc.Bytes()
	if err != nil {
		return err
	}
	if err := a.settings.Check(b); err != nil {
		return err
	}
	_, err = a.settings.Store.PutClusterSettings(ctx, cur.Version, string(b), by, note)
	return err
}

func (a *API) registerPresets(mux Mux) {
	mux.HandleFunc("GET /api/v1/presets", func(w http.ResponseWriter, r *http.Request) {
		v, err := a.Presets(r.Context())
		if err != nil {
			writeError(w, presetStatus(err), err.Error())
			return
		}
		writeJSON(w, http.StatusOK, v)
	})
	mux.HandleFunc("POST /api/v1/presets/{name}", func(w http.ResponseWriter, r *http.Request) {
		if err := a.AddPreset(r.Context(), r.PathValue("name"), actorOf(r)); err != nil {
			writeError(w, presetStatus(err), err.Error())
			return
		}
		writeJSON(w, http.StatusOK, ModelAdded{Added: r.PathValue("name")})
	})
	mux.HandleFunc("PUT /api/v1/presets/auto", func(w http.ResponseWriter, r *http.Request) {
		var req PresetsAutoRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		if err := a.SetAutoPresets(r.Context(), req.On, actorOf(r)); err != nil {
			writeError(w, presetStatus(err), err.Error())
			return
		}
		v, err := a.Presets(r.Context())
		if err != nil {
			writeError(w, presetStatus(err), err.Error())
			return
		}
		writeJSON(w, http.StatusOK, v)
	})
}

// errNoSuchPreset is a name no preset has.
var errNoSuchPreset = errors.New("no preset of that name")

func presetStatus(err error) int {
	switch {
	case errors.Is(err, errPresetsUnavailable):
		return http.StatusNotImplemented
	case errors.Is(err, errNoSuchPreset):
		return http.StatusNotFound
	default:
		return http.StatusBadRequest
	}
}