use std::collections::HashMap;
use std::io::{self, BufRead, Write};
use std::time::Instant;
use gaze::{
Action, ClassRule, CleanDocument, DefaultRule, LeakKind, LocaleTag, PiiClass, Pipeline,
RawDocument, RawMatch, RecognizerSpec, Rulepack, RulepackSource, SafetyNetFallback,
SafetyNetMode, SafetyNetPolicy, Scope, Session,
};
use gaze_recognizers::{embedded, NormalizerKind, RegexDetector, ValidatorKind};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BenchConfig {
RuleFloorCore,
RuleFloorExtended,
Pass3Kiji,
Pass3Opf,
Pass3LocaleAware,
}
#[derive(Debug, Deserialize)]
struct Request {
fixture_id: String,
locale_chain: Vec<String>,
text: String,
}
#[derive(Debug, Serialize)]
struct ManifestSpan {
raw_start: usize,
raw_end: usize,
clean_start: usize,
clean_end: usize,
class: String,
}
#[derive(Debug, Serialize)]
struct LeakSuspectSpan {
clean_start: usize,
clean_end: usize,
class: String,
safety_net_id: String,
kind: String,
}
#[derive(Debug, Serialize)]
struct Timing {
total_ms: f64,
pass1_ms: f64,
pass2_ms: Option<f64>,
pass3_ms: f64,
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = parse_config()?;
let stdin = io::stdin();
let mut stdout = io::BufWriter::new(io::stdout().lock());
let mut full_pipelines = HashMap::<String, Pipeline>::new();
let mut floor_pipelines = HashMap::<String, Pipeline>::new();
for line in stdin.lock().lines() {
let line = line?;
if line.trim().is_empty() {
continue;
}
let request: Request = serde_json::from_str(&line)?;
let locale_chain = request
.locale_chain
.iter()
.map(|locale| LocaleTag::parse(locale))
.collect::<Result<Vec<_>, _>>()?;
let cache_key = request.locale_chain.join(",");
if !full_pipelines.contains_key(&cache_key) {
full_pipelines.insert(cache_key.clone(), build_pipeline(config, &locale_chain)?);
}
if !floor_pipelines.contains_key(&cache_key) {
floor_pipelines.insert(
cache_key.clone(),
build_pipeline(floor_config(config), &locale_chain)?,
);
}
let full = full_pipelines
.get(&cache_key)
.expect("full pipeline cached");
let floor = floor_pipelines
.get(&cache_key)
.expect("floor pipeline cached");
let session = Session::new(Scope::Ephemeral)?;
let full_start = Instant::now();
let (clean_doc, manifest, report) = full.clean_with_safety_net_policy_detect_context(
&session,
RawDocument::Text(request.text.clone()),
&locale_chain,
&Default::default(),
SafetyNetPolicy::new(SafetyNetMode::Strict, SafetyNetFallback::Redact),
)?;
let total_ms = full_start.elapsed().as_secs_f64() * 1000.0;
let floor_start = Instant::now();
let _ = floor.clean_with_safety_net_policy_detect_context(
&session,
RawDocument::Text(request.text),
&locale_chain,
&Default::default(),
SafetyNetPolicy::new(SafetyNetMode::Strict, SafetyNetFallback::Redact),
)?;
let pass1_ms = floor_start.elapsed().as_secs_f64() * 1000.0;
let CleanDocument::Text(clean_text) = clean_doc else {
return Err("expected text clean document".into());
};
let manifest_spans = manifest
.into_iter()
.map(|span| ManifestSpan {
raw_start: span.raw_span.start,
raw_end: span.raw_span.end,
clean_start: span.clean_span.start,
clean_end: span.clean_span.end,
class: span.class.to_canonical_str(),
})
.collect();
let leak_suspects = report
.suspects
.into_iter()
.map(|suspect| LeakSuspectSpan {
clean_start: suspect.span.start,
clean_end: suspect.span.end,
class: suspect.class.to_canonical_str(),
safety_net_id: suspect.safety_net_id,
kind: leak_kind_name(&suspect.kind).to_string(),
})
.collect::<Vec<_>>();
serde_json::to_writer(
&mut stdout,
&Response {
fixture_id: request.fixture_id,
clean_text,
manifest_spans,
leak_suspects,
timing: Timing {
total_ms,
pass1_ms,
pass2_ms: None,
pass3_ms: (total_ms - pass1_ms).max(0.0),
},
},
)?;
stdout.write_all(b"\n")?;
stdout.flush()?;
}
Ok(())
}
#[derive(Debug, Serialize)]
struct Response {
fixture_id: String,
clean_text: String,
manifest_spans: Vec<ManifestSpan>,
leak_suspects: Vec<LeakSuspectSpan>,
timing: Timing,
}
fn parse_config() -> Result<BenchConfig, Box<dyn std::error::Error>> {
let mut args = std::env::args().skip(1);
let mut config = BenchConfig::RuleFloorExtended;
while let Some(arg) = args.next() {
if arg == "--config" {
let value = args.next().ok_or("--config requires a value")?;
config = match value.as_str() {
"rule-floor-core" => BenchConfig::RuleFloorCore,
"rule-floor-extended" => BenchConfig::RuleFloorExtended,
"pass3-kiji" => BenchConfig::Pass3Kiji,
"pass3-opf" => BenchConfig::Pass3Opf,
"pass3-locale-aware" => BenchConfig::Pass3LocaleAware,
_ => return Err(format!("unknown --config {value}").into()),
};
}
}
Ok(config)
}
fn floor_config(config: BenchConfig) -> BenchConfig {
match config {
BenchConfig::RuleFloorCore => BenchConfig::RuleFloorCore,
BenchConfig::RuleFloorExtended
| BenchConfig::Pass3Kiji
| BenchConfig::Pass3Opf
| BenchConfig::Pass3LocaleAware => BenchConfig::RuleFloorExtended,
}
}
fn build_pipeline(
config: BenchConfig,
_locale_chain: &[LocaleTag],
) -> Result<Pipeline, Box<dyn std::error::Error>> {
let mut pipeline = rule_floor_pipeline(config)?;
match config {
BenchConfig::RuleFloorCore | BenchConfig::RuleFloorExtended => {}
BenchConfig::Pass3Kiji => {
pipeline = register_kiji(pipeline)?;
}
BenchConfig::Pass3Opf => {
pipeline = register_opf(pipeline)?;
}
BenchConfig::Pass3LocaleAware => {
pipeline = register_locale_aware(pipeline)?;
}
}
Ok(pipeline)
}
fn rule_floor_pipeline(config: BenchConfig) -> Result<Pipeline, Box<dyn std::error::Error>> {
let bundle = if config == BenchConfig::RuleFloorCore {
"core"
} else {
"core-extended"
};
let rulepack = Rulepack::load(RulepackSource::Embedded(
embedded(bundle).ok_or("missing embedded rulepack")?,
))?;
let mut builder = Pipeline::builder()
.rule(DefaultRule::new(Action::Tokenize))
.rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
.rule(ClassRule::new(PiiClass::Name, Action::Tokenize))
.rule(ClassRule::new(
PiiClass::Custom("phone".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("ip_address".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("eth_address".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("postal_code".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("iban".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("credit_card".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("ssn".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("nino".to_string()),
Action::Tokenize,
))
.rule(ClassRule::new(
PiiClass::Custom("pan".to_string()),
Action::Tokenize,
));
for spec in rulepack
.recognizers
.iter()
.filter(|recognizer| recognizer.enabled)
{
if let Some(collision) = spec.collision.clone() {
builder = builder.register_collision(spec.id.clone(), collision);
}
if matches!(spec.matcher, RawMatch::Regex { .. }) {
builder = builder.recognizer(regex_from_spec(&rulepack, spec)?);
}
}
Ok(builder.build()?)
}
fn regex_from_spec(
rulepack: &Rulepack,
spec: &RecognizerSpec,
) -> Result<RegexDetector, Box<dyn std::error::Error>> {
let RawMatch::Regex {
pattern,
pattern_template,
capture_groups,
} = &spec.matcher
else {
unreachable!("caller filters non-regex recognizers");
};
let pattern = match (pattern.as_deref(), pattern_template.as_deref()) {
(Some(pattern), None) => pattern.to_string(),
(None, Some(template)) => lower_pattern_template(rulepack, template)?,
_ => return Err(format!("invalid regex recognizer pattern shape for {}", spec.id).into()),
};
let exclusions = spec
.context
.as_ref()
.map(|context| context.exclusions.clone())
.unwrap_or_default();
let validator_kind = spec
.validator
.as_ref()
.map(|validator| ValidatorKind::parse(&validator.kind))
.transpose()?;
let normalizer_kind = spec
.normalizer
.as_ref()
.map(|normalizer| NormalizerKind::parse(&normalizer.kind))
.transpose()?;
Ok(RegexDetector::with_rulepack_fields(
&pattern,
spec.class.clone(),
&spec.id,
spec.locales.clone(),
spec.scoring.base,
spec.scoring.priority,
spec.token.family.as_deref().unwrap_or("counter"),
capture_groups.clone(),
exclusions,
validator_kind,
normalizer_kind,
)?)
}
fn lower_pattern_template(
rulepack: &Rulepack,
template: &str,
) -> Result<String, Box<dyn std::error::Error>> {
let Some(locale) = &rulepack.locale else {
return Err("pattern template requires locale buckets".into());
};
let mut pattern = template.to_string();
for (bucket, values) in &locale.buckets {
let marker = format!("{{locale.{bucket}}}");
if pattern.contains(&marker) {
let alternation = values
.names
.iter()
.map(|value| regex::escape(value))
.collect::<Vec<_>>()
.join("|");
pattern = pattern.replace(&marker, &alternation);
}
}
if pattern.contains("{locale.") {
return Err(format!("unresolved locale pattern template: {pattern}").into());
}
Ok(pattern)
}
#[cfg(feature = "safety-net-kiji")]
fn register_kiji(pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
use gaze_recognizers::safety_net::kiji_distilbert::KijiDistilbertSafetyNet;
Ok(pipeline.with_safety_net(KijiDistilbertSafetyNet::from_env()?))
}
#[cfg(not(feature = "safety-net-kiji"))]
fn register_kiji(_pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
Err("compile with gaze-recognizers feature safety-net-kiji".into())
}
#[cfg(feature = "safety-net-openai")]
fn register_opf(pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
use gaze_recognizers::safety_net::openai_filter::OpenAiFilterSafetyNet;
Ok(pipeline.with_safety_net(OpenAiFilterSafetyNet::from_env()?))
}
#[cfg(not(feature = "safety-net-openai"))]
fn register_opf(_pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
Err("compile with gaze-recognizers feature safety-net-openai".into())
}
#[cfg(all(feature = "safety-net-kiji", feature = "safety-net-openai"))]
fn register_locale_aware(pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
use gaze_recognizers::safety_net::kiji_distilbert::{
KijiDistilbertSafetyNet, SubprocessKijiConfig,
};
use gaze_recognizers::safety_net::openai_filter::{
OpenAiFilterSafetyNet, SubprocessOpenAiFilterConfig,
};
let kiji_command = std::env::var_os("GAZE_KIJI_DISTILBERT_COMMAND")
.ok_or("GAZE_KIJI_DISTILBERT_COMMAND is not set")?;
let kiji_model_dir = std::env::var_os("GAZE_KIJI_DISTILBERT_MODEL_DIR")
.ok_or("GAZE_KIJI_DISTILBERT_MODEL_DIR is not set")?;
let opf_command =
std::env::var_os("GAZE_OPENAI_FILTER_OPF").ok_or("GAZE_OPENAI_FILTER_OPF is not set")?;
let opf_checkpoint = std::env::var_os("OPF_CHECKPOINT").ok_or("OPF_CHECKPOINT is not set")?;
Ok(pipeline
.with_safety_net(
OpenAiFilterSafetyNet::new(
SubprocessOpenAiFilterConfig::new(opf_command)
.with_checkpoint_path(opf_checkpoint)
.with_timeout(std::time::Duration::from_secs(30))
.with_args([
"--format",
"json",
"--output-mode",
"typed",
"--no-print-color-coded-text",
"--device",
"cpu",
]),
)
.with_locales(vec![LocaleTag::Global, LocaleTag::EnUs]),
)
.with_safety_net(
KijiDistilbertSafetyNet::new(
SubprocessKijiConfig::new(kiji_command)
.with_model_dir(kiji_model_dir)
.with_timeout(std::time::Duration::from_secs(30)),
)
.with_locales(vec![LocaleTag::DeDe]),
))
}
#[cfg(not(all(feature = "safety-net-kiji", feature = "safety-net-openai")))]
fn register_locale_aware(_pipeline: Pipeline) -> Result<Pipeline, Box<dyn std::error::Error>> {
Err("compile with gaze-recognizers features safety-net-kiji,safety-net-openai".into())
}
fn leak_kind_name(kind: &LeakKind) -> &'static str {
match kind {
LeakKind::Uncovered => "uncovered",
LeakKind::PartialBleed { .. } => "partial_bleed",
LeakKind::ClassMismatch { .. } => "class_mismatch",
_ => "unknown",
}
}