Mask contact details in a chat message

"""
Catch obfuscated contact details with a light classifier.

Rung N1. The regular expressions of N0 see one spelling of a phone number.
This sees the shape of one: a run of tokens that is mostly digits, mostly
short, and sitting next to words like "call" or "reach".

The weights are small enough to keep in the repository next to this file, and
there is no service to run: the model loads with the process.
"""

import re
import unicodedata

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

# Digits written as words, and the characters used to stand in for digits.
DIGIT_WORDS = ("zero|one|two|three|four|five|six|seven|eight|nine|ten|"
               "zéro|un|deux|trois|quatre|cinq|six|sept|huit|neuf|dix")
LOOKALIKES = str.maketrans({"O": "0", "o": "0", "I": "1", "i": "1", "l": "1", "S": "5", "s": "5"})


def _is_digits_in_disguise(token: str) -> bool:
    """
    True when folding lookalike characters turns the whole token into digits.

    The token must already contain one real digit. Folding unconditionally
    would be a mistake: "loll" would fold to 1011 and read as a fragment of
    a phone number.
    """
    if not any(c.isdigit() for c in token):
        return False
    return token.translate(LOOKALIKES).isdigit()


def shape(text: str) -> str:
    """
    Turn a message into the features that matter, and drop the rest.

    Every hidden digit becomes the same symbol, whatever its spelling, so the
    classifier learns what a hidden number looks like rather than which digits
    it held.
    """
    out = []
    # Compatibility folding first: full-width digits become digits, and a
    # decomposed "zéro" becomes one word again.
    for token in re.findall(r"[^\W_]+", unicodedata.normalize("NFKC", text)):
        lowered = token.lower()
        # Case matters for lookalikes, so fold before lowering.
        if re.fullmatch(DIGIT_WORDS, lowered):
            out.append("D")
        elif _is_digits_in_disguise(token):
            out.append("D" * len(token))
        elif len(lowered) >= 2:
            out.append(lowered)
    return " ".join(out)


def train(messages: list[str], labels: list[int]):
    """`labels` is 1 when the message hides contact details, 0 when it does not."""
    model = make_pipeline(
        TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 4), min_df=1),
        LogisticRegression(class_weight="balanced", max_iter=1000),
    )
    model.fit([shape(m) for m in messages], labels)
    return model


def is_hiding_contact_details(model, message: str, threshold: float = 0.5) -> bool:
    """
    Returns a decision, and the threshold is yours to set.

    Move it towards 1 if a false positive means blocking a legitimate message.
    Move it towards 0 if letting one through is the worse outcome.
    """
    score = model.predict_proba([shape(message)])[0][1]
    return bool(score >= threshold)

/**
 * Catch obfuscated contact details with a light classifier.
 *
 * Rung N1. The regular expressions of N0 see one spelling of a phone number.
 * This sees the shape of one: a run of tokens that is mostly digits, mostly
 * short, sitting next to words like "call" or "reach".
 *
 * Written out in full rather than pulled from a library: hashing the character
 * n-grams, training and deciding take three short functions. That is the
 * whole argument of this rung: the classical tool is small enough to read.
 */

// Digits written as words, in English and in French.
const DIGIT_WORDS = new Set(
  ('zero one two three four five six seven eight nine ten ' +
    'zéro un deux trois quatre cinq six sept huit neuf dix').split(' '),
);

// Characters used to stand in for digits.
const LOOKALIKES = { O: '0', o: '0', I: '1', i: '1', l: '1', S: '5', s: '5' };

const BUCKETS = 512; // hashing trick: no vocabulary to build or ship

/**
 * True when folding lookalike characters turns the whole token into digits.
 * The token must already carry one real digit; folding unconditionally would
 * turn "loll" into 1011.
 */
function isDigitsInDisguise(token) {
  if (!/\d/.test(token)) return false;
  const folded = [...token].map((c) => LOOKALIKES[c] ?? c).join('');
  return /^\d+$/.test(folded);
}

/** Turn a message into the features that matter, and drop the rest. */
export function shape(text) {
  const out = [];
  // Compatibility folding first: full-width digits become digits, and a
  // decomposed "zéro" becomes one word again.
  for (const token of text.normalize('NFKC').match(/[\p{L}\p{N}]+/gu) ?? []) {
    const lowered = token.toLowerCase();
    if (DIGIT_WORDS.has(lowered)) out.push('D');
    else if (isDigitsInDisguise(token)) out.push('D'.repeat(token.length));
    else if (lowered.length >= 2) out.push(lowered);
  }
  return out.join(' ');
}

/** Character n-grams of the shaped message, hashed into a fixed vector. */
function features(text) {
  const shaped = ` ${shape(text)} `;
  const vector = new Float64Array(BUCKETS);
  for (let n = 2; n <= 4; n += 1) {
    for (let i = 0; i + n <= shaped.length; i += 1) {
      let h = 2166136261;
      for (const c of shaped.slice(i, i + n)) h = ((h ^ c.codePointAt(0)) * 16777619) >>> 0;
      vector[h % BUCKETS] += 1;
    }
  }
  const norm = Math.hypot(...vector);
  return norm ? vector.map((v) => v / norm) : vector;
}

/** `labels` is 1 when the message hides contact details, 0 when it does not. */
export function train(messages, labels, { epochs = 400, rate = 0.5 } = {}) {
  // An untrained model would score every message 0.5 and block them all.
  if (!labels.includes(0) || !labels.includes(1) || labels.length !== messages.length) {
    throw new RangeError('training needs messages of both labels, one label per message');
  }
  const rows = messages.map(features);
  const weights = new Float64Array(BUCKETS);
  let bias = 0;

  for (let epoch = 0; epoch < epochs; epoch += 1) {
    for (let i = 0; i < rows.length; i += 1) {
      let z = bias;
      for (let j = 0; j < BUCKETS; j += 1) z += weights[j] * rows[i][j];
      const error = 1 / (1 + Math.exp(-z)) - labels[i];
      for (let j = 0; j < BUCKETS; j += 1) weights[j] -= rate * error * rows[i][j];
      bias -= rate * error;
    }
  }
  return { weights, bias };
}

/**
 * Returns a decision, and the threshold is yours to set.
 *
 * Move it towards 1 if a false positive means blocking a legitimate message.
 * Move it towards 0 if letting one through is the worse outcome.
 */
export function isHidingContactDetails(model, message, threshold = 0.5) {
  const vector = features(message);
  let z = model.bias;
  for (let j = 0; j < BUCKETS; j += 1) z += model.weights[j] * vector[j];
  return 1 / (1 + Math.exp(-z)) >= threshold;
}

"""
Mask personal data by asking a general-purpose model.

Rung N3. It is here so you can see what it costs, not because this entry
recommends it.

Note what the code has to do that N0 did not: retry on failure, cap the input
size, parse an answer that is only probably valid JSON, check every item of it
against the message, and refuse to pass the message through unmasked when the
answer is unusable. A prompt is only a request: nothing in it holds the model to
the format asked for. That plumbing is the real cost of this rung, and it is the
part your tests have to cover, because the model itself is not testable.
"""

from __future__ import annotations

import json

PROMPT = (
    "Find every piece of personal contact information in the message below.\n"
    "Answer with JSON only: a list of objects with keys `text` and `kind`,\n"
    "where `kind` is one of email, phone, iban, address.\n"
    "If there is none, answer with an empty list.\n\n"
    "Message:\n{message}"
)
KINDS = ("email", "phone", "iban", "address")

MAX_CHARACTERS = 8000

# The provider named here is an example, not a recommendation: the reasoning
# holds for any general-purpose model API, and the client is swappable. Pass
# any object with a `complete(prompt=..., temperature=...)` method.
MODEL = "gpt-4.1-mini"  # an example id: check the parameters your model accepts


class MaskingUnavailable(Exception):
    """The provider could not be reached, or answered something unusable."""


class ProviderClient:
    """The one call this snippet makes, on top of the provider's SDK."""

    def __init__(self, sdk=None, model: str = MODEL):
        if sdk is None:  # pragma: no cover - needs a key and a network
            from openai import OpenAI

            sdk = OpenAI()
        self.sdk, self.model = sdk, model

    def complete(self, *, prompt: str, temperature: float) -> str:
        response = self.sdk.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=temperature,
        )
        return response.choices[0].message.content


def mask(message: str, client=None, *, attempts: int = 3) -> str:
    """
    Replace contact details with a label naming what was removed.

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.
    """
    client = client or ProviderClient()

    # The provider bills every token of the prompt. Refusing oversized input
    # before the call is a cost control; the cap counts characters, not tokens.
    if len(message) > MAX_CHARACTERS:
        raise ValueError(f"message longer than {MAX_CHARACTERS} characters")

    found = _ask(client, message, attempts)

    # Replace the longest matches first, so a substring never eats its parent.
    for item in sorted(found, key=lambda i: len(i["text"]), reverse=True):
        message = message.replace(item["text"], f"[{item['kind']}]")
    return message


def _is_usable(parsed, message: str) -> bool:
    """A list of {text, kind}, every text found as is in the message, every kind one we asked for."""
    return isinstance(parsed, list) and all(
        isinstance(item, dict)
        and isinstance(item.get("text"), str)
        and item["text"] != ""
        and item["text"] in message
        and item.get("kind") in KINDS
        for item in parsed
    )


def _unfenced(answer: str) -> str:
    """
    Strip a code fence that wraps the whole answer, and nothing else.

    Models often hand back `\u0060\u0060\u0060json …\u0060\u0060\u0060`, and refusing that form would
    pay for a second call for nothing. Any other departure — text before or
    after, two blocks, a fence never closed — is left alone, and fails to
    parse, which is the point.
    """
    text = answer.strip()
    if text.startswith("```") and text.endswith("```") and text.count("```") == 2:
        return text[3:-3].removeprefix("json")
    return text


def _ask(client, message: str, attempts: int) -> list[dict]:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(prompt=PROMPT.format(message=message), temperature=0)
            if answer is None:  # a refusal carries no content: unusable, not empty
                last_error = ValueError("the model returned no content")
                continue
            parsed = json.loads(_unfenced(answer))
            if _is_usable(parsed, message):
                return parsed
            # An item the message does not contain would mask nothing, and the message would leave in clear.
            last_error = ValueError("the model answered something other than items found in the message")
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise MaskingUnavailable(str(last_error))

/**
 * Mask personal data by asking a general-purpose model.
 *
 * Rung N3. It is here so you can see what it costs, not because this entry
 * recommends it.
 *
 * Note what the code has to do that N0 did not: retry on failure, cap the
 * input size, parse an answer that is only probably valid JSON, check every
 * item of it against the message, and refuse to pass the message through
 * unmasked when the answer is unusable. A prompt is only a request: nothing in
 * it holds the model to the format asked for. That plumbing is the real cost
 * of this rung, and it is the part your tests have to cover, because the model
 * itself is not testable.
 */

const PROMPT = [
  'Find every piece of personal contact information in the message below.',
  'Answer with JSON only: a list of objects with keys `text` and `kind`,',
  'where `kind` is one of email, phone, iban, address.',
  'If there is none, answer with an empty list.',
  '',
  'Message:',
].join('\n');
const KINDS = ['email', 'phone', 'iban', 'address'];

export const MAX_CHARACTERS = 8000;

// The provider named here is an example, not a recommendation: the reasoning
// holds for any general-purpose model API, and the client is swappable. Pass
// any object with a `complete({ prompt, temperature })` method.
export const MODEL = 'gpt-4.1-mini'; // an example id: check the parameters your model accepts

export class MaskingUnavailable extends Error {}

export async function providerClient(sdk, model = MODEL) {
  if (!sdk) {
    const { OpenAI } = await import('openai');
    sdk = new OpenAI();
  }
  return {
    async complete({ prompt, temperature }) {
      const response = await sdk.chat.completions.create({
        model,
        messages: [{ role: 'user', content: prompt }],
        temperature,
      });
      return response.choices[0].message.content;
    },
  };
}

/**
 * Replace contact details with a label naming what was removed.
 *
 * @param {string} message
 * @param {object} options
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 */
export async function mask(message, { client, attempts = 3 } = {}) {
  client ??= await providerClient();

  // The provider bills every token of the prompt. Refusing oversized input
  // before the call is a cost control; the cap counts characters (code points,
  // as Python does), not tokens.
  if ([...message].length > MAX_CHARACTERS) {
    throw new RangeError(`message longer than ${MAX_CHARACTERS} characters`);
  }

  const found = await ask(client, message, attempts);

  // Replace the longest matches first, so a substring never eats its parent.
  let out = message;
  for (const { text, kind } of [...found].sort((a, b) => b.text.length - a.text.length)) {
    out = out.split(text).join(`[${kind}]`);
  }
  return out;
}

/** A list of {text, kind}, every text found as is in the message, every kind one we asked for. */
function isUsable(parsed, message) {
  return Array.isArray(parsed) && parsed.every((item) => (
    typeof item?.text === 'string' && item.text !== '' && message.includes(item.text) && KINDS.includes(item.kind)
  ));
}

/**
 * Strip a code fence that wraps the whole answer, and nothing else.
 *
 * Models often hand back ```json … ```, and refusing that form would pay for a
 * second call for nothing. Any other departure — text before or after, two
 * blocks, a fence never closed — is left alone, and fails to parse, which is
 * the point.
 */
function unfenced(answer) {
  const text = answer.trim();
  const fences = text.match(/```/g)?.length ?? 0;
  if (text.startsWith('```') && text.endsWith('```') && fences === 2) {
    return text.slice(3, -3).replace(/^json/, '');
  }
  return text;
}

async function ask(client, message, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      const answer = await client.complete({
        prompt: `${PROMPT}\n${message}`,
        // Temperature zero, because a masking decision that changes between
        // two identical calls cannot be reviewed.
        temperature: 0,
      });
      if (answer === null || answer === undefined) {
        // A refusal carries no content: unusable, not empty.
        lastError = new Error('the model returned no content');
        continue;
      }
      const parsed = JSON.parse(unfenced(answer));
      // An item the message does not contain would mask nothing, and the message would leave in clear.
      if (isUsable(parsed, message)) return parsed;
      lastError = new Error('the model answered something other than items found in the message');
    } catch (error) {
      lastError = error;
    }
  }
  throw new MaskingUnavailable(String(lastError));
}