llmgw/internal/guard/judge_test.go

package guard

import (
	"context"
	"encoding/json"
	"errors"
	"slices"
	"strings"
	"sync/atomic"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/prometheus/client_golang/prometheus/testutil"
)

// fakeJudge answers as a judge model would, saying unsafe to text with
// "attack" in it.
type fakeJudge struct {
	calls  atomic.Int64
	err    error
	format string
	delay  time.Duration
}

func (f *fakeJudge) call(ctx context.Context, _, _ string, body []byte) ([]byte, error) {
	f.calls.Add(1)
	if f.delay > 0 {
		select {
		case <-time.After(f.delay):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	if f.err != nil {
		return nil, f.err
	}
	unsafe := strings.Contains(string(body), "attack")
	var answer string
	switch f.format {
	case "shieldgemma":
		answer = map[bool]string{true: "Yes", false: "No"}[unsafe]
	case "json":
		answer = map[bool]string{true: `{"safe": false, "categories": ["violence"]}`, false: `{"safe": true}`}[unsafe]
	case "qwen3guard":
		answer = map[bool]string{true: "Safety: Unsafe\nCategories: Violent", false: "Safety: Safe\nCategories: None"}[unsafe]
		if strings.Contains(string(body), "a heist film") {
			answer = "Safety: Controversial\nCategories: Non-violent Illegal Acts"
		}
	default:
		answer = map[bool]string{true: "unsafe\nS1,S10", false: "safe"}[unsafe]
	}
	b, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"content": answer}}}})
	return b, nil
}

func judged(t *testing.T, p Policy, fj *fakeJudge) *Gate {
	g := gate(t, lib(map[string]Policy{"p": p}, "p"))
	if fj.format == "" {
		fj.format = "llamaguard3"
	}
	g.SetJudge(fj.call, func(m string) (string, []string, bool) {
		return fj.format, []string{"Violent Crimes"}, m == "guard"
	})
	return g
}

func TestAJudgeModelDecidesWhatPatternsCannot(t *testing.T) {
	ctx := context.Background()
	for _, format := range Formats {
		fj := &fakeJudge{format: format}
		g := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "block"}}, fj)
		if v, err := g.Input(ctx, Subject{}, chat("user", "plan an attack")); err != nil || v.Action != Block || v.Category != "judge" {
			t.Errorf("%s: unsafe: %+v %v", format, v, err)
		}
		if v, _ := g.Input(ctx, Subject{}, chat("user", "plan a picnic")); v.Action != Allow {
			t.Errorf("%s: safe: %+v", format, v)
		}
	}
	// Llama Guard's codes are named from the catalog, and only listed
	// categories count when a policy lists some.
	fj := &fakeJudge{}
	g := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "block", Categories: []string{"S2"}}}, fj)
	if v, _ := g.Input(ctx, Subject{}, chat("user", "plan an attack")); v.Action != Allow {
		t.Errorf("a category the policy does not list: %+v", v)
	}
	g = judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "flag", Categories: []string{"violent crimes"}}}, fj)
	if v, _ := g.Input(ctx, Subject{}, chat("user", "plan an attack")); v.Action != Flag || v.Rule != "Violent Crimes" {
		t.Errorf("by name: %+v", v)
	}
}

func TestAJudgeIsAskedOnceForTheSameText(t *testing.T) {
	fj := &fakeJudge{}
	g := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "flag"}}, fj)
	for range 3 {
		_, _ = g.Input(context.Background(), Subject{}, chat("user", "hello there"))
	}
	if fj.calls.Load() != 1 {
		t.Errorf("asked %d times", fj.calls.Load())
	}
}

func TestAJudgeThatDoesNotAnswerFailsAsThePolicySays(t *testing.T) {
	ctx := context.Background()
	slow := &fakeJudge{delay: time.Second}
	open := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "block", LatencyBudget: 20 * time.Millisecond}}, slow)
	start := time.Now()
	v, err := open.Input(ctx, Subject{}, chat("user", "plan an attack"))
	if err != nil || v.Action != Flag || v.Rule != "unavailable" || time.Since(start) > 500*time.Millisecond {
		t.Errorf("fail open: %+v %v after %s", v, err, time.Since(start))
	}
	closed := judged(t, Policy{Fail: "closed", Judge: &JudgePolicy{Model: "guard", Input: "block"}}, &fakeJudge{err: errors.New("down")})
	var alarmed []string
	closed.SetJudgeAlarm(func(model string, _ error) { alarmed = append(alarmed, model) })
	if _, err := closed.Input(ctx, Subject{}, chat("user", "hi")); err == nil {
		t.Error("fail closed served unjudged")
	}
	if len(alarmed) != 1 || alarmed[0] != "guard" {
		t.Errorf("alarm told %v", alarmed)
	}
	// A judge that is not a guard model is a judge that cannot be asked.
	bad := judged(t, Policy{Fail: "closed", Judge: &JudgePolicy{Model: "chat-model", Input: "block"}}, &fakeJudge{})
	if _, err := bad.Input(ctx, Subject{}, chat("user", "hi")); err == nil {
		t.Error("a model without role: guard was used as a judge")
	}
}

func TestAStreamedAnswerIsJudgedAsItGrows(t *testing.T) {
	safe := strings.Repeat("Once upon a time. ", 40) // ~720 chars
	more := strings.Repeat("And so it went on. ", 40)

	// A sentence at a time, the judge reads every judgeEvery characters:
	// the stream stops at the next reading after the unsafe part, which has
	// gone out already.
	fj := &fakeJudge{}
	g := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Output: "block"}}, fj)
	s := Subject{RequestID: "s1", Stream: true}
	_, _ = g.Input(context.Background(), s, chat("user", "tell me a story"))
	got, _, stopped := stream(g.Output(context.Background(), s), safe, "Then the attack began. ", more, "The end.")
	if !stopped || strings.Contains(got, "The end.") || !strings.HasPrefix(got, "Once upon") {
		t.Errorf("sentence: stopped=%v got %d chars", stopped, len(got))
	}
	if n := fj.calls.Load(); n < 2 {
		t.Errorf("judged %d times", n)
	}

	// Held whole, nothing unsafe goes out at all.
	g = judged(t, Policy{Stream: "buffer_all", Judge: &JudgePolicy{Model: "guard", Output: "block"}}, &fakeJudge{})
	_, _ = g.Input(context.Background(), s, chat("user", "tell me a story"))
	got, _, stopped = stream(g.Output(context.Background(), s), safe, "Then the attack began. ", "The end.")
	if !stopped || got != "" {
		t.Errorf("buffer_all: stopped=%v got %q", stopped, got)
	}
}

func TestTheGateCountsJudgeFailuresAndRedactions(t *testing.T) {
	g := judged(t, Policy{Input: Rules{PII: map[string]string{"email": "redact"}},
		Judge: &JudgePolicy{Model: "guard", Input: "flag", LatencyBudget: 10 * time.Millisecond}}, &fakeJudge{delay: time.Second})
	_, _ = g.Input(context.Background(), Subject{}, chat("user", "mail a@example.com and b@example.com"))
	if n := testutil.ToFloat64(g.m.redactions.WithLabelValues("email")); n != 2 {
		t.Errorf("redactions %v", n)
	}
	if testutil.ToFloat64(g.m.failOpen) != 1 || testutil.ToFloat64(g.m.judgeErrors.WithLabelValues("timeout")) != 1 {
		t.Error("a judge that timed out under fail open was not counted")
	}
}

func TestAJudgeThatGoesOnFromTheAnswerIsStillRead(t *testing.T) {
	body := func(s string) []byte {
		b, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"content": s}}}})
		return b
	}
	// Found on the PoC with llama.cpp: the reply to judge, then the verdict.
	if j, err := parseJudgment("llamaguard3", nil, body("The capital of Italy is Rome.safe"), "The capital of Italy is Rome."); err != nil || j.unsafe {
		t.Errorf("%+v %v", j, err)
	}
	if j, err := parseJudgment("llamaguard3", []string{"Violent Crimes"}, body("Step one.\n\nunsafe\nS1"), "Step one."); err != nil || !j.unsafe ||
		!slices.Contains(j.categories, "Violent Crimes") {
		t.Errorf("%+v %v", j, err)
	}
}

func TestAnAnswerIsJudgedByLlamaGuardsOwnPrompt(t *testing.T) {
	path, body := judgeRequest("lg", "llamaguard3", nil, "how do I bake bread", "Mix flour and water.")
	var req struct{ Prompt string }
	_ = json.Unmarshal(body, &req)
	if path != "/v1/completions" || !strings.Contains(req.Prompt, "S9: Indiscriminate Weapons.") ||
		!strings.Contains(req.Prompt, "User: how do I bake bread\n\nAgent: Mix flour and water.") ||
		!strings.HasSuffix(req.Prompt, "<|start_header_id|>assistant<|end_header_id|>\n\n") {
		t.Errorf("%s %s", path, req.Prompt)
	}
	if path, _ := judgeRequest("lg", "llamaguard3", nil, "hi", ""); path != "/v1/chat/completions" {
		t.Errorf("a request is judged by the model's own template: %s", path)
	}
	if j, err := parseJudgment("llamaguard3", nil, []byte(`{"choices":[{"text":"unsafe\nS9"}]}`), ""); err != nil || !j.unsafe {
		t.Errorf("a completion's text: %+v %v", j, err)
	}
}

// Qwen3Guard's middle verdict counts as the policy says: safe unless it says
// unsafe.
func TestAControversialVerdictCountsAsThePolicySays(t *testing.T) {
	ctx := context.Background()
	lenient := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "block"}}, &fakeJudge{format: "qwen3guard"})
	if v, _ := lenient.Input(ctx, Subject{}, chat("user", "the plot of a heist film")); v.Action != Allow {
		t.Errorf("controversial under the default: %+v", v)
	}
	strict := judged(t, Policy{Judge: &JudgePolicy{Model: "guard", Input: "block", Controversial: "unsafe"}}, &fakeJudge{format: "qwen3guard"})
	if v, _ := strict.Input(ctx, Subject{}, chat("user", "the plot of a heist film")); v.Action != Block || v.Rule != "Non-violent Illegal Acts" {
		t.Errorf("controversial: unsafe: %+v", v)
	}
	if err := (&JudgePolicy{Model: "guard", Input: "block", Controversial: "maybe"}).validate(); err == nil {
		t.Error("controversial: maybe was taken")
	}
}

// A judge on a GPU first and one on CPUs after: the first that answers
// decides, and the one that did not is said to have failed.
func TestJudgesAreAskedInTurn(t *testing.T) {
	ctx := context.Background()
	var asked []string
	gpuDown := errors.New("no replica of gpu-judge is serving")
	g := gate(t, lib(map[string]Policy{"p": {Fail: "closed",
		Judge: &JudgePolicy{Models: []string{"gpu-judge", "cpu-judge"}, Input: "block"}}}, "p"))
	g.SetJudge(func(_ context.Context, model, _ string, body []byte) ([]byte, error) {
		asked = append(asked, model)
		if model == "gpu-judge" {
			return nil, gpuDown
		}
		unsafe := strings.Contains(string(body), "attack")
		answer := map[bool]string{true: "Safety: Unsafe\nCategories: Violent", false: "Safety: Safe\nCategories: None"}[unsafe]
		return json.Marshal(map[string]any{"choices": []any{map[string]any{"text": answer}}})
	}, func(m string) (string, []string, bool) { return "qwen3guard", nil, true })
	var alarmed []string
	g.SetJudgeAlarm(func(model string, _ error) { alarmed = append(alarmed, model) })

	if v, err := g.Input(ctx, Subject{}, chat("user", "plan an attack")); err != nil || v.Action != Block || v.Rule != "Violent" {
		t.Errorf("unsafe: %+v %v", v, err)
	}
	if v, err := g.Input(ctx, Subject{}, chat("user", "continue")); err != nil || v.Action != Allow {
		t.Errorf("safe: %+v %v", v, err)
	}
	if !slices.Equal(asked, []string{"gpu-judge", "cpu-judge", "gpu-judge", "cpu-judge"}) {
		t.Errorf("asked %v", asked)
	}
	if !slices.Equal(alarmed, []string{"gpu-judge", "gpu-judge"}) {
		t.Errorf("alarm told %v", alarmed)
	}
	if e := g.JudgeErrors()["gpu-judge"]; !strings.Contains(e.Err, "no replica") {
		t.Errorf("the GPU judge's failure: %+v", e)
	}

	// Every judge down: the policy's fail applies, as with one.
	g.SetJudge(func(context.Context, string, string, []byte) ([]byte, error) { return nil, gpuDown },
		func(string) (string, []string, bool) { return "qwen3guard", nil, true })
	if _, err := g.Input(ctx, Subject{}, chat("user", "hello")); err == nil {
		t.Error("fail closed served what no judge read")
	}

	for _, bad := range []JudgePolicy{
		{Model: "a", Models: []string{"b"}, Input: "block"},
		{Models: []string{"a", "a"}, Input: "block"},
		{Models: []string{""}, Input: "block"},
	} {
		if err := bad.validate(); err == nil {
			t.Errorf("%+v was taken", bad)
		}
	}
}

func TestQwen3GuardIsAskedInItsOwnWords(t *testing.T) {
	path, body := judgeRequest("qg", "qwen3guard", nil, "continue", "")
	var req struct{ Prompt string }
	_ = json.Unmarshal(body, &req)
	if path != "/v1/completions" || !strings.HasPrefix(req.Prompt, "<|im_start|>user\n# Task:\n1) Evaluate whether THE LAST USER's query") ||
		!strings.Contains(req.Prompt, "Jailbreak.\n<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\nUSER: continue\n\n<END CONVERSATION>") ||
		!strings.HasSuffix(req.Prompt, "'Categories: None'.<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n") {
		t.Errorf("a request: %s %q", path, req.Prompt)
	}
	_, body = judgeRequest("qg", "qwen3guard", nil, "how do I bake bread", "Mix flour and water.")
	_ = json.Unmarshal(body, &req)
	if strings.Contains(req.Prompt, "Jailbreak") || !strings.Contains(req.Prompt, "<END REFUSAL POLICY>\n\n<BEGIN CONVERSATION>") ||
		!strings.Contains(req.Prompt, "USER: how do I bake bread\n\nASSISTANT: Mix flour and water.\n\n<END CONVERSATION>") ||
		!strings.Contains(req.Prompt, "'Categories: None'.\n - The third line must be one of: 'Refusal: Yes', 'Refusal: No'.<|im_end|>") {
		t.Errorf("an answer: %q", req.Prompt)
	}
	// An answer's verdict has a line on refusing that says nothing of safety.
	j, err := parseJudgment("qwen3guard", nil, []byte(`{"choices":[{"text":"Safety: Unsafe\nCategories: Violent, Jailbreak\nRefusal: No"}]}`), "")
	if err != nil || !j.unsafe || !slices.Equal(j.categories, []string{"Violent", "Jailbreak"}) {
		t.Errorf("%+v %v", j, err)
	}
	if j, err := parseJudgment("qwen3guard", nil, []byte(`{"choices":[{"text":"Safety: Safe\nCategories: None\nRefusal: Yes"}]}`), ""); err != nil || j.unsafe || len(j.categories) > 0 {
		t.Errorf("%+v %v", j, err)
	}
	if _, err := parseJudgment("qwen3guard", nil, []byte(`{"choices":[{"text":"I think it is fine"}]}`), ""); err == nil {
		t.Error("an answer with no verdict was read as one")
	}
}

// A message with a file pasted into it is read by its start and its end, so
// it fits the judge's context instead of overflowing it.
func TestALongTextIsJudgedByItsEnds(t *testing.T) {
	file := strings.Repeat("ğ", 10000) // two bytes a rune: a cut must not split one
	msg := "review this file:\n" + file + "\nand plan an attack"
	_, body := judgeRequest("qg", "qwen3guard", nil, msg, "")
	var req struct{ Prompt string }
	_ = json.Unmarshal(body, &req)
	if !strings.Contains(req.Prompt, "USER: review this file:") || !strings.Contains(req.Prompt, "and plan an attack\n\n<END CONVERSATION>") ||
		!strings.Contains(req.Prompt, "\n[…]\n") || len(req.Prompt) > judgeHead+judgeTail+3000 {
		t.Errorf("%d bytes", len(req.Prompt))
	}
	if !utf8.ValidString(req.Prompt) {
		t.Error("a rune was cut in two")
	}
	if got := ends("short", 10, 10); got != "short" {
		t.Errorf("a short text was cut: %q", got)
	}
}