diff --git a/BACKLOG.md b/BACKLOG.md index 0832eae..cd042c5 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -1700,7 +1700,9 @@ provider. Never prompt or answer text (`TestTracesSpeakTheGenAIConventionsAndHol Prometheus already counted tokens by direction, cached apart. Left: the MCP calls' spans as `gen_ai.tool.*`, once /mcp is traced at all. -### L2. Redacted personal data put back in the answer (S) +### L2. Redacted personal data put back in the answer (S) — done 2026-10-25 +`restore: true` (`guard.Rules.Restore`), `Verdict.Restore`, `internal/router/restore.go` +ahead of the output filter. The original item follows. The guardrails rewrite personal data as `[KIND_n]` (`internal/guard/gate.go`, `placeholders`), the same value always the same placeholder. A model that answers "Dear [PERSON_1]" leaves it so. Kong's sanitizer restores. diff --git a/CHANGELOG.md b/CHANGELOG.md index 134fe04..c20c27a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,11 @@ APIs change between minors without a deprecation period. ### Added +- **Redacted values put back in the answer.** `restore: true` on a guardrail + policy's input: the model reads `[EMAIL_1]`, the client reads its own + address in the reply, streamed or not. Only the request's own values, held + in memory for its life; the output rules read the answer as the client + will. - **Traces in the OpenTelemetry GenAI conventions.** An inference request and each attempt at an engine carry `gen_ai.operation.name`, the requested and the answering model, `max_tokens`, `temperature` and `top_p` when set, diff --git a/docs/guardrails.md b/docs/guardrails.md index 73f7a59..8225311 100644 --- a/docs/guardrails.md +++ b/docs/guardrails.md @@ -91,6 +91,21 @@ guard: value is the same placeholder throughout a request). - **`block`** stops it. +**Putting them back.** With `restore: true` in a policy's `input`, the +values a request had redacted are put back in its answer: the model reads +"write to [PERSON_1] at [EMAIL_1]", the client reads the name and address +it sent. Streamed answers too — a placeholder split across two events is +held until it is whole. Only that request's own values, only into its +answer; the map is kept in memory for the request and nowhere else. The +output rules then read the answer as the client will, so a policy that +redacts personal data in answers still does. + +```yaml +input: + pii: {email: redact, phone: redact} + restore: true +``` + **Shapes.** - `system_prompt` is put first. diff --git a/docs/reference.md b/docs/reference.md index 20d6f91..ee440c1 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -320,6 +320,7 @@ rather than repeated. | `injection` | `string` | [Rules](#rules) | | `deny` | `[]Pattern` | [Rules](#rules) | | `max_chars` | `int` | [Rules](#rules) | +| `restore` | `bool` | [Rules](#rules) | | `token_file` | `string` | [SCIM](#scim) | | `python` | `string` | [SGLangConfig](#sglangconfig) | | `ttft_p99` | `time.Duration` | [SLO](#slo) | @@ -2769,6 +2770,14 @@ Deny are patterns, matched against the text as a model reads it. MaxChars refuses text longer than this; zero is no limit. +### `restore` + +`bool` + +Restore, on input, puts what redact rewrote as placeholders back into +the answer: the model reads "[EMAIL_1]", the client reads its own +address. Only the values that request sent, only into its answer. + ## SCIM ### `token_file` diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index bc9c554..9f97c16 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -406,6 +406,7 @@ var Reference = []Block{ {Key: "injection", Type: "string", Doc: "Injection is what is done on signs of prompt injection: flag or block.", Nested: ""}, {Key: "deny", Type: "[]Pattern", Doc: "Deny are patterns, matched against the text as a model reads it.", Nested: ""}, {Key: "max_chars", Type: "int", Doc: "MaxChars refuses text longer than this; zero is no limit.", Nested: ""}, + {Key: "restore", Type: "bool", Doc: "Restore, on input, puts what redact rewrote as placeholders back into\nthe answer: the model reads \"[EMAIL_1]\", the client reads its own\naddress. Only the values that request sent, only into its answer.", Nested: ""}, }}, {Package: "config", Name: "SCIM", Title: "SCIM", Doc: "", Fields: []Field{ {Key: "token_file", Type: "string", Doc: "TokenFile holds the bearer token the provider sends. Empty turns\nprovisioning off.", Nested: ""}, diff --git a/internal/guard/gate.go b/internal/guard/gate.go index a2c7959..70c8333 100644 --- a/internal/guard/gate.go +++ b/internal/guard/gate.go @@ -307,6 +307,15 @@ func (p *placeholders) rewrite(text string, spans []detect.Span) string { return b.String() } +// inverse is each placeholder and the value it replaced. +func (p *placeholders) inverse() map[string]string { + out := make(map[string]string, len(p.byValue)) + for v, ph := range p.byValue { + out[ph] = v + } + return out +} + var _ Guard = (*Gate)(nil) // Directives are the policy in force's demands of the gateway. Not in @@ -432,6 +441,9 @@ func (g *Gate) Input(ctx context.Context, s Subject, body []byte) (Verdict, erro rw.seg.set(rw.text) changed = true } + if p.Input.Restore && len(rewrites) > 0 { + v.Restore = ph.inverse() + } if d.pin(p.SystemPrompt, p.DropClientSystem) { changed = true } diff --git a/internal/guard/gate_test.go b/internal/guard/gate_test.go index 4d4f974..17545a8 100644 --- a/internal/guard/gate_test.go +++ b/internal/guard/gate_test.go @@ -309,3 +309,27 @@ func TestAJudgesRefusalNamesTheHazard(t *testing.T) { t.Errorf("%q", got) } } + +func TestRestoreSaysWhichPlaceholderWasWhichOnlyWhenAsked(t *testing.T) { + body := []byte(`{"messages":[{"role":"user","content":"mail ayse@example.com and ayse@example.com"}]}`) + for _, c := range []struct { + mode string + restore bool + want int + }{{"enforce", true, 1}, {"enforce", false, 0}, {"shadow", true, 0}} { + p := Policy{Mode: c.mode, Input: Rules{PII: map[string]string{"email": "redact"}, Restore: c.restore}} + g := gate(t, lib(map[string]Policy{"pii": p}, "pii")) + v, err := g.Input(context.Background(), Subject{}, body) + if err != nil { + t.Fatal(err) + } + if len(v.Restore) != c.want { + t.Errorf("%s restore=%v: %v", c.mode, c.restore, v.Restore) + } + for ph, val := range v.Restore { + if val != "ayse@example.com" || !strings.HasPrefix(ph, "[EMAIL_") { + t.Errorf("%s -> %s", ph, val) + } + } + } +} diff --git a/internal/guard/guard.go b/internal/guard/guard.go index 1016f00..163467b 100644 --- a/internal/guard/guard.go +++ b/internal/guard/guard.go @@ -75,6 +75,10 @@ type Verdict struct { // Shadow says the policy only records: the action is what it would have // done, and nothing is done. Shadow bool + // Restore maps each placeholder the input was given to the value it + // stands for, when the policy puts them back in the answer. Held for + // the request's life only: never stored, logged or traced. + Restore map[string]string } // Stricter reports whether v outranks w: block over redact over flag over diff --git a/internal/guard/policy.go b/internal/guard/policy.go index d762b00..06377dd 100644 --- a/internal/guard/policy.go +++ b/internal/guard/policy.go @@ -86,6 +86,10 @@ type Rules struct { Deny []Pattern `koanf:"deny"` // MaxChars refuses text longer than this; zero is no limit. MaxChars int `koanf:"max_chars"` + // Restore, on input, puts what redact rewrote as placeholders back into + // the answer: the model reads "[EMAIL_1]", the client reads its own + // address. Only the values that request sent, only into its answer. + Restore bool `koanf:"restore"` } // Pattern is a named regular expression and what its match does. @@ -253,6 +257,9 @@ func mergeRules(dst *Rules, r Rules) { } } dst.MaxChars = minSet(dst.MaxChars, r.MaxChars) + // Any policy may ask: what is put back is the caller's own, to the + // caller, and never reaches the model either way. + dst.Restore = dst.Restore || r.Restore } func stricterOf(a, b string, order []string) string { diff --git a/internal/router/restore.go b/internal/router/restore.go new file mode 100644 index 0000000..d508755 --- /dev/null +++ b/internal/router/restore.go @@ -0,0 +1,192 @@ +package router + +import ( + "bytes" + "encoding/json" + "net/http" + "strings" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" +) + +// Putting a request's own redacted values back into its answer: the model +// read "[EMAIL_1]", the client reads the address it sent. A guardrail policy +// asks for it with restore: true on its input rules. +// +// Before the output guardrails, not after: they then judge the answer as the +// client will read it, so a policy that redacts personal data in answers +// still does — and a placeholder the output guardrails make of a value the +// model invented is never mistaken for one of the request's own. + +// restorer swaps placeholders back, in a whole answer or a streamed one. +type restorer struct { + values map[string]string + longest int + // held is, per choice, the tail of the text so far that may be the start + // of a placeholder the next event completes: "[EMA" before "IL_1]". + held map[int]string + // last is the most recent event, the shape a final release is sent in. + last map[string]any +} + +func newRestorer(values map[string]string) *restorer { + r := &restorer{values: values, held: map[int]string{}} + for ph := range values { + r.longest = max(r.longest, len(ph)) + } + return r +} + +// text swaps every whole placeholder in s. +func (r *restorer) text(s string) string { + if !strings.Contains(s, "[") { + return s + } + for ph, v := range r.values { + s = strings.ReplaceAll(s, ph, v) + } + return s +} + +// split is s with its placeholders swapped, less a tail that could still +// become one: what to send now, and what to hold for the next event. +func (r *restorer) split(s string) (send, hold string) { + s = r.text(s) + i := strings.LastIndexByte(s, '[') + if i < 0 || strings.IndexByte(s[i:], ']') >= 0 || len(s)-i >= r.longest { + return s, "" + } + tail := s[i:] + for ph := range r.values { + if strings.HasPrefix(ph, tail) { + return s[:i], tail + } + } + return s, "" +} + +// body swaps the placeholders in a whole JSON answer. A placeholder is only +// ever inside a string there, so each value goes in escaped for one. +func (r *restorer) body(b []byte) []byte { + if !bytes.Contains(b, []byte("[")) { + return b + } + for ph, v := range r.values { + q, _ := json.Marshal(v) + b = bytes.ReplaceAll(b, []byte(ph), q[1:len(q)-1]) + } + return b +} + +// OnEvent swaps the placeholders in one streamed chat-completion event, +// holding back the start of one an event has split. +func (r *restorer) OnEvent(data string) ([]string, bool) { + if !strings.Contains(data, "[") && len(r.held) == 0 { + return []string{data}, false + } + var ev map[string]any + dec := json.NewDecoder(strings.NewReader(data)) + dec.UseNumber() + if err := dec.Decode(&ev); err != nil { + return []string{data}, false + } + choices, _ := ev["choices"].([]any) + for i, c := range choices { + ch, _ := c.(map[string]any) + delta, _ := ch["delta"].(map[string]any) + content, ok := delta["content"].(string) + if !ok { + continue + } + idx := i + if n, ok := ch["index"].(json.Number); ok { + if v, err := n.Int64(); err == nil { + idx = int(v) + } + } + send, hold := r.split(r.held[idx] + content) + // The last event of a choice releases everything. + if ch["finish_reason"] != nil { + send, hold = send+hold, "" + } + delta["content"] = send + if hold == "" { + delete(r.held, idx) + } else { + r.held[idx] = hold + } + } + r.last = ev + out, err := json.Marshal(ev) + if err != nil { + return []string{data}, false + } + return []string{string(out)}, false +} + +// OnEnd releases what is still held, as one more event per choice. +func (r *restorer) OnEnd() []string { + if len(r.held) == 0 || r.last == nil { + return nil + } + var out []string + for idx, text := range r.held { + ev := map[string]any{} + for k, v := range r.last { + if k != "choices" && k != "usage" { + ev[k] = v + } + } + ev["choices"] = []any{map[string]any{"index": idx, "delta": map[string]any{"content": text}, "finish_reason": nil}} + if b, err := json.Marshal(ev); err == nil { + out = append(out, string(b)) + } + } + r.held = map[int]string{} + return out +} + +// streamChain runs one stream filter's output through another. +type streamChain struct{ first, then proxy.StreamFilter } + +func (c streamChain) OnEvent(data string) ([]string, bool) { + first, stop := c.first.OnEvent(data) + var out []string + for _, d := range first { + e, s := c.then.OnEvent(d) + out = append(out, e...) + stop = stop || s + } + return out, stop +} + +func (c streamChain) OnEnd() []string { + var out []string + for _, d := range c.first.OnEnd() { + e, _ := c.then.OnEvent(d) + out = append(out, e...) + } + return append(out, c.then.OnEnd()...) +} + +// restoreInto puts the input's placeholders back in this attempt's answer, +// ahead of whatever output filter the attempt already has. +func restoreInto(values map[string]string, stream bool, preq *proxy.Request) { + if len(values) == 0 { + return + } + rs := newRestorer(values) + if f := preq.Filter; f != nil { + preq.Filter = func(b []byte, h http.Header) []byte { return f(rs.body(b), h) } + } else { + preq.Filter = func(b []byte, _ http.Header) []byte { return rs.body(b) } + } + if !stream { + return + } + if sf := preq.StreamFilter; sf != nil { + preq.StreamFilter = streamChain{first: rs, then: sf} + } else { + preq.StreamFilter = rs + } +} diff --git a/internal/router/restore_test.go b/internal/router/restore_test.go new file mode 100644 index 0000000..6d3cc49 --- /dev/null +++ b/internal/router/restore_test.go @@ -0,0 +1,110 @@ +package router + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/guard" +) + +// restoreGuard redacts alice's address as the gate does, and asks for it +// back in the answer. +type restoreGuard struct{} + +func (restoreGuard) Input(_ context.Context, _ guard.Subject, body []byte) (guard.Verdict, error) { + if !bytes.Contains(body, []byte("alice@example.com")) { + return guard.Verdict{}, nil + } + return guard.Verdict{Action: guard.Redact, Policy: "pii", Category: "email", + Body: bytes.ReplaceAll(body, []byte("alice@example.com"), []byte("[EMAIL_1]")), + Restore: map[string]string{"[EMAIL_1]": `alice@example.com`}}, nil +} +func (restoreGuard) Output(context.Context, guard.Subject) guard.OutputFilter { return nil } +func (restoreGuard) Directives(guard.Subject) guard.Directives { return guard.Directives{} } + +// echoEngine answers with the placeholder, in one piece or, streamed, split +// across two events. +func echoEngine(t *testing.T) *httptest.Server { + mux := http.NewServeMux() + mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, "vllm:num_requests_running 0\nvllm:num_requests_waiting 0\nvllm:kv_cache_usage_perc 0.1\n") + }) + mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + b := new(bytes.Buffer) + _, _ = b.ReadFrom(r.Body) + if bytes.Contains(b.Bytes(), []byte("alice@example.com")) { + t.Errorf("the engine saw the address: %s", b) + } + if !bytes.Contains(b.Bytes(), []byte(`"stream":true`)) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"Sent to [EMAIL_1]."},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":4}}`) + return + } + w.Header().Set("Content-Type", "text/event-stream") + for _, piece := range []string{"Sent to [EMA", "IL_1]", "."} { + fmt.Fprintf(w, "data: {\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":%q},\"finish_reason\":null}]}\n\n", piece) + } + fmt.Fprint(w, "data: {\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +func TestTheAnswerGetsBackTheValuesItsRequestHadRedacted(t *testing.T) { + srv := echoEngine(t) + rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}})) + rt.SetGuard(restoreGuard{}) + + rec := post(t, rt.Handler(), say("write to alice@example.com"), nil) + if !strings.Contains(rec.Body.String(), "Sent to alice@example.com.") || strings.Contains(rec.Body.String(), "[EMAIL_1]") { + t.Errorf("one-shot answer: %s", rec.Body.String()) + } + + streamed := say("write to alice@example.com") + streamed["stream"] = true + rec = post(t, rt.Handler(), streamed, nil) + var text strings.Builder + for _, line := range strings.Split(rec.Body.String(), "\n") { + if i := strings.Index(line, `"content":"`); i >= 0 { + rest := line[i+len(`"content":"`):] + text.WriteString(rest[:strings.Index(rest, `"`)]) + } + } + if text.String() != "Sent to alice@example.com." { + t.Errorf("streamed answer reads %q\n%s", text.String(), rec.Body.String()) + } +} + +func TestARestorerHoldsOnlyWhatCouldBeAPlaceholder(t *testing.T) { + r := newRestorer(map[string]string{"[EMAIL_1]": "a@b.c", "[PHONE_1]": "+46 70"}) + for _, c := range []struct{ in, send, hold string }{ + {"Dear [EMA", "Dear ", "[EMA"}, + {"a [1] b", "a [1] b", ""}, + {"x [EMAIL_1] y", "x a@b.c y", ""}, + {"end [", "end ", "["}, + {"[X", "[X", ""}, + } { + send, hold := r.split(c.in) + if send != c.send || hold != c.hold { + t.Errorf("split(%q) = %q, %q; want %q, %q", c.in, send, hold, c.send, c.hold) + } + } + // A value with a quote stays valid JSON in a whole answer. + q := newRestorer(map[string]string{"[PERSON_1]": `O"Brien`}) + if got := string(q.body([]byte(`{"c":"hi [PERSON_1]"}`))); got != `{"c":"hi O\"Brien"}` { + t.Errorf("body: %s", got) + } + // What is held at the end is still sent. + e := newRestorer(map[string]string{"[EMAIL_1]": "a@b.c"}) + e.OnEvent(`{"choices":[{"index":0,"delta":{"content":"to [EM"}}]}`) + if out := e.OnEnd(); len(out) != 1 || !strings.Contains(out[0], `"content":"[EM"`) { + t.Errorf("OnEnd: %v", out) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index b9a1600..92ea837 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -1249,6 +1249,7 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa } // A filter per attempt, for the model this attempt is served by. filter := r.outputFilter(uctx, req, model, &preq) + restoreInto(req.guardIn.Restore, req.Stream, &preq) var out proxy.Outcome var perr error if dec.Endpoint.Engine == endpoint.EngineComfyUI {