llmgw/internal/control/adminapi/judges.go

package adminapi

import (
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
)

// JudgeView is one judge model the guardrail policies ask, and whether it
// can answer them.
type JudgeView struct {
	Model string `json:"model"`
	// Policies are those that ask it; Closed those that refuse what it
	// cannot read (fail: closed).
	Policies []string `json:"policies"`
	Closed   []string `json:"closed,omitempty"`
	// Then are the judges its policies ask after it, when it gives no
	// verdict: a CPU judge behind a GPU one.
	Then []string `json:"then,omitempty"`
	// InCatalog is whether a model of that name is in the catalog, and
	// Judge whether it is there as a judge (role: guard).
	InCatalog bool `json:"in_catalog"`
	Judge     bool `json:"judge"`
	Ready     int  `json:"ready_replicas"`
	Wanted    int  `json:"wanted_replicas"`
	// BudgetMs is how long a request waits for it: the shortest of its
	// policies' latency budgets.
	BudgetMs int64 `json:"budget_ms,omitempty"`
	// LastError is its last failure to give a verdict, if any.
	LastError   string     `json:"last_error,omitempty"`
	LastErrorAt *time.Time `json:"last_error_at,omitempty"`
	// Problem says, in words, what is wrong and what to do; empty when
	// nothing is.
	Problem string `json:"problem,omitempty"`
}

// judgeErrorFresh is how recent a judge's failure must be to count as a
// problem now.
const judgeErrorFresh = 5 * time.Minute

// Judges are the judge models the policies in force ask, each with what is
// wrong with it, if anything.
func (a *API) Judges() []JudgeView {
	if a.guard == nil || a.guard.Library == nil {
		return nil
	}
	lib := a.guard.Library()
	if !lib.On() {
		// Off: no judge is asked, so none is missing.
		return nil
	}
	byModel := map[string]*JudgeView{}
	for name, p := range lib.Policies {
		judges := p.Judge.Judges()
		for i, model := range judges {
			v := byModel[model]
			if v == nil {
				v = &JudgeView{Model: model}
				byModel[model] = v
			}
			if !slices.Contains(v.Policies, name) {
				v.Policies = append(v.Policies, name)
			}
			if b := p.Judge.Budget().Milliseconds(); v.BudgetMs == 0 || b < v.BudgetMs {
				v.BudgetMs = b
			}
			if p.Fail == "closed" && p.Mode != "shadow" && !slices.Contains(v.Closed, name) {
				v.Closed = append(v.Closed, name)
			}
			for _, after := range judges[i+1:] {
				if !slices.Contains(v.Then, after) {
					v.Then = append(v.Then, after)
				}
			}
		}
	}
	ready := map[string]int{}
	for _, rep := range a.ctl.Reconciler().Replicas() {
		if rep.State == reconciler.Ready {
			ready[rep.Model]++
		}
	}
	var errs map[string]guard.JudgeError
	if a.guard.JudgeErrors != nil {
		errs = a.guard.JudgeErrors()
	}
	out := make([]JudgeView, 0, len(byModel))
	for _, v := range byModel {
		sort.Strings(v.Policies)
		sort.Strings(v.Closed)
		if m, ok := a.ctl.Catalog().Resolve(v.Model); ok && m.Name == v.Model {
			v.InCatalog, v.Judge = true, m.IsGuard()
			v.Wanted, v.Ready = m.Scaling.Min, ready[m.Name]
		}
		if e, ok := errs[v.Model]; ok {
			at := e.At
			v.LastError, v.LastErrorAt = e.Err, &at
		}
		v.Problem = judgeProblem(*v)
		out = append(out, *v)
	}
	slices.SortFunc(out, func(x, y JudgeView) int {
		switch {
		case x.Model < y.Model:
			return -1
		case x.Model > y.Model:
			return 1
		}
		return 0
	})
	return out
}

func judgeProblem(v JudgeView) string {
	consequence := "requests its policies judge are served unjudged"
	switch {
	case len(v.Then) > 0:
		consequence = "its policies ask " + strings.Join(v.Then, ", else ") + " instead"
	case len(v.Closed) > 0:
		consequence = "requests its policies judge are refused (guard_unavailable)"
	}
	switch {
	case !v.InCatalog:
		return fmt.Sprintf("not in the catalog, so %s: llmgw ctl guard judge enable", consequence)
	case !v.Judge:
		return fmt.Sprintf("in the catalog as a model clients use, not as a judge (role: guard), so %s: "+
			"llmgw ctl guard judge enable --replace", consequence)
	case v.Ready == 0:
		return fmt.Sprintf("no replica ready, so %s: llmgw ctl replicas --model %s", consequence, v.Model)
	case v.LastErrorAt != nil && time.Since(*v.LastErrorAt) < judgeErrorFresh:
		return fmt.Sprintf("gave no verdict %s ago (%s)", time.Since(*v.LastErrorAt).Round(time.Second), v.LastError)
	}
	return ""
}

func (a *API) listJudges(w http.ResponseWriter, _ *http.Request) {
	if a.guard == nil {
		writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
		return
	}
	writeJSON(w, http.StatusOK, a.Judges())
}

// JudgeAskRequest is a text to have a judge read.
type JudgeAskRequest struct {
	// Model is the judge; empty is the one the policies ask, when they ask one.
	Model string `json:"model,omitempty"`
	Text  string `json:"text"`
	// Answer, when given, is read as the assistant's reply to Text.
	Answer string `json:"answer,omitempty"`
}

// JudgeAskView is what the judge answered.
type JudgeAskView struct {
	Model string `json:"model"`
	// Verdict is safe, unsafe or controversial, as the gateway reads the
	// answer; a policy's judge.controversial says what controversial counts as.
	Verdict    string   `json:"verdict"`
	Categories []string `json:"categories,omitempty"`
	// Raw is the judge's answer as it gave it.
	Raw    string `json:"raw"`
	TookMs int64  `json:"took_ms"`
}

func (a *API) askJudge(w http.ResponseWriter, r *http.Request) {
	if a.guard == nil || a.guard.Ask == nil {
		writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
		return
	}
	var req JudgeAskRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || strings.TrimSpace(req.Text) == "" {
		writeError(w, http.StatusBadRequest, "give the text for the judge to read")
		return
	}
	if req.Model == "" {
		judges := a.Judges()
		if len(judges) == 0 {
			writeError(w, http.StatusBadRequest, "no policy asks a judge: give the model")
			return
		}
		req.Model = judges[0].Model
	}
	ans, err := a.guard.Ask(r.Context(), req.Model, req.Text, req.Answer)
	v := JudgeAskView{Model: req.Model, Raw: ans.Raw, Categories: ans.Categories, TookMs: ans.Took.Milliseconds()}
	if err != nil {
		writeError(w, http.StatusBadGateway, fmt.Sprintf("the judge %s gave no verdict: %v (it answered %q)", req.Model, err, ans.Raw))
		return
	}
	switch {
	case ans.Unsafe:
		v.Verdict = "unsafe"
	case ans.Controversial:
		// Qwen3Guard's middle verdict: what it counts as is the policy's.
		v.Verdict = "controversial"
	default:
		v.Verdict = "safe"
	}
	writeJSON(w, http.StatusOK, v)
}

// JudgeCheckRequest asks how a judge keeps up: Burst judgments at once after
// three in turn; zero is two for each ready replica.
type JudgeCheckRequest struct {
	Model string `json:"model,omitempty"`
	Burst int    `json:"burst,omitempty"`
}

// JudgeCheckView is how long a judgment took, alone and among others, set
// against the time a request waits for one.
type JudgeCheckView struct {
	Model    string `json:"model"`
	BudgetMs int64  `json:"budget_ms"`
	Ready    int    `json:"ready_replicas"`
	// TookMs are three judgments of the same text, one after another.
	TookMs    []int64 `json:"took_ms"`
	SlowestMs int64   `json:"slowest_ms"`
	// Burst judgments at once: the slowest of them, how many gave no
	// verdict, and how many a second the judges read.
	Burst          int     `json:"burst"`
	BurstSlowestMs int64   `json:"burst_slowest_ms"`
	BurstFailed    int     `json:"burst_failed"`
	PerSecond      float64 `json:"judgments_per_second"`
	// Verdict is ok, tight (past half the budget alone, or past it among
	// others) or slow (past it alone); Advice what to do when not ok.
	Verdict string   `json:"verdict"`
	Advice  []string `json:"advice,omitempty"`
}

// judgeCheckText is a prompt of about 280 tokens, the length the judges
// were measured at: ordinary work, so the verdict is safe and the time is
// what a request pays.
const judgeCheckText = "Our support team has been asked to write a short guide for new customers who are setting " +
	"up the gateway for the first time. The guide should explain how to create an API key in the portal, how to " +
	"choose a model from the list the portal shows, and how to send a first request with curl or with the OpenAI " +
	"Python client by changing only the base URL and the key. It should also say what the usage page shows, how " +
	"budgets and rate limits are set by an administrator for each team, and what a person should do when a " +
	"request is refused because a limit was reached. Please keep the tone friendly and plain, use short " +
	"paragraphs, and add a small table that lists the three most common error codes with what each means and " +
	"what to try next. The audience is technical but new to the product, so avoid internal names, and end with " +
	"a short checklist the reader can follow on their first day. Write it in English, about four hundred words."

func (a *API) checkJudge(w http.ResponseWriter, r *http.Request) {
	if a.guard == nil || a.guard.Ask == nil {
		writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
		return
	}
	var req JudgeCheckRequest
	_ = json.NewDecoder(r.Body).Decode(&req)
	var jv *JudgeView
	for _, j := range a.Judges() {
		if req.Model == "" || j.Model == req.Model {
			j := j
			jv = &j
			break
		}
	}
	if jv == nil {
		writeError(w, http.StatusBadRequest, "no policy asks a judge of that name")
		return
	}
	if jv.Ready == 0 {
		writeError(w, http.StatusConflict, "the judge "+jv.Model+" has no replica ready: "+jv.Problem)
		return
	}
	v := JudgeCheckView{Model: jv.Model, BudgetMs: jv.BudgetMs, Ready: jv.Ready, Burst: req.Burst}
	if v.Burst <= 0 {
		v.Burst = 2 * jv.Ready
	}
	v.Burst = min(v.Burst, 32)
	var mu sync.Mutex
	start := time.Now()
	// Each text begins differently: the engine keeps what it read before,
	// and the same text again is answered from that in a tenth of the time
	// a request's own prompt takes.
	n := 0
	text := func() string {
		mu.Lock()
		defer mu.Unlock()
		n++
		return fmt.Sprintf("Request %d of a check begun at %s.\n", n, start.Format(time.RFC3339Nano)) + judgeCheckText
	}
	for range 3 {
		ans, err := a.guard.Ask(r.Context(), jv.Model, text(), "")
		if err != nil {
			writeError(w, http.StatusBadGateway, fmt.Sprintf("the judge %s gave no verdict: %v", jv.Model, err))
			return
		}
		v.TookMs = append(v.TookMs, ans.Took.Milliseconds())
		v.SlowestMs = max(v.SlowestMs, ans.Took.Milliseconds())
	}
	var wg sync.WaitGroup
	start = time.Now()
	for range v.Burst {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ans, err := a.guard.Ask(r.Context(), jv.Model, text(), "")
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				v.BurstFailed++
				return
			}
			v.BurstSlowestMs = max(v.BurstSlowestMs, ans.Took.Milliseconds())
		}()
	}
	wg.Wait()
	if wall := time.Since(start).Seconds(); wall > 0 {
		v.PerSecond = float64(v.Burst-v.BurstFailed) / wall
	}
	v.Verdict, v.Advice = judgeVerdict(v, *jv)
	writeJSON(w, http.StatusOK, v)
}

// judgeVerdict reads a check: a judgment alone past the budget refuses every
// request under fail: closed; past half of it, or past it among others, a
// busy moment does.
func judgeVerdict(v JudgeCheckView, jv JudgeView) (string, []string) {
	verdict := "ok"
	switch {
	case v.BudgetMs > 0 && v.SlowestMs > v.BudgetMs:
		verdict = "slow"
	case v.BudgetMs > 0 && (v.SlowestMs*2 > v.BudgetMs || v.BurstSlowestMs > v.BudgetMs || v.BurstFailed > 0):
		verdict = "tight"
	}
	if verdict == "ok" {
		return verdict, nil
	}
	refused := "served unjudged"
	if len(jv.Closed) > 0 {
		refused = "refused"
	}
	advice := []string{
		fmt.Sprintf("a judgment past %d ms is %s: each request is judged once for its prompt and again every 600 characters of its answer", v.BudgetMs, refused),
		"more cores for each judge: llmgw ctl model resources " + v.Model + " --cores 4, then llmgw ctl restart --model " + v.Model,
		"more judges, on machines of their own: llmgw ctl scale " + v.Model + " N, or judge machines at install",
		"a judge on a GPU: llmgw ctl guard judge use qwen3guard-4b " + v.Model,
		"a longer wait: guard.policies.<policy>.judge.latency_budget in llmgw ctl config edit",
	}
	if len(jv.Closed) > 0 {
		advice = append(advice, "or fail: open, which serves what the judge did not read, flagged")
	}
	return verdict, advice
}