llmgw/internal/guard/judge.go

package guard

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard/detect"
)

// JudgePolicy has a judge model on the cluster read text for what patterns
// cannot: violence, self-harm, abuse — whatever the model was trained on.
type JudgePolicy struct {
	// Model is a catalog model with role: guard.
	Model string `koanf:"model"`
	// Input and Output are what an unsafe verdict does: flag or block.
	// Empty does not ask the judge about that direction.
	Input  string `koanf:"input"`
	Output string `koanf:"output"`
	// Categories, when set, are the only ones that count: Llama Guard's
	// codes (S1, S10) or the names the model reports.
	Categories []string `koanf:"categories"`
	// LatencyBudget is how long the judge is waited for before the policy's
	// fail applies. Zero is 250ms.
	LatencyBudget time.Duration `koanf:"latency_budget"`
}

// DefaultJudgeBudget is how long a judge is waited for when a policy does
// not say.
const DefaultJudgeBudget = 250 * time.Millisecond

func (j *JudgePolicy) budget() time.Duration {
	if j.LatencyBudget > 0 {
		return j.LatencyBudget
	}
	return DefaultJudgeBudget
}

func (j *JudgePolicy) validate() error {
	if j.Model == "" {
		return errors.New("judge.model is required")
	}
	for k, a := range map[string]string{"input": j.Input, "output": j.Output} {
		switch a {
		case "", "flag", "block":
		default:
			return fmt.Errorf("judge.%s is %q: flag or block", k, a)
		}
	}
	if j.Input == "" && j.Output == "" {
		return errors.New("judge: set input or output (flag or block), or it judges nothing")
	}
	return nil
}

// mergeJudge keeps the first judge model asked for, the strictest actions,
// every category and the shortest wait.
func mergeJudge(a, b *JudgePolicy) *JudgePolicy {
	if b == nil {
		return a
	}
	if a == nil {
		c := *b
		c.Categories = slices.Clone(b.Categories)
		return &c
	}
	out := *a
	out.Input = stricterOf(a.Input, b.Input, []string{"", "flag", "block"})
	out.Output = stricterOf(a.Output, b.Output, []string{"", "flag", "block"})
	out.Categories = union(slices.Clone(a.Categories), b.Categories)
	if b.LatencyBudget > 0 && (a.LatencyBudget == 0 || b.LatencyBudget < a.LatencyBudget) {
		out.LatencyBudget = b.LatencyBudget
	}
	return &out
}

// Caller sends a body to a model on the cluster at a path —
// /v1/chat/completions or /v1/completions — and returns the answer's body:
// the router, past admission, accounting and guardrails.
type Caller func(ctx context.Context, model, path string, body []byte) ([]byte, error)

// ModelFormat says how a judge model's answers are read, and the category
// names it reports; ok is false for a model that is not a judge.
type ModelFormat func(model string) (format string, categories []string, ok bool)

// judgment is what a judge said.
type judgment struct {
	unsafe     bool
	categories []string
}

// judge asks a judge model about text, with a cache of recent answers.
type judge struct {
	call   Caller
	format ModelFormat
	m      *gateMetrics

	mu    sync.Mutex
	cache map[[32]byte]cached
}

type cached struct {
	j  judgment
	at time.Time
}

const (
	judgeCacheTTL  = 5 * time.Minute
	judgeCacheSize = 10_000
)

// SetJudge lets policies with a judge: section ask a model on the cluster.
func (g *Gate) SetJudge(call Caller, format ModelFormat) {
	g.judger = &judge{call: call, format: format, m: g.m, cache: map[[32]byte]cached{}}
}

// errNoJudge is a policy naming a judge the gateway cannot ask, and
// errJudgeFormat a judge whose answer could not be read.
var (
	errNoJudge     = errors.New("no judge model to ask")
	errJudgeFormat = errors.New("the judge's answer could not be read")
)

// ask has the judge read conversation turns — the user's, and the
// assistant's answer when judging output.
func (j *judge) ask(ctx context.Context, jp *JudgePolicy, user, answer string) (res judgment, err error) {
	if j == nil {
		return judgment{}, errNoJudge
	}
	start := time.Now()
	cachedHit := false
	defer func() {
		if !cachedHit && j.m != nil {
			j.m.judged(jp.Model, time.Since(start), err)
		}
	}()
	format, names, ok := j.format(jp.Model)
	if !ok {
		return judgment{}, fmt.Errorf("%w: %q is not a model with role: guard", errNoJudge, jp.Model)
	}
	key := sha256.Sum256([]byte(jp.Model + "\x00" + detect.Normalize(user) + "\x00" + detect.Normalize(answer)))
	j.mu.Lock()
	if c, ok := j.cache[key]; ok && time.Since(c.at) < judgeCacheTTL {
		j.mu.Unlock()
		cachedHit = true
		return c.j, nil
	}
	j.mu.Unlock()

	ctx, cancel := context.WithTimeout(ctx, jp.budget())
	defer cancel()
	path, body := judgeRequest(jp.Model, format, names, user, answer)
	out, err := j.call(ctx, jp.Model, path, body)
	if err != nil {
		return judgment{}, err
	}
	res, err = parseJudgment(format, names, out, answer)
	if err != nil {
		return judgment{}, err
	}
	j.mu.Lock()
	if len(j.cache) >= judgeCacheSize {
		for k, c := range j.cache {
			if time.Since(c.at) >= judgeCacheTTL || len(j.cache) >= judgeCacheSize {
				delete(j.cache, k)
			}
		}
	}
	j.cache[key] = cached{j: res, at: time.Now()}
	j.mu.Unlock()
	return res, nil
}

// counts reports whether a judgment is one the policy acts on.
func (jp *JudgePolicy) counts(res judgment) (bool, string) {
	if !res.unsafe {
		return false, ""
	}
	if len(jp.Categories) == 0 {
		return true, strings.Join(res.categories, ",")
	}
	for _, c := range res.categories {
		if slices.ContainsFunc(jp.Categories, func(w string) bool { return strings.EqualFold(w, c) }) {
			return true, c
		}
	}
	return false, ""
}

// llamaGuard3Categories are Llama Guard 3's hazards, in the order of their
// codes.
var llamaGuard3Categories = []string{"Violent Crimes", "Non-Violent Crimes", "Sex Crimes", "Child Exploitation",
	"Defamation", "Specialized Advice", "Privacy", "Intellectual Property", "Indiscriminate Weapons", "Hate",
	"Self-Harm", "Sexual Content", "Elections"}

// llamaGuard3Answer is Llama Guard 3's prompt for judging an answer, written
// out rather than left to the model's chat template: an engine given a
// conversation that ends on the assistant's turn — llama.cpp — takes it as a
// reply to go on with, and the verdict comes out as the reply's next words, or
// not at all.
func llamaGuard3Answer(categories []string, user, answer string) string {
	if len(categories) == 0 {
		categories = llamaGuard3Categories
	}
	var b strings.Builder
	// No <|begin_of_text|>: the engine puts it first itself.
	b.WriteString("<|start_header_id|>user<|end_header_id|>\n\nTask: Check if there is unsafe content in " +
		"'Agent' messages in conversations according our safety policy with the below categories.\n\n" +
		"<BEGIN UNSAFE CONTENT CATEGORIES>\n")
	for i, c := range categories {
		fmt.Fprintf(&b, "S%d: %s.\n", i+1, c)
	}
	b.WriteString("<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\n")
	if user != "" {
		b.WriteString("User: " + user + "\n\n")
	}
	b.WriteString("Agent: " + answer + "\n\n<END CONVERSATION>\n\n" +
		"Provide your safety assessment for ONLY THE LAST Agent message in the above conversation:\n" +
		" - First line must read 'safe' or 'unsafe'.\n" +
		" - If unsafe, a second line must include a comma-separated list of violated categories." +
		"<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n")
	return b.String()
}

// judgeRequest is the request a judge model is sent, and where.
func judgeRequest(model, format string, categories []string, user, answer string) (string, []byte) {
	if format == "llamaguard3" && answer != "" {
		b, _ := json.Marshal(map[string]any{"model": model, "prompt": llamaGuard3Answer(categories, user, answer),
			"max_tokens": 16, "temperature": 0, "stream": false})
		return "/v1/completions", b
	}
	var msgs []map[string]string
	switch format {
	case "llamaguard3":
		// The model's own chat template turns the turns into its prompt.
		msgs = append(msgs, map[string]string{"role": "user", "content": user})
		if answer != "" {
			msgs = append(msgs, map[string]string{"role": "assistant", "content": answer})
		}
	case "shieldgemma":
		who, text := "Human Question", user
		if answer != "" {
			who, text = "Chatbot Response", answer
		}
		principles := "* " + strings.Join(categories, "\n* ")
		if len(categories) == 0 {
			principles = "* No dangerous, hateful, harassing or sexually explicit content."
		}
		msgs = append(msgs, map[string]string{"role": "user", "content": "You are a policy expert trying to help determine " +
			"whether a message is in violation of the defined safety policies.\n\n<start_of_turn>\n" + who + ": " + text +
			"\n<end_of_turn>\n\nOur safety principle is defined in the below:\n\n" + principles +
			"\n\nDoes the " + strings.ToLower(who) + " violate the above principle? Your answer must start with 'Yes' or 'No'."})
	default: // json
		subject := "user message"
		text := user
		if answer != "" {
			subject, text = "assistant answer", answer
		}
		msgs = append(msgs,
			map[string]string{"role": "system", "content": "Classify whether the " + subject + " is unsafe. " +
				`Answer with JSON only: {"safe": true|false, "categories": ["..."]}.`},
			map[string]string{"role": "user", "content": text})
	}
	b, _ := json.Marshal(map[string]any{"model": model, "messages": msgs, "max_tokens": 32, "temperature": 0, "stream": false})
	return "/v1/chat/completions", b
}

// parseJudgment reads a judge model's answer.
//
// answer is the text that was judged, when it was an answer: llama.cpp
// takes a conversation ending on the assistant's turn as a reply to go on
// with, and its answer then begins with that turn — "Rome.safe".
func parseJudgment(format string, names []string, body []byte, answer string) (judgment, error) {
	var resp struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
			Text string `json:"text"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(body, &resp); err != nil || len(resp.Choices) == 0 {
		return judgment{}, fmt.Errorf("%w: not a completion", errJudgeFormat)
	}
	text := strings.TrimSpace(resp.Choices[0].Message.Content + resp.Choices[0].Text)
	if a := strings.TrimSpace(answer); a != "" {
		text = strings.TrimSpace(strings.TrimPrefix(text, a))
	}
	switch format {
	case "llamaguard3":
		lines := strings.Split(text, "\n")
		switch strings.ToLower(strings.TrimSpace(lines[0])) {
		case "safe":
			return judgment{}, nil
		case "unsafe":
			res := judgment{unsafe: true}
			if len(lines) > 1 {
				for _, c := range strings.Split(lines[1], ",") {
					c = strings.TrimSpace(c)
					if c == "" {
						continue
					}
					res.categories = append(res.categories, c)
					// And the name it stands for, when the model has them.
					var n int
					if _, err := fmt.Sscanf(c, "S%d", &n); err == nil && n >= 1 && n <= len(names) {
						res.categories = append(res.categories, names[n-1])
					}
				}
			}
			return res, nil
		}
	case "shieldgemma":
		l := strings.ToLower(text)
		switch {
		case strings.HasPrefix(l, "yes"):
			return judgment{unsafe: true, categories: []string{"policy"}}, nil
		case strings.HasPrefix(l, "no"):
			return judgment{}, nil
		}
	default:
		start, end := strings.Index(text, "{"), strings.LastIndex(text, "}")
		if start >= 0 && end > start {
			var v struct {
				Safe       *bool    `json:"safe"`
				Categories []string `json:"categories"`
			}
			if json.Unmarshal([]byte(text[start:end+1]), &v) == nil && v.Safe != nil {
				return judgment{unsafe: !*v.Safe, categories: v.Categories}, nil
			}
		}
	}
	// The judge's own answer, cut short, so an operator can see what form it
	// came in: a verdict, not the text judged.
	seen := []rune(text)
	if len(seen) > 40 {
		seen = seen[:40]
	}
	return judgment{}, fmt.Errorf("%w: not in the form %s has: %q", errJudgeFormat, format, string(seen))
}

// Moderate asks a judge model about one text, as a user's message, outside
// any policy: for /v1/moderations. It answers whether the model found it
// unsafe, and under which of its categories.
func Moderate(ctx context.Context, call Caller, model, format string, names []string, text string) (bool, []string, error) {
	path, body := judgeRequest(model, format, names, text, "")
	out, err := call(ctx, model, path, body)
	if err != nil {
		return false, nil, err
	}
	res, err := parseJudgment(format, names, out, "")
	if err != nil {
		return false, nil, err
	}
	return res.unsafe, res.categories, nil
}

// DefaultCategories are the names a judge format reports its categories by
// when the catalog gives none: Llama Guard 3's S1 to S13.
func DefaultCategories(format string) []string {
	if format == "llamaguard3" {
		return slices.Clone(llamaGuard3Categories)
	}
	return nil
}