feat: add PII redaction plugin crate with local backend

Cargo.toml

@@ -5,6 +5,7 @@
 members = [
     "crates/core",
     "crates/adaptive",
+    "crates/pii-redaction",
     "crates/cli",
     # Language Bindings
     "crates/python",
@@ -24,6 +25,7 @@ repository = "https://github.com/NVIDIA/NeMo-Relay"
 [workspace.dependencies]
 nemo-relay = { version = "0.4.0", path = "crates/core", default-features = false }
 nemo-relay-adaptive = { version = "0.4.0", path = "crates/adaptive" }
+nemo-relay-pii-redaction = { version = "0.4.0", path = "crates/pii-redaction" }
 nemo-relay-ffi = { version = "0.4.0", path = "crates/ffi" }
 nemo-relay-cli = { version = "0.4.0", path = "crates/cli" }
 opentelemetry = { version = "0.31", default-features = false }

crates/cli/Cargo.toml

@@ -27,6 +27,7 @@ atof-streaming = ["nemo-relay/atof-streaming"]
 [dependencies]
 nemo-relay = { workspace = true, features = ["guardrails-remote", "object-store", "openinference"] }
 nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] }
+nemo-relay-pii-redaction.workspace = true
 async-stream = "0.3"
 axum = "0.8"
 bytes = "1"

crates/cli/src/doctor.rs

@@ -19,6 +19,7 @@ use nemo_relay::codec::pricing::{PricingCatalog, PricingConfig, PricingSourceCon
 use nemo_relay::observability::plugin_component::OBSERVABILITY_PLUGIN_KIND;
 use nemo_relay::plugin::{DiagnosticLevel, PluginConfig, validate_plugin_config};
 use nemo_relay_adaptive::plugin_component::register_adaptive_component;
+use nemo_relay_pii_redaction::component::register_pii_redaction_component;
 use serde::Serialize;
 use serde_json::{Value, json};
 use tokio::time::timeout;
@@ -607,6 +608,14 @@ async fn collect_observability(gateway: &GatewayConfig) -> Vec<Check> {
         });
         return checks;
     }
+    if let Err(error) = register_pii_redaction_component() {
+        checks.push(Check {
+            name: "PII redaction plugin",
+            status: Status::Fail,
+            details: format!("registration failed: {error}"),
+        });
+        return checks;
+    }
     let report = validate_plugin_config(&plugin_config);
     if report.diagnostics.is_empty() {
         checks.push(Check {

crates/cli/src/plugins.rs

@@ -280,6 +280,10 @@ fn edit_component_field(
             edit_config_field(theme, &mut state.config, field)?;
             state.mark_config_touched();
         }
+        EditableComponent::PiiRedaction(state) => {
+            edit_config_field(theme, &mut state.config, field)?;
+            state.mark_config_touched();
+        }
     }
     Ok(())
 }

crates/cli/src/plugins/config_io.rs

@@ -8,6 +8,7 @@ use std::path::{Path, PathBuf};
 use console::style;
 use nemo_relay::plugin::{ConfigPolicy, PluginConfig, validate_plugin_config};
 use nemo_relay_adaptive::plugin_component::register_adaptive_component;
+use nemo_relay_pii_redaction::component::register_pii_redaction_component;
 use serde_json::{Map, Value};
 
 use crate::config::{
@@ -119,6 +120,9 @@ pub(crate) fn validate_config(config: &PluginConfig) -> Result<(), CliError> {
     register_adaptive_component().map_err(|error| {
         CliError::Config(format!("adaptive plugin registration failed: {error}"))
     })?;
+    register_pii_redaction_component().map_err(|error| {
+        CliError::Config(format!("PII redaction plugin registration failed: {error}"))
+    })?;
     let report = validate_plugin_config(config);
     if report.has_errors() {
         let messages = report

crates/cli/src/plugins/editor_model.rs

@@ -13,6 +13,7 @@ use nemo_relay::plugins::nemo_guardrails::component::{
 };
 use nemo_relay_adaptive::AdaptiveConfig;
 use nemo_relay_adaptive::plugin_component::ADAPTIVE_PLUGIN_KIND;
+use nemo_relay_pii_redaction::component::{PII_REDACTION_PLUGIN_KIND, PiiRedactionConfig};
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use serde_json::{Map, Value, json};
@@ -36,6 +37,7 @@ pub(super) enum EditableComponent {
     Observability(Box<ComponentEditorState<ObservabilityConfig>>),
     Adaptive(Box<ComponentEditorState<AdaptiveConfig>>),
     NemoGuardrails(Box<ComponentEditorState<NeMoGuardrailsConfig>>),
+    PiiRedaction(Box<ComponentEditorState<PiiRedactionConfig>>),
 }
 
 impl EditableComponent {
@@ -44,6 +46,7 @@ impl EditableComponent {
             Self::Observability(_) => "Observability",
             Self::Adaptive(_) => "Adaptive",
             Self::NemoGuardrails(_) => "NeMo Guardrails",
+            Self::PiiRedaction(_) => "PII Redaction",
         }
     }
 
@@ -52,6 +55,7 @@ impl EditableComponent {
             Self::Observability(_) => ObservabilityConfig::editor_schema().fields,
             Self::Adaptive(_) => AdaptiveConfig::editor_schema().fields,
             Self::NemoGuardrails(_) => NeMoGuardrailsConfig::editor_schema().fields,
+            Self::PiiRedaction(_) => PiiRedactionConfig::editor_schema().fields,
         }
     }
 
@@ -60,6 +64,7 @@ impl EditableComponent {
             Self::Observability(state) => state.enabled,
             Self::Adaptive(state) => state.enabled,
             Self::NemoGuardrails(state) => state.enabled,
+            Self::PiiRedaction(state) => state.enabled,
         }
     }
 
@@ -68,6 +73,7 @@ impl EditableComponent {
             Self::Observability(state) => state.toggle_enabled(),
             Self::Adaptive(state) => state.toggle_enabled(),
             Self::NemoGuardrails(state) => state.toggle_enabled(),
+            Self::PiiRedaction(state) => state.toggle_enabled(),
         }
     }
 
@@ -76,6 +82,7 @@ impl EditableComponent {
             Self::Observability(state) => state.set_enabled(enabled),
             Self::Adaptive(state) => state.set_enabled(enabled),
             Self::NemoGuardrails(state) => state.set_enabled(enabled),
+            Self::PiiRedaction(state) => state.set_enabled(enabled),
         }
     }
 
@@ -84,6 +91,7 @@ impl EditableComponent {
             Self::Observability(state) => state.reset_enabled(),
             Self::Adaptive(state) => state.reset_enabled(),
             Self::NemoGuardrails(state) => state.reset_enabled(),
+            Self::PiiRedaction(state) => state.reset_enabled(),
         }
     }
 
@@ -92,6 +100,7 @@ impl EditableComponent {
             Self::Observability(state) => observability_summary(state),
             Self::Adaptive(state) => adaptive_summary(state),
             Self::NemoGuardrails(state) => nemo_guardrails_summary(state),
+            Self::PiiRedaction(state) => pii_redaction_summary(state),
         }
     }
 
@@ -102,6 +111,9 @@ impl EditableComponent {
             Self::NemoGuardrails(state) => {
                 config_field_configured(&state.config, field).unwrap_or(false)
             }
+            Self::PiiRedaction(state) => {
+                config_field_configured(&state.config, field).unwrap_or(false)
+            }
         }
     }
 
@@ -119,6 +131,10 @@ impl EditableComponent {
                 reset_config_field(&mut state.config, field)?;
                 state.mark_config_touched();
             }
+            Self::PiiRedaction(state) => {
+                reset_config_field(&mut state.config, field)?;
+                state.mark_config_touched();
+            }
         }
         Ok(())
     }
@@ -128,6 +144,7 @@ impl EditableComponent {
             Self::Observability(state) => store_observability_state(config, state),
             Self::Adaptive(state) => store_adaptive_state(config, state),
             Self::NemoGuardrails(state) => store_nemo_guardrails_state(config, state),
+            Self::PiiRedaction(state) => store_pii_redaction_state(config, state),
         }
     }
 }
@@ -151,6 +168,7 @@ pub(super) fn editable_components(
         EditableComponent::Observability(Box::new(component_observability_state(config)?)),
         EditableComponent::Adaptive(Box::new(component_adaptive_state(config)?)),
         EditableComponent::NemoGuardrails(Box::new(component_nemo_guardrails_state(config)?)),
+        EditableComponent::PiiRedaction(Box::new(component_pii_redaction_state(config)?)),
     ])
 }
 
@@ -334,6 +352,12 @@ pub(super) fn component_nemo_guardrails_state(
     component_editor_state(config, NEMO_GUARDRAILS_PLUGIN_KIND, false)
 }
 
+pub(super) fn component_pii_redaction_state(
+    config: &PluginConfig,
+) -> Result<ComponentEditorState<PiiRedactionConfig>, CliError> {
+    component_editor_state(config, PII_REDACTION_PLUGIN_KIND, false)
+}
+
 pub(super) fn store_observability_state(
     config: &mut PluginConfig,
     state: &ComponentEditorState<ObservabilityConfig>,
@@ -382,6 +406,22 @@ pub(super) fn store_nemo_guardrails_state(
     Ok(())
 }
 
+pub(super) fn store_pii_redaction_state(
+    config: &mut PluginConfig,
+    state: &ComponentEditorState<PiiRedactionConfig>,
+) -> Result<(), CliError> {
+    if state.should_store(state.config_touched || pii_redaction_configured(&state.config)) {
+        store_component_editor_config(
+            config,
+            PII_REDACTION_PLUGIN_KIND,
+            state.enabled,
+            pii_redaction_config_map(&state.config)?,
+            merge_pii_redaction_editor_config,
+        );
+    }
+    Ok(())
+}
+
 fn store_component_editor_config(
     config: &mut PluginConfig,
     kind: &str,
@@ -713,6 +753,23 @@ pub(super) fn nemo_guardrails_config_map(
     }
 }
 
+pub(super) fn pii_redaction_config_map(
+    config: &PiiRedactionConfig,
+) -> Result<Map<String, Value>, CliError> {
+    let value = serde_json::to_value(config).map_err(serde_error)?;
+    match value {
+        Value::Object(mut map) => {
+            if is_version_one(map.get("version")) {
+                map.remove("version");
+            }
+            Ok(map)
+        }
+        _ => Err(CliError::Config(
+            "pii_redaction config must serialize to an object".into(),
+        )),
+    }
+}
+
 pub(super) fn merge_observability_editor_config(
     existing: &mut Map<String, Value>,
     edited: Map<String, Value>,
@@ -755,6 +812,21 @@ pub(super) fn merge_nemo_guardrails_editor_config(
     );
 }
 
+pub(super) fn merge_pii_redaction_editor_config(
+    existing: &mut Map<String, Value>,
+    edited: Map<String, Value>,
+) {
+    if is_version_one(existing.get("version")) {
+        existing.remove("version");
+    }
+    merge_known_editor_object(
+        existing,
+        edited,
+        &nested_editor_keys(PiiRedactionConfig::editor_schema()),
+        PiiRedactionConfig::editor_schema(),
+    );
+}
+
 fn is_version_one(value: Option<&Value>) -> bool {
     value.and_then(Value::as_u64) == Some(1)
 }
@@ -911,3 +983,30 @@ pub(super) fn nemo_guardrails_summary(
         }
     )
 }
+
+pub(super) fn pii_redaction_configured(config: &PiiRedactionConfig) -> bool {
+    PiiRedactionConfig::editor_schema()
+        .fields
+        .iter()
+        .filter(|field| field.name != POLICY_SECTION)
+        .any(|field| config_field_configured(config, *field).unwrap_or(false))
+}
+
+pub(super) fn pii_redaction_summary(state: &ComponentEditorState<PiiRedactionConfig>) -> String {
+    let configured_fields = PiiRedactionConfig::editor_schema()
+        .fields
+        .iter()
+        .filter(|field| field.name != POLICY_SECTION)
+        .filter(|field| config_field_configured(&state.config, **field).unwrap_or(false))
+        .map(|field| field.label)
+        .collect::<Vec<_>>();
+    format!(
+        "component {}, fields {}",
+        if state.enabled { "enabled" } else { "disabled" },
+        if configured_fields.is_empty() {
+            "none".into()
+        } else {
+            configured_fields.join(", ")
+        }
+    )
+}

crates/cli/src/server.rs

@@ -10,6 +10,7 @@ use axum::routing::{get, post};
 use axum::{Json, Router};
 use nemo_relay::plugin::{PluginConfig, clear_plugin_configuration, initialize_plugins_exact};
 use nemo_relay_adaptive::plugin_component::register_adaptive_component;
+use nemo_relay_pii_redaction::component::register_pii_redaction_component;
 use reqwest::Client;
 use serde_json::Value;
 use tokio::net::TcpListener;
@@ -203,6 +204,9 @@ impl PluginActivation {
         register_adaptive_component().map_err(|error| {
             CliError::Config(format!("adaptive plugin registration failed: {error}"))
         })?;
+        register_pii_redaction_component().map_err(|error| {
+            CliError::Config(format!("PII redaction plugin registration failed: {error}"))
+        })?;
         // Gateway already resolved its config; activate exactly (no re-discovery).
         let plugin_config: PluginConfig = serde_json::from_value(config)
             .map_err(|error| CliError::Config(format!("invalid plugin config: {error}")))?;