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
7 changes: 7 additions & 0 deletions core/config/meta/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -816,6 +816,13 @@ func DefaultRegistry() map[string]FieldMetaOverride {
AutocompleteProvider: "models:token_classify",
Order: 201,
},
"pii.reverse_in_response": {
Section: "pii",
Label: "Restore PII In Response",
Description: "Replace masked values with request-scoped pseudonyms and restore them when the model returns those pseudonyms. Supports streaming responses and never persists the substitution map.",
Component: "toggle",
Order: 202,
},

// --- PII detection policy (on a token_classify detector model) ---
"pii_detection.min_score": {
Expand Down
6 changes: 6 additions & 0 deletions core/config/model_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -440,8 +440,14 @@ 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"`

// ReverseInResponse replaces request PII with stable, request-scoped
// pseudonyms and restores those values when they appear in the response.
ReverseInResponse bool `yaml:"reverse_in_response,omitempty" json:"reverse_in_response,omitempty"`
}

func (c ModelConfig) PIIReverseInResponse() bool { return c.PII.ReverseInResponse }

// @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
18 changes: 16 additions & 2 deletions core/services/routing/pii/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,7 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa

texts := adapter.Scan(parsed)
updates := make([]ScannedText, 0, len(texts))
pseudonyms := newPseudonymizer()
var blocked bool
var firstEventID string

Expand Down Expand Up @@ -259,7 +260,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.PIIReverseInResponse() {
redacted = pseudonyms.replace(st.Text, res.Spans)
}
updates = append(updates, ScannedText{Index: st.Index, Text: redacted})
}

if blocked {
Expand All @@ -279,7 +284,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
33 changes: 31 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,12 @@ func setRequestOnContext(req *fakeRequest) echo.MiddlewareFunc {
type fakeModelPIIConfig struct {
enabled bool
detectors []string
reverse bool
}

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) PIIReverseInResponse() bool { return f.reverse }

func withModelConfig(cfg fakeModelPIIConfig) echo.MiddlewareFunc {
return func(next echo.HandlerFunc) echo.HandlerFunc {
Expand Down Expand Up @@ -129,6 +131,33 @@ 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 EMAIL_001 or EMAIL_002"))
_, err := c.Response().Write([]byte(`data: {"delta":"EMAIL_0`))
Expect(err).ToNot(HaveOccurred())
_, err = c.Response().Write([]byte(`01 and 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\":\"alice@example.com and bob@example.com\"}\n\n"))
})

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
131 changes: 131 additions & 0 deletions core/services/routing/pii/pseudonymizer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,131 @@
package pii

import (
"encoding/json"
"fmt"
"net/http"
"strings"
"unicode"
)

type responsePIIConfig interface {
PIIReverseInResponse() bool
}

type pseudonymizer struct {
byValue map[string]string
original map[string]string
counts map[string]int
}

func newPseudonymizer() *pseudonymizer {
return &pseudonymizer{
byValue: map[string]string{},
original: map[string]string{},
counts: map[string]int{},
}
}

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_%03d", group, p.counts[group])
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:]
}
14 changes: 12 additions & 2 deletions docs/content/operations/middleware.md
Original file line number Diff line number Diff line change
Expand Up @@ -165,15 +165,25 @@ pii:
enabled: true # default-on for cloud-proxy; explicit for audit
detectors:
- privacy-filter-multilingual
reverse_in_response: true # restore request PII if the model echoes its pseudonym
```

`reverse_in_response` enables bijective, request-scoped replacement. Each
masked value is sent to the backend as a stable pseudonym such as `EMAIL_001`
instead of a generic redaction marker. If the model includes that pseudonym 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.

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. Bijective response restoration currently
applies to LocalAI API routes; the MITM proxy keeps its own output-redaction
policy.

### Instance-wide default detector

Expand Down
Loading