Compare commits

...

2 Commits

Author SHA1 Message Date
localai-org-maint-bot
625ca71464 fix(pii): wrap reversible redaction tokens
Use configurable token delimiters to avoid restoring ordinary model text that happens to match an internal identifier. Rename the option and document the confidentiality tradeoff.

Assisted-by: Codex:gpt-5
2026-08-02 03:05:22 +00:00
localai-org-maint-bot
2e6c63daec feat(pii): restore request-scoped pseudonyms
Replace masked request values with unique per-request tokens when response restoration is enabled, then restore them across JSON and SSE write boundaries. Document the opt-in model setting and expose it in config metadata.\n\nAssisted-by: Codex:gpt-5
2026-08-02 03:05:22 +00:00
6 changed files with 276 additions and 6 deletions

View File

@@ -816,6 +816,25 @@ func DefaultRegistry() map[string]FieldMetaOverride {
AutocompleteProvider: "models:token_classify",
Order: 201,
},
"pii.reversible_redactions": {
Section: "pii",
Label: "Reversible Redactions",
Description: "Replace masked values with wrapped request-scoped tokens and restore them when the model returns those tokens. Supports streaming responses and never persists the substitution map.",
Component: "toggle",
Order: 202,
},
"pii.reversible_token_prefix": {
Section: "pii",
Label: "Reversible Token Prefix",
Description: "Prefix for reversible redaction tokens. Defaults to [REDACTED:.",
Order: 203,
},
"pii.reversible_token_suffix": {
Section: "pii",
Label: "Reversible Token Suffix",
Description: "Suffix for reversible redaction tokens. Defaults to ].",
Order: 204,
},
// --- PII detection policy (on a token_classify detector model) ---
"pii_detection.min_score": {

View File

@@ -440,8 +440,18 @@ type PIIConfig struct {
// model just opts in by listing detectors. Multiple detectors union
// their hits; overlapping spans resolve to the strongest action.
Detectors []string `yaml:"detectors,omitempty" json:"detectors,omitempty"`
// ReversibleRedactions replaces request PII with stable, request-scoped
// tokens and restores those values when the wrapped tokens appear in the response.
ReversibleRedactions bool `yaml:"reversible_redactions,omitempty" json:"reversible_redactions,omitempty"`
ReversibleTokenPrefix string `yaml:"reversible_token_prefix,omitempty" json:"reversible_token_prefix,omitempty"`
ReversibleTokenSuffix string `yaml:"reversible_token_suffix,omitempty" json:"reversible_token_suffix,omitempty"`
}
func (c ModelConfig) PIIReversibleRedactions() bool { return c.PII.ReversibleRedactions }
func (c ModelConfig) PIIReversibleTokenPrefix() string { return c.PII.ReversibleTokenPrefix }
func (c ModelConfig) PIIReversibleTokenSuffix() string { return c.PII.ReversibleTokenSuffix }
// @Description Detection policy for a token-classification (NER) model
// used as a PII detector. Lives on the detector model's own config so the
// model is a self-describing policy unit: consuming models reference it by

View File

@@ -194,6 +194,16 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
texts := adapter.Scan(parsed)
updates := make([]ScannedText, 0, len(texts))
prefix, suffix := defaultReversibleTokenPrefix, defaultReversibleTokenSuffix
if cfg, ok := rawCfg.(responsePIIConfig); ok {
if cfg.PIIReversibleTokenPrefix() != "" {
prefix = cfg.PIIReversibleTokenPrefix()
}
if cfg.PIIReversibleTokenSuffix() != "" {
suffix = cfg.PIIReversibleTokenSuffix()
}
}
pseudonyms := newPseudonymizer(prefix, suffix)
var blocked bool
var firstEventID string
@@ -259,7 +269,11 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
if res.Blocked {
blocked = true
}
updates = append(updates, ScannedText{Index: st.Index, Text: res.Redacted})
redacted := res.Redacted
if cfg, ok := rawCfg.(responsePIIConfig); ok && cfg.PIIReversibleRedactions() {
redacted = pseudonyms.replace(st.Text, res.Spans)
}
updates = append(updates, ScannedText{Index: st.Index, Text: redacted})
}
if blocked {
@@ -279,7 +293,16 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
if firstEventID != "" {
c.Set(ctxKeyPIIEventID, firstEventID)
}
return next(c)
if len(pseudonyms.original) == 0 {
return next(c)
}
writer := newRestoringWriter(c.Response().Writer, pseudonyms.original)
c.Response().Writer = writer
err := next(c)
if finishErr := writer.Finish(); err == nil {
err = finishErr
}
return err
}
}
}

View File

@@ -60,10 +60,16 @@ func setRequestOnContext(req *fakeRequest) echo.MiddlewareFunc {
type fakeModelPIIConfig struct {
enabled bool
detectors []string
reverse bool
prefix string
suffix string
}
func (f fakeModelPIIConfig) PIIIsEnabled() bool { return f.enabled }
func (f fakeModelPIIConfig) PIIDetectors() []string { return f.detectors }
func (f fakeModelPIIConfig) PIIIsEnabled() bool { return f.enabled }
func (f fakeModelPIIConfig) PIIDetectors() []string { return f.detectors }
func (f fakeModelPIIConfig) PIIReversibleRedactions() bool { return f.reverse }
func (f fakeModelPIIConfig) PIIReversibleTokenPrefix() string { return f.prefix }
func (f fakeModelPIIConfig) PIIReversibleTokenSuffix() string { return f.suffix }
func withModelConfig(cfg fakeModelPIIConfig) echo.MiddlewareFunc {
return func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -129,6 +135,58 @@ var _ = Describe("RequestMiddleware (NER)", func() {
Expect(events[0].Direction).To(Equal(DirectionIn))
})
It("restores distinct pseudonyms across streaming write boundaries", func() {
body := &fakeRequest{Messages: []string{"Email alice@example.com or bob@example.com"}}
mw := RequestMiddleware(&Redactor{}, store(), fakeAdapter(), nil,
WithNERResolver(resolverFor(map[string]NERConfig{
"privacy-filter": nerCfg(ActionMask,
NEREntity{Group: "EMAIL", Start: 6, End: 23, Score: 0.95},
NEREntity{Group: "EMAIL", Start: 27, End: 42, Score: 0.95}),
})))
e := echo.New()
e.POST("/chat", func(c echo.Context) error {
Expect(body.Messages[0]).To(Equal("Email [REDACTED:EMAIL_001] or [REDACTED:EMAIL_002]"))
_, err := c.Response().Write([]byte(`data: {"delta":"EMAIL_001 and [REDACTED:EMAIL_0`))
Expect(err).ToNot(HaveOccurred())
_, err = c.Response().Write([]byte(`01] and [REDACTED:EMAIL_002]"}` + "\n\n"))
return err
}, setRequestOnContext(body), withModelConfig(fakeModelPIIConfig{
enabled: true, detectors: []string{"privacy-filter"}, reverse: true,
}), mw)
req := httptest.NewRequest(http.MethodPost, "/chat", strings.NewReader(`{}`))
w := httptest.NewRecorder()
e.ServeHTTP(w, req)
Expect(w.Code).To(Equal(http.StatusOK))
Expect(w.Body.String()).To(Equal("data: {\"delta\":\"EMAIL_001 and alice@example.com and bob@example.com\"}\n\n"))
})
It("uses configured reversible redaction token delimiters", func() {
body := &fakeRequest{Messages: []string{"Email alice@example.com"}}
mw := RequestMiddleware(&Redactor{}, store(), fakeAdapter(), nil,
WithNERResolver(resolverFor(map[string]NERConfig{
"privacy-filter": nerCfg(ActionMask,
NEREntity{Group: "EMAIL", Start: 6, End: 23, Score: 0.95}),
})))
e := echo.New()
e.POST("/chat", func(c echo.Context) error {
Expect(body.Messages[0]).To(Equal("Email <PII:EMAIL_001>"))
_, err := c.Response().Write([]byte(`{"text":"<PII:EMAIL_001>"}`))
return err
}, setRequestOnContext(body), withModelConfig(fakeModelPIIConfig{
enabled: true, detectors: []string{"privacy-filter"}, reverse: true,
prefix: "<PII:", suffix: ">",
}), mw)
req := httptest.NewRequest(http.MethodPost, "/chat", strings.NewReader(`{}`))
w := httptest.NewRecorder()
e.ServeHTTP(w, req)
Expect(w.Code).To(Equal(http.StatusOK))
Expect(w.Body.String()).To(Equal(`{"text":"alice@example.com"}`))
})
It("blocks (400) when a detected entity's action is block", func() {
st := store()
body := &fakeRequest{Messages: []string{"my password is hunter2 ok"}}

View File

@@ -0,0 +1,142 @@
package pii
import (
"encoding/json"
"fmt"
"net/http"
"strings"
"unicode"
)
type responsePIIConfig interface {
PIIReversibleRedactions() bool
PIIReversibleTokenPrefix() string
PIIReversibleTokenSuffix() string
}
const (
defaultReversibleTokenPrefix = "[REDACTED:"
defaultReversibleTokenSuffix = "]"
)
type pseudonymizer struct {
byValue map[string]string
original map[string]string
counts map[string]int
prefix string
suffix string
}
func newPseudonymizer(prefix, suffix string) *pseudonymizer {
return &pseudonymizer{
byValue: map[string]string{},
original: map[string]string{},
counts: map[string]int{},
prefix: prefix,
suffix: suffix,
}
}
func (p *pseudonymizer) replace(text string, spans []Span) string {
var b strings.Builder
last := 0
for _, span := range spans {
if span.Action != ActionMask || span.Start < last || span.End > len(text) {
continue
}
b.WriteString(text[last:span.Start])
value := text[span.Start:span.End]
token, ok := p.byValue[value]
if !ok {
group := pseudonymGroup(span.Pattern)
p.counts[group]++
token = fmt.Sprintf("%s%s_%03d%s", p.prefix, group, p.counts[group], p.suffix)
p.byValue[value] = token
p.original[token] = value
}
b.WriteString(token)
last = span.End
}
b.WriteString(text[last:])
return b.String()
}
func pseudonymGroup(pattern string) string {
if i := strings.LastIndexByte(pattern, ':'); i >= 0 {
pattern = pattern[i+1:]
}
var b strings.Builder
for _, r := range strings.ToUpper(pattern) {
if unicode.IsLetter(r) || unicode.IsDigit(r) {
b.WriteRune(r)
} else {
b.WriteByte('_')
}
}
if b.Len() == 0 {
return "PII"
}
return b.String()
}
type restoringWriter struct {
http.ResponseWriter
pending string
replacements map[string]string
}
func newRestoringWriter(w http.ResponseWriter, originals map[string]string) *restoringWriter {
replacements := make(map[string]string, len(originals))
for token, original := range originals {
encoded, _ := json.Marshal(original)
replacements[token] = string(encoded[1 : len(encoded)-1])
}
return &restoringWriter{ResponseWriter: w, replacements: replacements}
}
func (w *restoringWriter) Write(data []byte) (int, error) {
w.pending += string(data)
ready, pending := w.splitReady(w.replace(w.pending))
w.pending = pending
if ready != "" {
if _, err := w.ResponseWriter.Write([]byte(ready)); err != nil {
return 0, err
}
}
return len(data), nil
}
func (w *restoringWriter) Flush() {
if f, ok := w.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (w *restoringWriter) Finish() error {
if w.pending == "" {
return nil
}
_, err := w.ResponseWriter.Write([]byte(w.replace(w.pending)))
w.pending = ""
return err
}
func (w *restoringWriter) replace(s string) string {
for token, original := range w.replacements {
s = strings.ReplaceAll(s, token, original)
}
return s
}
func (w *restoringWriter) splitReady(s string) (string, string) {
keep := 0
for token := range w.replacements {
limit := min(len(token)-1, len(s))
for n := 1; n <= limit; n++ {
if strings.HasSuffix(s, token[:n]) && n > keep {
keep = n
}
}
}
return s[:len(s)-keep], s[len(s)-keep:]
}

View File

@@ -165,15 +165,33 @@ pii:
enabled: true # default-on for cloud-proxy; explicit for audit
detectors:
- privacy-filter-multilingual
reversible_redactions: true # restore request PII if the model echoes its wrapped token
reversible_token_prefix: "[REDACTED:" # optional; this is the default
reversible_token_suffix: "]" # optional; this is the default
```
`reversible_redactions` enables bijective, request-scoped replacement. Each
masked value is sent to the backend as a stable wrapped token such as
`[REDACTED:EMAIL_001]` instead of a generic redaction marker. If the model includes that token in
its response, LocalAI restores the original value before returning JSON or SSE
to the caller. The substitution map exists only for that request and is never
logged or persisted. Leave the option unset (the default) for irreversible
`[REDACTED:...]` masking.
The prefix and suffix reduce collisions with ordinary model output and can be
customized with `reversible_token_prefix` and `reversible_token_suffix`.
Reversible redactions provide less confidentiality than irreversible masking:
any third party that can observe both the redacted request and restored response
may be able to infer the original values.
Multiple detectors **union** their detections; overlapping spans resolve to
the strongest action (`block` > `mask` > `allow`). A configured detector
that can't be loaded **fails the request closed** (HTTP 503,
`error.type=pii_ner_unavailable`) rather than silently skipping the check.
The same NER path runs on the [MITM proxy]({{< relref "mitm-proxy.md" >}})
request body for intercepted hosts. Response/output redaction is out of
scope for now.
request body for intercepted hosts. Reversible response restoration currently
applies to LocalAI API routes; the MITM proxy keeps its own output-redaction
policy.
### Instance-wide default detector