diff --git a/internal/guard/judge.go b/internal/guard/judge.go index 290d507..0096c0d 100644 --- a/internal/guard/judge.go +++ b/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" + + "\n") + for i, c := range categories { + fmt.Fprintf(&b, "S%d: %s.\n", i+1, c) + } + b.WriteString("\n\n\n\n") + if user != "" { + b.WriteString("User: " + user + "\n\n") + } + b.WriteString("Agent: " + answer + "\n\n\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)) } diff --git a/internal/guard/judge_test.go b/internal/guard/judge_test.go index 2b38125..6b36526 100644 --- a/internal/guard/judge_test.go +++ b/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) + } +} diff --git a/internal/router/guard_test.go b/internal/router/guard_test.go index 70829db..e75c238 100644 --- a/internal/router/guard_test.go +++ b/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) } diff --git a/internal/router/judge.go b/internal/router/judge.go index 39e709b..c16376f 100644 --- a/internal/router/judge.go +++ b/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 }