llmgw/internal/cli/guard_judge.go

package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"slices"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
)

func newGuardJudgeCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "judge",
		Short: "The judge models the policies ask: whether each can answer, and serving one",
		Long: "A policy with a judge has a model read what patterns cannot. The judge must\n" +
			"be in the catalog under the name the policy gives, as a judge (role:\n" +
			"guard), with a replica ready; otherwise what the policy judges is refused\n" +
			"(fail: closed, guard_unavailable) or served unjudged (fail: open).\n\n" +
			"A policy may name several, asked in turn (judge.models): a judge on a GPU\n" +
			"first and one on CPUs after it is the GPU's speed while it is there, and\n" +
			"the CPUs' when it is not.\n\n" +
			"Judges llmgw knows where to fetch, which a policy may name and the leader\n" +
			"then fetches and serves by itself: " + knownJudgeNames() + ".",
	}
	c.AddCommand(newGuardJudgeStatusCmd(client, asJSON), newGuardJudgeEnableCmd(client),
		newGuardJudgeUseCmd(client), newGuardJudgeAskCmd(client, asJSON))
	return c
}

func newGuardJudgeStatusCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "status",
		Short: "Each judge the policies ask, and what is wrong with it",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var js []adminapi.JudgeView
			if err := client().get(cmd.Context(), "/api/v1/guard/judges", &js); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(js)
			}
			out := cmd.OutOrStdout()
			if len(js) == 0 {
				fmt.Fprintln(out, "no policy asks a judge")
				return nil
			}
			w := tabTo(out)
			fmt.Fprintln(w, "JUDGE\tPOLICIES\tIN CATALOG\tAS JUDGE\tREADY\tSTATE")
			for _, j := range js {
				state := "answering"
				if j.Problem != "" {
					state = j.Problem
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d\t%s\n", j.Model, strings.Join(j.Policies, ","),
					yesNo(j.InCatalog), yesNo(j.Judge), j.Ready, j.Wanted, state)
			}
			for _, j := range js {
				if len(j.Then) > 0 {
					fmt.Fprintf(w, "\n%s gives no verdict: %s is asked instead", j.Model, strings.Join(j.Then, ", else "))
				}
			}
			if err := w.Flush(); err != nil {
				return err
			}
			strays, _ := strayJudges(cmd, client, js)
			sayStrays(out, strays)
			return nil
		},
	}
}

// strayJudges are models whose weights are a judge's, under a name no policy
// asks: a judge the policies moved on from, or one served from the mirror
// page as a model clients use, as before that page knew better.
func strayJudges(cmd *cobra.Command, client func() *ctlClient, js []adminapi.JudgeView) ([]adminapi.ModelView, error) {
	var models []adminapi.ModelView
	if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
		return nil, err
	}
	asked := map[string]bool{}
	for _, j := range js {
		asked[j.Model] = true
	}
	var out []adminapi.ModelView
	for _, m := range models {
		if _, judge := guard.FormatOfWeights(m.Source); !asked[m.Name] && judge {
			out = append(out, m)
		}
	}
	return out, nil
}

// sayStrays says what each stray judge is and how to be rid of it.
func sayStrays(out io.Writer, strays []adminapi.ModelView) {
	for _, m := range strays {
		if m.Role == catalog.RoleGuard {
			fmt.Fprintf(out, "\n%s is a judge no policy asks, still running: llmgw ctl model remove %s\n", m.Name, m.Name)
			continue
		}
		fmt.Fprintf(out, "\n%s looks like a judge's weights but is served as a model clients use, and no "+
			"policy asks it: llmgw ctl model remove %s\n", m.Name, m.Name)
	}
}

func newGuardJudgeEnableCmd(client func() *ctlClient) *cobra.Command {
	var (
		model, repo, file, format, gpuClass string
		replicas, cores                     int
		gpu, replace, detach                bool
	)
	c := &cobra.Command{
		Use:   "enable",
		Short: "Serve a judge: Qwen3Guard-Gen 0.6B on CPUs, or with --gpu the 4B on a GPU",
		Long: "Fetches a judge's weights into the mirror (once; what is there already, its\n" +
			"checksums matching, is not fetched again), adds it to the catalog as a judge\n" +
			"(role: guard) and waits for a replica to answer. Which judge: --model, or\n" +
			"with --gpu " + defaultGPUJudge + ", or else the first the policies ask that is\n" +
			"missing, or else " + defaultJudge + ".\n\n" +
			"A judge on CPUs takes no card. One on a GPU takes one card, as a standard\n" +
			"model: from nothing more important, so it waits while the cards are busy;\n" +
			"put a judge on CPUs after it in the policies (llmgw ctl guard judge use) so\n" +
			"there is always one to ask.\n\n" +
			"Known judges: " + knownJudgeNames() + ". Another is served with --repo,\n" +
			"--file and --format. A model of that name already in the catalog as a\n" +
			"model clients use is refused; --replace removes it first. An offline\n" +
			"cluster gets its judge from its bundle with server init.",
		Example: "  llmgw ctl guard judge enable\n" +
			"  llmgw ctl guard judge enable --gpu\n" +
			"  llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			ctx := cmd.Context()
			out := cmd.OutOrStdout()
			var js []adminapi.JudgeView
			if err := client().get(ctx, "/api/v1/guard/judges", &js); err != nil {
				return err
			}
			name := model
			if name == "" && gpu {
				name = defaultGPUJudge
			}
			if name == "" {
				for _, j := range js {
					if !j.InCatalog || !j.Judge {
						name = j.Model
						break
					}
				}
			}
			if name == "" {
				name = defaultJudge
			}

			k, known := knownJudgeNamed(name)
			switch {
			case repo != "":
				f := format
				if f == "" {
					f, _ = guard.FormatOfWeights(repo + "/" + file)
				}
				if f == "" || !guard.KnownFormat(f) {
					return fmt.Errorf("say how its answers are read: --format %s", strings.Join(guard.Formats, "|"))
				}
				k = knownJudge{Name: name, Repo: repo, File: file, Format: f, GPU: gpu}
			case !known:
				return fmt.Errorf("%s is not a judge llmgw knows where to fetch (%s): give --repo, --file and --format", name, knownJudgeNames())
			}
			if gpu {
				k.GPU = true
			}

			var models []adminapi.ModelView
			if err := client().get(ctx, "/api/v1/models", &models); err != nil {
				return err
			}
			if i := slices.IndexFunc(models, func(m adminapi.ModelView) bool { return m.Name == name }); i >= 0 {
				if models[i].Role == catalog.RoleGuard {
					fmt.Fprintf(out, "%s is in the catalog as a judge already; llmgw ctl guard judge status says how it is\n", name)
					return nil
				}
				if !replace {
					return fmt.Errorf("%s is in the catalog as a model clients use, not as a judge: "+
						"--replace removes it and serves it again as a judge", name)
				}
				var res map[string]string
				if err := client().delete(ctx, "/api/v1/models/"+name, &res); err != nil {
					return fmt.Errorf("removing %s: %w", name, err)
				}
				fmt.Fprintf(out, "%s removed\n", name)
			}

			opts := judgeServeOptions{Replicas: replicas, Cores: cores, GPUClass: gpuClass}
			if k.GPU && opts.GPUClass == "" {
				var s adminapi.Summary
				if err := client().get(ctx, "/api/v1/summary", &s); err != nil {
					return err
				}
				if opts.GPUClass = roomiestGPUClass(s.GPUs); opts.GPUClass == "" {
					return fmt.Errorf("this cluster has no GPU for %s: serve %s on CPUs instead", name, defaultJudge)
				}
			}
			fmt.Fprintf(out, "serving %s\n", k)
			body, _ := json.Marshal(k.serveRequest(opts))
			var job modelserve.Job
			if err := client().patient(30*time.Second).post(ctx, "/api/v1/models/serve", body, &job); err != nil {
				return err
			}
			if detach {
				return detached(cmd, job)
			}
			if err := followJob(cmd, client, job, "is the judge"); err != nil {
				return err
			}
			if !slices.ContainsFunc(js, func(j adminapi.JudgeView) bool { return j.Model == name }) {
				fmt.Fprintf(out, "\nno policy asks %s yet: llmgw ctl guard judge use %s", name, name)
				if k.GPU {
					fmt.Fprintf(out, " %s", defaultJudge)
				}
				fmt.Fprintln(out)
			}
			if strays, _ := strayJudges(cmd, client, js); len(strays) > 0 {
				sayStrays(out, slices.DeleteFunc(strays, func(m adminapi.ModelView) bool { return m.Name == name }))
			}
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&model, "model", "", "the judge to serve: "+knownJudgeNames()+", or a name for --repo's")
	f.BoolVar(&gpu, "gpu", false, "serve it on a GPU ("+defaultGPUJudge+" when no --model)")
	f.StringVar(&gpuClass, "gpu-class", "", "the GPUs it runs on (default: the class with the most free cards)")
	f.StringVar(&repo, "repo", "", "the registry repository of a judge llmgw does not know")
	f.StringVar(&file, "file", "", "the GGUF file in --repo")
	f.StringVar(&format, "format", "", "how --repo's judge answers: "+strings.Join(guard.Formats, ", "))
	f.IntVar(&replicas, "replicas", 0, "replicas (default: 2 on CPUs, as every judged request waits on one; 1 on a GPU)")
	f.IntVar(&cores, "cores", 4, "CPU cores a replica on CPUs computes on")
	f.BoolVar(&replace, "replace", false, "remove a model of that name that is not a judge first")
	f.BoolVar(&detach, "detach", false, "start it and return; llmgw ctl jobs follows it")
	return c
}

// roomiestGPUClass is the class with the most free cards, else the most
// cards; empty when the cluster has none.
func roomiestGPUClass(classes []adminapi.ClassUse) string {
	best := -1
	for i, c := range classes {
		if c.Total == 0 {
			continue
		}
		if best < 0 || c.Free() > classes[best].Free() || (c.Free() == classes[best].Free() && c.Total > classes[best].Total) {
			best = i
		}
	}
	if best < 0 {
		return ""
	}
	return classes[best].Class
}

func newGuardJudgeUseCmd(client func() *ctlClient) *cobra.Command {
	var policies []string
	var note string
	c := &cobra.Command{
		Use:   "use JUDGE...",
		Short: "Set the judges the policies ask, in the order they are asked",
		Long: "Sets judge.models of every policy that has a judge, or of those given with\n" +
			"--policy: the first judge that answers gives the verdict, and the next is\n" +
			"asked when one has no replica serving, fails, or does not answer within\n" +
			"the policy's latency_budget. A judge on a GPU first and one on CPUs after\n" +
			"it is the GPU's speed while it is there, and the CPUs' when it is not.\n\n" +
			"A judge llmgw knows (" + knownJudgeNames() + ")\n" +
			"that is not in the catalog is fetched and served by the leader; another\n" +
			"is served first with llmgw ctl guard judge enable --repo. In force on\n" +
			"every control plane within seconds.",
		Example: "  llmgw ctl guard judge use " + defaultJudge + "\n" +
			"  llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge + "\n" +
			"  llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge + " --policy strict",
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			for i, a := range args {
				if slices.Contains(args[:i], a) {
					return fmt.Errorf("%s is named twice", a)
				}
			}
			if note == "" {
				note = "judges: " + strings.Join(args, ", else ")
			}
			list := "[" + strings.Join(args, ", ") + "]"
			return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
				names := policies
				if len(names) == 0 {
					for _, p := range doc.MapKeys([]string{"guard", "policies"}) {
						if doc.Node([]string{"guard", "policies", p, "judge"}) != nil {
							names = append(names, p)
						}
					}
					if len(names) == 0 {
						return errors.New("no policy has a judge: name one with --policy, and it is given one " +
							"that blocks what the judge finds unsafe in requests (judge.input: block)")
					}
				}
				for _, p := range names {
					base := []string{"guard", "policies", p}
					if doc.Node(base) == nil {
						return fmt.Errorf("there is no policy %s (llmgw ctl guard list)", p)
					}
					judge := append(slices.Clone(base), "judge")
					if doc.Node(judge) == nil {
						if err := doc.SetYAML(append(slices.Clone(judge), "input"), "block"); err != nil {
							return err
						}
					}
					doc.Unset(append(slices.Clone(judge), "model"))
					if err := doc.SetYAML(append(slices.Clone(judge), "models"), list); err != nil {
						return err
					}
					cats := append(slices.Clone(judge), "categories")
					var have []string
					if n := doc.Node(cats); n != nil {
						for _, c := range n.Content {
							have = append(have, c.Value)
						}
					}
					if next := categoriesFor(args, have); next != nil {
						if err := doc.SetYAML(cats, "["+strings.Join(next, ", ")+"]"); err != nil {
							return err
						}
						fmt.Fprintf(cmd.OutOrStdout(), "%s: categories [%s] are not what %s name; it counts [%s] instead\n",
							p, strings.Join(have, ", "), strings.Join(args, " or "), strings.Join(next, ", "))
					}
				}
				return nil
			})
		},
	}
	c.Flags().StringSliceVar(&policies, "policy", nil, "the policies to change (default: every one with a judge)")
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

func newGuardJudgeAskCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var model, answer string
	c := &cobra.Command{
		Use:   "ask TEXT",
		Short: "Have a judge read a text and say what it answered",
		Long: "For finding out why a request was refused: the judge reads TEXT as a user's\n" +
			"message (with --answer, ANSWER as the reply to it), past the cache, and its\n" +
			"answer is shown as it gave it and as the gateway reads it — \"unsafe\" and\n" +
			"the hazards it named, S6 Specialized Advice and the like. Nothing is kept:\n" +
			"not in the audit, the usage or the logs.",
		Example: "  llmgw ctl guard judge ask 'How do I sort a list in Python?'\n" +
			"  llmgw ctl guard judge ask 'Write a poem' --answer 'Roses are red'",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			body, _ := json.Marshal(adminapi.JudgeAskRequest{Model: model, Text: args[0], Answer: answer})
			var v adminapi.JudgeAskView
			if err := client().patient(45*time.Second).post(cmd.Context(), "/api/v1/guard/judges/ask", body, &v); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(v)
			}
			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s: %s", v.Model, v.Verdict)
			if len(v.Categories) > 0 {
				fmt.Fprintf(out, " (%s)", strings.Join(v.Categories, ", "))
			}
			fmt.Fprintf(out, ", %d ms\nit answered: %q\n", v.TookMs, v.Raw)
			return nil
		},
	}
	c.Flags().StringVar(&model, "model", "", "the judge (default: the one the policies ask)")
	c.Flags().StringVar(&answer, "answer", "", "an assistant's reply to TEXT, judged as output is")
	return c
}