Judge an answer with Llama Guard 3's own prompt, as a completion, not a conversation ending on the answer

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

internal/guard/judge.go

@@ -80,9 +80,10 @@ func mergeJudge(a, b *JudgePolicy) *JudgePolicy {
 	return &out
 }
 
-// Caller sends a chat completion body to a model on the cluster and returns
-// the answer's body: the router, past admission, accounting and guardrails.
-type Caller func(ctx context.Context, model string, body []byte) ([]byte, error)
+// 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.
@@ -154,8 +155,8 @@ func (j *judge) ask(ctx context.Context, jp *JudgePolicy, user, answer string) (
 
 	ctx, cancel := context.WithTimeout(ctx, jp.budget())
 	defer cancel()
-	body := judgeRequest(jp.Model, format, names, user, answer)
-	out, err := j.call(ctx, jp.Model, body)
+	path, body := judgeRequest(jp.Model, format, names, user, answer)
+	out, err := j.call(ctx, jp.Model, path, body)
 	if err != nil {
 		return judgment{}, err
 	}
@@ -192,8 +193,48 @@ func (jp *JudgePolicy) counts(res judgment) (bool, string) {
 	return false, ""
 }
 
-// judgeRequest is the chat completion a judge model is sent.
-func judgeRequest(model, format string, categories []string, user, answer string) []byte {
+// 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":
@@ -227,7 +268,7 @@ func judgeRequest(model, format string, categories []string, user, answer string
 			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 b
+	return "/v1/chat/completions", b
 }
 
 // parseJudgment reads a judge model's answer.
@@ -241,12 +282,13 @@ func parseJudgment(format string, names []string, body []byte, answer string) (j
 			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 chat completion", errJudgeFormat)
+		return judgment{}, fmt.Errorf("%w: not a completion", errJudgeFormat)
 	}
-	text := strings.TrimSpace(resp.Choices[0].Message.Content)
+	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))
 	}

internal/guard/judge_test.go

@@ -22,7 +22,7 @@ type fakeJudge struct {
 	delay  time.Duration
 }
 
-func (f *fakeJudge) call(ctx context.Context, _ string, body []byte) ([]byte, error) {
+func (f *fakeJudge) call(ctx context.Context, _, _ string, body []byte) ([]byte, error) {
 	f.calls.Add(1)
 	if f.delay > 0 {
 		select {
@@ -169,3 +169,20 @@ func TestAJudgeThatGoesOnFromTheAnswerIsStillRead(t *testing.T) {
 		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)
+	}
+}

internal/router/guard_test.go

@@ -221,7 +221,7 @@ func TestAJudgeModelIsNobodysToAskButTheGuardrails(t *testing.T) {
 		t.Errorf("listed: %s", rec.Body.String())
 	}
 	calls := e.calls.Load()
-	out, err := rt.Judge(t.Context(), "llama-guard", []byte(`{"model":"llama-guard","messages":[{"role":"user","content":"x"}]}`))
+	out, err := rt.Judge(t.Context(), "llama-guard", "/v1/chat/completions", []byte(`{"model":"llama-guard","messages":[{"role":"user","content":"x"}]}`))
 	if err != nil || !strings.Contains(string(out), `"hi"`) || e.calls.Load() != calls+1 {
 		t.Errorf("the guardrails asking: %s %v", out, err)
 	}

internal/router/judge.go

@@ -15,12 +15,15 @@ import (
 // judge. It goes past admission — a judge call queued behind the tenants it
 // is judging would wait on them — and past accounting and the guardrails
 // themselves, but through the scheduler and the replicas like any request.
-func (r *Router) Judge(ctx context.Context, model string, body []byte) ([]byte, error) {
-	hreq, err := http.NewRequestWithContext(ctx, http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
+func (r *Router) Judge(ctx context.Context, model, path string, body []byte) ([]byte, error) {
+	if path != "/v1/chat/completions" && path != "/v1/completions" {
+		return nil, fmt.Errorf("a judge is asked at /v1/chat/completions or /v1/completions, not %s", path)
+	}
+	hreq, err := http.NewRequestWithContext(ctx, http.MethodPost, path, bytes.NewReader(body))
 	if err != nil {
 		return nil, err
 	}
-	req, err := r.parse(hreq, "/v1/chat/completions")
+	req, err := r.parse(hreq, path)
 	if err != nil {
 		return nil, err
 	}