llmgw/internal/router/chat_load.go

package router

import (
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// modelStatus is how far a model is from answering, for the chat page's
// progress bar. A person who asks for a model that is not loaded should see
// it being fetched and started, not a spinner that ends in an error.
type modelStatus struct {
	// State is ready, downloading, starting, waiting, failed or stopped.
	State string `json:"state"`
	// Percent, Done and Total describe a download; Percent is absent when the
	// size is unknown, rather than a number that stops moving.
	Percent *int  `json:"percent,omitempty"`
	Done    int64 `json:"done,omitempty"`
	Total   int64 `json:"total,omitempty"`
	// Seconds is how long the current stage has run.
	Seconds int64  `json:"seconds,omitempty"`
	Message string `json:"message,omitempty"`
}

// loadStatus says where a model stands, from what the control plane reports.
func (r *Router) loadStatus(model string) modelStatus {
	if len(r.reg.Snapshots(model)) > 0 {
		return modelStatus{State: "ready"}
	}
	if r.clusterView == nil {
		return modelStatus{State: "stopped", Message: "nothing is serving this model"}
	}
	var best *modelStatus
	for _, rep := range r.clusterView.Replicas(model) {
		st := modelStatus{Seconds: rep.AgeSec}
		switch rep.State {
		case "ready":
			// Up, and about to be published to the router.
			return modelStatus{State: "starting", Message: "almost ready"}
		case "failed":
			st.State, st.Message = "failed", rep.Err
		case "placing", "starting":
			st.State = "starting"
			for _, sh := range rep.Shards {
				if sh.Fetch == nil {
					continue
				}
				st.State = "downloading"
				st.Done += sh.Fetch.BytesDone
				st.Total += sh.Fetch.BytesTotal
			}
			if st.State == "downloading" && st.Total > 0 {
				p := int(st.Done * 100 / st.Total)
				st.Percent = &p
			}
		default:
			continue
		}
		// A replica making progress says more than one that failed.
		if best == nil || best.State == "failed" {
			s := st
			best = &s
		}
	}
	if best != nil {
		return *best
	}
	if _, m := r.resolve(model); m.Scaling.OnDemand() {
		return modelStatus{State: "waiting", Message: "waiting for a node with room for it"}
	}
	return modelStatus{State: "stopped", Message: "this model is not running"}
}

// chatModel resolves the model a chat request names, checking the person may
// use it in the context they chose. Empty when they may not, which answers the
// same as a model that does not exist.
func (r *Router) chatModel(req *http.Request, user store.User, asked, context string) string {
	name, _ := r.resolve(asked)
	if _, known := r.modelMap()[name]; !known && len(r.reg.Snapshots(name)) == 0 {
		return ""
	}
	p, err := r.principalFor(req, user, context)
	if err != nil || !r.permits(p, name) {
		return ""
	}
	return name
}

func (r *Router) handleChatStatus(w http.ResponseWriter, req *http.Request, user store.User) {
	name := r.chatModel(req, user, req.URL.Query().Get("model"), req.URL.Query().Get("context"))
	if name == "" {
		http.NotFound(w, req)
		return
	}
	writeJSON(w, http.StatusOK, r.loadStatus(name))
}

// handleChatWake asks for an on-demand model and returns at once: the page
// shows the load happening rather than holding a request open until it ends.
func (r *Router) handleChatWake(w http.ResponseWriter, req *http.Request, user store.User) {
	var in struct {
		Model   string `json:"model"`
		Context string `json:"context"`
	}
	body, _ := io.ReadAll(io.LimitReader(req.Body, 4096))
	if err := json.Unmarshal(body, &in); err != nil || strings.TrimSpace(in.Model) == "" {
		http.Error(w, "which model?", http.StatusBadRequest)
		return
	}
	name := r.chatModel(req, user, in.Model, in.Context)
	if name == "" {
		http.NotFound(w, req)
		return
	}
	if r.activator != nil && len(r.reg.Snapshots(name)) == 0 {
		r.activator.Want(name)
	}
	writeJSON(w, http.StatusOK, r.loadStatus(name))
}