diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index f386bd1a5fef..ec1e360424d8 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -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": { diff --git a/core/config/model_config.go b/core/config/model_config.go index 21bc42facc63..1b49a2024ad5 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -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 diff --git a/core/services/routing/pii/middleware.go b/core/services/routing/pii/middleware.go index 94a446fbc5c9..aeb3868687ea 100644 --- a/core/services/routing/pii/middleware.go +++ b/core/services/routing/pii/middleware.go @@ -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 @@ -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 { @@ -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 } } } diff --git a/core/services/routing/pii/middleware_test.go b/core/services/routing/pii/middleware_test.go index 6b8416876d06..e87e586c0d64 100644 --- a/core/services/routing/pii/middleware_test.go +++ b/core/services/routing/pii/middleware_test.go @@ -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 { @@ -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"}} diff --git a/core/services/routing/pii/pseudonymizer.go b/core/services/routing/pii/pseudonymizer.go new file mode 100644 index 000000000000..e09b5cf0837e --- /dev/null +++ b/core/services/routing/pii/pseudonymizer.go @@ -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:] +} diff --git a/docs/content/operations/middleware.md b/docs/content/operations/middleware.md index 46a528528711..245e2a52d32d 100644 --- a/docs/content/operations/middleware.md +++ b/docs/content/operations/middleware.md @@ -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