llmgw/internal/guard/judge_test.go

package guard

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

	"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]
	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 []string{"llamaguard3", "shieldgemma", "json"} {
		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)
	}
}