prompt-sanitizer 0.1.0

data/lib/prompt_sanitizer/engines/ner_engine.rb

@@ -0,0 +1,279 @@
+# frozen_string_literal: true
+
+module PromptSanitizer
+  module Engines
+    # NER Engine — Layer 2 of prompt-sanitizer (SMART / FULL mode only).
+    #
+    # Detects context-dependent PII that regex cannot catch: person names,
+    # organisations, locations, and miscellaneous named entities.
+    #
+    # Two backends are supported:
+    #
+    #   :informers — uses the `informers` gem with Xenova/distilbert-NER or
+    #                Xenova/bert-base-NER (ONNX int8).  F1 92.17.  ~25–50 ms.
+    #                Model auto-downloaded to ~/.cache/huggingface/ on first use.
+    #                Recommended default.
+    #
+    #   :mitie     — uses the `mitie` gem with the MITIE C++ NER model.
+    #                F1 88.10.  ~2 ms.  Requires separate model download (~600 MB).
+    #                Use when sub-5 ms latency is required.
+    #
+    # Both backends are optional runtime dependencies.  When neither gem is
+    # installed, NEREngine#available? returns false and #detect returns [].
+    # The Sanitizer falls back to FAST mode silently in this case.
+    #
+    # Thread safety: both ONNX Runtime sessions and MITIE models are immutable
+    # after loading — safe to share across Puma threads.
+    class NEREngine
+      # Maximum number of characters per chunk when splitting long prompts.
+      # distilbert / bert-base have a 512 subword-token limit; ~300 words
+      # ≈ 1,800 characters is a safe conservative ceiling.
+      CHUNK_SIZE    = 1_800
+      CHUNK_OVERLAP = 200  # overlap between chunks to avoid edge-case misses
+
+      # BIO tag → EntityType mapping (CoNLL-2003 schema)
+      TAG_MAP = {
+        "PER"  => EntityType::PERSON,
+        "ORG"  => EntityType::ORGANIZATION,
+        "LOC"  => EntityType::LOCATION,
+        "MISC" => EntityType::MISC,
+      }.freeze
+
+      # ── Construction ────────────────────────────────────────────────────────
+
+      # @param backend [Symbol]  :informers (default) or :mitie
+      # @param model   [String]  "distilbert" (default) or "bert-base" for
+      #                          the informers backend; path to ner_model.dat
+      #                          for the mitie backend.
+      def initialize(backend: :informers, model: "distilbert")
+        @backend = backend
+        @model   = model
+        @pipeline = nil
+        @mutex    = Mutex.new
+        _load_backend
+      end
+
+      # Returns true when the chosen backend gem is installed and the model
+      # is ready to use.
+      def available?
+        !@pipeline.nil?
+      end
+
+      # Detect named entities in +text+ and return Array<DetectedEntity>.
+      # Returns [] immediately when the backend is unavailable.
+      #
+      # Long texts are automatically chunked and results are merged.
+      def detect(text)
+        return [] unless available?
+
+        safe_text = text.encode("UTF-8", invalid: :replace, undef: :replace, replace: "")
+        return [] if safe_text.strip.empty?
+
+        if safe_text.length > CHUNK_SIZE
+          _detect_chunked(safe_text)
+        else
+          _detect_single(safe_text)
+        end
+      rescue => e
+        # Never let NER failures break the sanitizer — degrade gracefully.
+        warn "[PromptSanitizer] NER error (#{@backend}): #{e.message}"
+        []
+      end
+
+      # ── Private ─────────────────────────────────────────────────────────────
+
+      private
+
+      def _load_backend
+        case @backend
+        when :informers then _load_informers
+        when :mitie     then _load_mitie
+        else
+          raise ConfigurationError, "Unknown NER backend: #{@backend.inspect}. Use :informers or :mitie."
+        end
+      end
+
+      # ── informers backend ────────────────────────────────────────────────────
+
+      def _load_informers
+        require "informers"
+
+        model_id = case @model
+                   when "distilbert" then "Xenova/distilbert-NER"
+                   when "bert-base"  then "Xenova/bert-base-NER"
+                   else @model  # allow fully-qualified HuggingFace model IDs
+                   end
+
+        # Load once and memoize — thread-safe after initialization.
+        @pipeline     = Informers.pipeline("ner", model_id, dtype: "int8")
+        @backend_type = :informers
+      rescue LoadError
+        # informers gem not installed — NER silently unavailable.
+        @pipeline = nil
+      rescue => e
+        warn "[PromptSanitizer] Failed to load informers NER model: #{e.message}"
+        @pipeline = nil
+      end
+
+      # ── mitie backend ────────────────────────────────────────────────────────
+
+      def _load_mitie
+        require "mitie"
+
+        model_path = @model == "distilbert" || @model == "bert-base" ? nil : @model
+        model_path ||= ENV.fetch("MITIE_MODEL_PATH", "ner_model.dat")
+
+        unless File.exist?(model_path)
+          warn "[PromptSanitizer] MITIE model not found at #{model_path}. " \
+               "Download from https://github.com/mit-nlp/MITIE/releases"
+          @pipeline = nil
+          return
+        end
+
+        @pipeline     = Mitie::NER.new(model_path)
+        @backend_type = :mitie
+      rescue LoadError
+        @pipeline = nil
+      rescue => e
+        warn "[PromptSanitizer] Failed to load MITIE NER model: #{e.message}"
+        @pipeline = nil
+      end
+
+      # ── Detection ────────────────────────────────────────────────────────────
+
+      def _detect_single(text)
+        case @backend_type
+        when :informers then _informers_detect(text, offset: 0)
+        when :mitie     then _mitie_detect(text, offset: 0)
+        else []
+        end
+      end
+
+      # Split long text into overlapping chunks, detect in each, then merge.
+      # Entities whose start_pos falls in the overlap zone of a later chunk
+      # are deduplicated by (original, start_pos) pair.
+      def _detect_chunked(text)
+        entities = []
+        seen     = {}
+        pos      = 0
+
+        while pos < text.length
+          chunk      = text[pos, CHUNK_SIZE]
+          chunk_hits = _detect_single(chunk).map do |e|
+            DetectedEntity.new(
+              entity_type: e.entity_type,
+              original:    e.original,
+              replacement: nil,
+              start_pos:   e.start_pos + pos,
+              end_pos:     e.end_pos   + pos,
+              confidence:  e.confidence,
+              layer:       :ner
+            )
+          end
+
+          chunk_hits.each do |e|
+            key = "#{e.original}:#{e.start_pos}"
+            next if seen[key]
+
+            seen[key] = true
+            entities << e
+          end
+
+          pos += CHUNK_SIZE - CHUNK_OVERLAP
+        end
+
+        entities
+      end
+
+      # ── informers result parsing ─────────────────────────────────────────────
+
+      # informers returns BIO-tagged word pieces; merge consecutive I- tags
+      # back into a single span.
+      def _informers_detect(text, offset:)
+        raw = @pipeline.(text)
+        return [] if raw.nil? || raw.empty?
+
+        entities  = []
+        current   = nil
+
+        Array(raw).each do |token|
+          tag_raw = token[:entity] || token["entity"] || ""
+          word    = token[:word]   || token["word"]   || ""
+          score   = (token[:score] || token["score"] || 0.0).to_f
+          t_start = (token[:start] || token["start"] || 0).to_i
+          t_end   = (token[:end]   || token["end"]   || 0).to_i
+
+          # Parse BIO prefix and base tag
+          bio, tag = tag_raw.split("-", 2)
+          entity_type = TAG_MAP[tag]
+          next if entity_type.nil?
+
+          if bio == "B" || (bio == "I" && current.nil?)
+            # Flush previous entity
+            entities << _build_entity(current, text) if current
+            current = { type: entity_type, start: t_start, end: t_end, score: score, tokens: [word] }
+
+          elsif bio == "I" && current && current[:type] == entity_type
+            # Continue current entity — extend span
+            current[:end]    = t_end
+            current[:score]  = [current[:score], score].min  # conservative: use lowest
+            current[:tokens] << word
+
+          else
+            # Tag changed mid-sequence — flush and start fresh
+            entities << _build_entity(current, text) if current
+            current = nil
+          end
+        end
+
+        entities << _build_entity(current, text) if current
+        entities.compact
+      end
+
+      def _build_entity(data, text)
+        return nil if data.nil?
+
+        raw_value = text[data[:start]...data[:end]]
+        return nil if raw_value.nil? || raw_value.strip.empty?
+
+        DetectedEntity.new(
+          entity_type: data[:type],
+          original:    raw_value.strip,
+          replacement: nil,
+          start_pos:   data[:start],
+          end_pos:     data[:end],
+          confidence:  data[:score],
+          layer:       :ner
+        )
+      end
+
+      # ── MITIE result parsing ─────────────────────────────────────────────────
+
+      def _mitie_detect(text, offset:)
+        doc = @pipeline.doc(text)
+        doc.entities.filter_map do |entity|
+          tag         = entity[:tag]&.upcase
+          entity_type = TAG_MAP[tag]
+          next unless entity_type
+
+          value = entity[:text]
+          next if value.nil? || value.strip.empty?
+
+          # MITIE provides character offset via :offset
+          char_start = entity[:offset] || text.index(value) || 0
+          char_end   = char_start + value.length
+
+          DetectedEntity.new(
+            entity_type: entity_type,
+            original:    value.strip,
+            replacement: nil,
+            start_pos:   char_start,
+            end_pos:     char_end,
+            confidence:  (entity[:score] || 0.80).to_f.clamp(0.0, 1.0),
+            layer:       :ner
+          )
+        end
+      end
+    end
+  end
+end

data/lib/prompt_sanitizer/engines/regex_engine.rb

@@ -0,0 +1,216 @@
+# frozen_string_literal: true
+
+module PromptSanitizer
+  module Engines
+    # Regex Engine — Layer 1 of prompt-sanitizer.
+    #
+    # Detects structured PII (email, phone, SSN, credit cards, IBANs, IPs,
+    # crypto addresses, MAC addresses, URLs, passport numbers, driving licences,
+    # and date patterns) using regular expressions with optional checksum
+    # validation (Luhn for credit cards, IBAN mod-97).
+    #
+    # All patterns run on every sanitize() call regardless of Mode.
+    class RegexEngine
+      Pattern = Struct.new(:entity_type, :regex, :confidence, :validator, keyword_init: true)
+
+      # ── Validators ──────────────────────────────────────────────────────────
+
+      # Luhn algorithm — validates credit/debit card numbers.
+      def self.luhn_valid?(card)
+        digits = card.gsub(/\D/, "").chars.map(&:to_i)
+        return false if digits.size < 13
+
+        total = digits.reverse.each_with_index.sum do |d, i|
+          i.odd? ? [d * 2 - 9, d * 2].min + (d * 2 >= 10 ? 0 : 0) : d
+          # Standard Luhn: double every second digit from the right
+          if i.odd?
+            doubled = d * 2
+            doubled > 9 ? doubled - 9 : doubled
+          else
+            d
+          end
+        end
+        total % 10 == 0
+      end
+
+      # IBAN mod-97 validation.
+      def self.iban_valid?(iban)
+        raw = iban.gsub(/[\s\-]/, "").upcase
+        return false unless raw.length.between?(15, 34)
+
+        rearranged = raw[4..] + raw[0..3]
+        numeric    = rearranged.chars.map { |c| c =~ /[A-Z]/ ? (c.ord - 55).to_s : c }.join
+        numeric.to_i % 97 == 1
+      rescue ArgumentError
+        false
+      end
+
+      # ── Pattern registry ─────────────────────────────────────────────────────
+
+      PATTERNS = [
+        # ── Email ──────────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::EMAIL,
+          regex:       /(?<![a-zA-Z0-9._%+\-])[a-zA-Z0-9._%+\-]{1,64}@[a-zA-Z0-9.\-]{1,253}\.[a-zA-Z]{2,}(?![a-zA-Z0-9._%+\-@])/i,
+          confidence:  0.99
+        ),
+
+        # ── US phone ───────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::PHONE,
+          regex:       /(?<!\d)(?:\+?1[\s.\-]?)?(?:\([2-9]\d{2}\)|[2-9]\d{2})[\s.\-]?\d{3}[\s.\-]?\d{4}(?!\d)/,
+          confidence:  0.85
+        ),
+
+        # ── International phone — compact E.164 e.g. +447946123456 ────────────
+        Pattern.new(
+          entity_type: EntityType::PHONE,
+          regex:       /(?<!\d)\+[1-9]\d{6,14}(?!\d)/,
+          confidence:  0.80
+        ),
+
+        # ── International phone — spaced/dashed e.g. +44 20 7946 0958 ─────────
+        Pattern.new(
+          entity_type: EntityType::PHONE,
+          regex:       /(?<!\d)\+[1-9]\d{0,3}(?:[\s.\-]\d{2,4}){2,4}(?!\d)/,
+          confidence:  0.78
+        ),
+
+        # ── US SSN ─────────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::SSN,
+          regex:       /(?<!\d)(?!000|666|9\d{2})\d{3}[\s\-](?!00)\d{2}[\s\-](?!0000)\d{4}(?!\d)/,
+          confidence:  0.95
+        ),
+
+        # ── Credit / debit card (Luhn-validated) ──────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::CREDIT_CARD,
+          regex:       /(?<!\d)(?:4[0-9]{3}|5[1-5][0-9]{2}|3[47][0-9]{2}|3(?:0[0-5]|[68][0-9])[0-9]|6(?:011|5[0-9]{2})|(?:2131|1800|35\d{3}))[\s\-]?(?:\d{4}[\s\-]?){2}\d{1,4}(?!\d)/,
+          confidence:  0.95,
+          validator:   ->(m) { luhn_valid?(m) }
+        ),
+
+        # ── IBAN (mod-97 validated) ────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::IBAN,
+          regex:       /\b[A-Z]{2}\d{2}(?:\s?[A-Z0-9]{4}){2,7}\s?[A-Z0-9]{1,4}\b/i,
+          confidence:  0.92,
+          validator:   ->(m) { iban_valid?(m) }
+        ),
+
+        # ── IPv4 ───────────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::IP_ADDRESS,
+          regex:       /(?<!\d)(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?)(?!\d)/,
+          confidence:  0.90
+        ),
+
+        # ── IPv6 (full and compressed forms) ──────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::IP_ADDRESS,
+          regex:       /(?<![:\w])(?:(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}|(?:[0-9a-fA-F]{1,4}:){1,7}:|(?:[0-9a-fA-F]{1,4}:){1,6}:[0-9a-fA-F]{1,4}|(?:[0-9a-fA-F]{1,4}:){1,5}(?::[0-9a-fA-F]{1,4}){1,2}|::(?:[0-9a-fA-F]{1,4}:){0,5}[0-9a-fA-F]{1,4}|[0-9a-fA-F]{1,4}::(?:[0-9a-fA-F]{1,4}:){0,4}[0-9a-fA-F]{1,4})(?![:\w])/i,
+          confidence:  0.90
+        ),
+
+        # ── MAC address ────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::MAC_ADDRESS,
+          regex:       /(?<![:\w])(?:[0-9a-fA-F]{2}[:\-]){5}[0-9a-fA-F]{2}(?![:\w])/i,
+          confidence:  0.90
+        ),
+
+        # ── URL (http/https) ───────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::URL,
+          regex:       /https?:\/\/(?:[a-zA-Z0-9\-._~:\/?#\[\]@!$&'()*+,;=%]|(?:%[0-9a-fA-F]{2}))+/i,
+          confidence:  0.85
+        ),
+
+        # ── Bitcoin address (P2PKH, P2SH, Bech32) ─────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::CRYPTO_ADDRESS,
+          regex:       /(?<![a-zA-Z0-9])(?:[13][a-km-zA-HJ-NP-Z1-9]{25,34}|bc1[qpzry9x8gf2tvdw0s3jn54khce6mua7l]{6,87})(?![a-zA-Z0-9])/,
+          confidence:  0.88
+        ),
+
+        # ── Ethereum address ───────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::CRYPTO_ADDRESS,
+          regex:       /(?<![a-fA-F0-9])0x[0-9a-fA-F]{40}(?![0-9a-fA-F])/,
+          confidence:  0.92
+        ),
+
+        # ── US Passport ────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::PASSPORT,
+          regex:       /(?<![A-Z0-9])[A-Z]{1,2}\d{7,9}(?![A-Z0-9])/,
+          confidence:  0.72
+        ),
+
+        # ── US ZIP code ────────────────────────────────────────────────────────
+        Pattern.new(
+          entity_type: EntityType::ZIP_CODE,
+          regex:       /(?<!\d)\d{5}(?:-\d{4})?(?!\d)/,
+          confidence:  0.55
+        ),
+
+        # ── Date patterns (DD/MM/YYYY, YYYY-MM-DD, Month DD YYYY, etc.) ────────
+        Pattern.new(
+          entity_type: EntityType::DATE,
+          regex:       /(?<!\d)(?:\d{1,2}[\/\-\.]\d{1,2}[\/\-\.]\d{2,4}|\d{4}[\/\-\.]\d{1,2}[\/\-\.]\d{1,2}|(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.?\s+\d{1,2},?\s+\d{4}|\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{4})(?!\d)/i,
+          confidence:  0.75
+        ),
+      ].freeze
+
+      # ── Instance ─────────────────────────────────────────────────────────────
+
+      def initialize
+        @patterns = PATTERNS.dup
+      end
+
+      # Register a custom regex pattern at runtime.
+      #
+      #   engine.add_pattern(:custom, /\bACME-\d{6}\b/, confidence: 0.90)
+      #   engine.add_pattern(:custom, "ACME-\\d{6}", confidence: 0.90)
+      def add_pattern(entity_type, regex, confidence: 0.80, validator: nil)
+        regex = Regexp.new(regex) if regex.is_a?(String)
+        @patterns << Pattern.new(
+          entity_type: entity_type,
+          regex:       regex,
+          confidence:  confidence,
+          validator:   validator
+        )
+      end
+
+      # Run all patterns against +text+ and return an Array of DetectedEntity.
+      # Overlapping matches from different patterns are kept — deduplication
+      # happens in the Sanitizer, which has the full multi-engine view.
+      def detect(text)
+        safe_text = text.encode("UTF-8", invalid: :replace, undef: :replace, replace: "")
+        entities  = []
+
+        @patterns.each do |pat|
+          safe_text.scan(pat.regex) do
+            m     = Regexp.last_match
+            value = m[0]
+
+            next if pat.validator && !self.class.instance_exec(value, &pat.validator)
+
+            entities << DetectedEntity.new(
+              entity_type: pat.entity_type,
+              original:    value,
+              replacement: nil,
+              start_pos:   m.begin(0),
+              end_pos:     m.end(0),
+              confidence:  pat.confidence,
+              layer:       :regex
+            )
+          end
+        end
+
+        entities
+      end
+    end
+  end
+end