llmgw/internal/cli/guard_judge.go

package cli

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

	"github.com/spf13/cobra"

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

// The judge the guardrails are installed with: Meta's Llama Guard 3 1B,
// from the repository server init fetches it from.
const (
	guardModelRepo = "QuantFactory/Llama-Guard-3-1B-GGUF"
	guardModelFile = "Llama-Guard-3-1B.Q4_K_M.gguf"
)

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).",
	}
	c.AddCommand(newGuardJudgeStatusCmd(client, asJSON), newGuardJudgeEnableCmd(client))
	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)
			}
			if err := w.Flush(); err != nil {
				return err
			}
			strays, _ := strayJudges(cmd, client, js)
			for _, s := range strays {
				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", s, s)
			}
			return nil
		},
	}
}

// strayJudges are models whose weights are a judge's but which are in the
// catalog as models clients use, under a name no policy asks: what serving
// Llama Guard from the mirror page did before it knew better.
func strayJudges(cmd *cobra.Command, client func() *ctlClient, js []adminapi.JudgeView) ([]string, 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 []string
	for _, m := range models {
		if !asked[m.Name] && strings.Contains(strings.ToLower(m.Source), "llama-guard") {
			out = append(out, m.Name)
		}
	}
	return out, nil
}

func newGuardJudgeEnableCmd(client func() *ctlClient) *cobra.Command {
	var (
		model, repo, file string
		replicas, cores   int
		replace, detach   bool
	)
	c := &cobra.Command{
		Use:   "enable",
		Short: "Serve the judge the policies ask, as a judge: Llama Guard 3 1B on CPUs",
		Long: "Fetches Llama Guard 3 1B into the mirror (once), adds it to the catalog under\n" +
			"the name the policies give its judge, as a judge (role: guard,\n" +
			"format llamaguard3), on CPUs so it takes no card, and waits for a replica\n" +
			"to answer. With no --model, the judge the policies ask that is missing.\n\n" +
			"A model of that name already in the catalog as a model clients use is\n" +
			"refused; --replace removes it first. An offline cluster gets its judge\n" +
			"from its bundle with server init.",
		Example: "  llmgw ctl guard judge enable\n" +
			"  llmgw ctl guard judge enable --replace",
		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
			}
			var target *adminapi.JudgeView
			for i, j := range js {
				if (model != "" && j.Model == model) || (model == "" && (!j.InCatalog || !j.Judge)) {
					target = &js[i]
					break
				}
			}
			switch {
			case target == nil && model != "":
				target = &adminapi.JudgeView{Model: model}
				fmt.Fprintf(out, "no policy asks %s yet; it is served, and a policy's judge.model names it\n", model)
			case target == nil && len(js) == 0:
				return errors.New("no policy asks a judge: give one a judge (llmgw ctl config edit, guard.policies.<name>.judge.model)")
			case target == nil:
				fmt.Fprintln(out, "every judge the policies ask is in the catalog, as a judge")
				return nil
			}
			if target.InCatalog && target.Judge {
				fmt.Fprintf(out, "%s is in the catalog as a judge already; llmgw ctl guard judge status says how it is\n", target.Model)
				return nil
			}
			if target.InCatalog && !target.Judge {
				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", target.Model)
				}
				var res map[string]string
				if err := client().delete(ctx, "/api/v1/models/"+target.Model, &res); err != nil {
					return fmt.Errorf("removing %s: %w", target.Model, err)
				}
				fmt.Fprintf(out, "%s removed\n", target.Model)
			}
			req := modelserve.Request{
				Repo: repo, File: file, Name: target.Model, GuardFormat: "llamaguard3",
				Description: "The guardrails' judge: reads prompts and answers for unsafe content. Not offered to clients.",
				CPU:         true, Replicas: replicas, CPUCores: cores,
				CtxPool: 8192, ContextLength: 4096,
			}
			body, _ := json.Marshal(req)
			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 strays, _ := strayJudges(cmd, client, js); len(strays) > 0 {
				fmt.Fprintf(out, "\nno policy asks %s, served from a judge's weights as a model clients use: "+
					"llmgw ctl model remove %s\n", strings.Join(strays, ", "), strings.Join(strays, " "))
			}
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&model, "model", "", "the judge to serve (default: the one the policies ask that is missing)")
	f.StringVar(&repo, "repo", guardModelRepo, "the registry repository its weights come from")
	f.StringVar(&file, "file", guardModelFile, "the file in it")
	f.IntVar(&replicas, "replicas", 2, "replicas: every judged request waits on one")
	f.IntVar(&cores, "cores", 4, "CPU cores a replica 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
}