Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions core/config/meta/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down
10 changes: 10 additions & 0 deletions core/config/model_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 25 additions & 2 deletions core/services/routing/pii/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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 {
Expand All @@ -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
}
}
}
Expand Down
62 changes: 60 additions & 2 deletions core/services/routing/pii/middleware_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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"}}
Expand Down
142 changes: 142 additions & 0 deletions core/services/routing/pii/pseudonymizer.go
Original file line number Diff line number Diff line change
@@ -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:]
}
22 changes: 20 additions & 2 deletions docs/content/operations/middleware.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading