diff --git a/backend/api.py b/backend/api.py index 248dee0..36eacb3 100644 --- a/backend/api.py +++ b/backend/api.py @@ -43,6 +43,7 @@ configure_rate_limiting, rate_limit) import requests from routes.analytics import analytics_bp, record_scan +from text_preparation import prepare_text from utils.spamSeverity import calculate_spam_severity @@ -826,6 +827,7 @@ def heuristic_url_is_malicious(url): tld = host.rsplit(".", 1)[-1] if "." in host else "" return tld in SUSPICIOUS_TLDS + MAX_MESSAGE_LENGTH = int(os.getenv("MAX_MESSAGE_LENGTH", 5000)) MAX_MESSAGE_LENGTH = settings.max_message_length @@ -1214,23 +1216,32 @@ def predict(): if text is None or (isinstance(text, str) and not text.strip()): with open(LOG_FILE, "a") as f: - f.write(f"WARNING: No text provided at {__import__('datetime').datetime.now()}\n") + f.write( + f"WARNING: No text provided at {__import__('datetime').datetime.now()}\n" + ) return jsonify({"error": "No text provided"}), 400 if not isinstance(text, str): - return jsonify({ - "error": f"'text' must be a string, got {type(text).__name__}" - }), 400 - + return ( + jsonify( + {"error": f"'text' must be a string, got {type(text).__name__}"} + ), + 400, + ) # Maximum-length validation before any vectorization/inference work. if len(text) > MAX_MESSAGE_LENGTH: - return jsonify({ - "error": ( - f"'text' exceeds maximum length of {MAX_MESSAGE_LENGTH} " - f"characters (got {len(text)})" - ) - }), 400 + return ( + jsonify( + { + "error": ( + f"'text' exceeds maximum length of {MAX_MESSAGE_LENGTH} " + f"characters (got {len(text)})" + ) + } + ), + 400, + ) # Read the live serving objects through the shared state so a # POST /reload-model hot-swap is picked up here without a restart # (#973). One snapshot per request keeps the model, vectorizer and @@ -1250,7 +1261,7 @@ def predict(): # request also repopulates the cache). cache_options = {"type": input_type} cache_key = predict_cache.make_cache_key( - normalizer.normalize(text), serving.version, cache_options + prepare_text(text), serving.version, cache_options ) cache_bypass = _cache_bypass_requested() if not cache_bypass: @@ -1294,7 +1305,9 @@ def predict(): if final_output == "safe" and heuristic_url_is_malicious(text): final_output = "malicious" else: - text_vector = serving.vectorizer.transform([text]) + # Prepared after translation, so the string handed to the vectorizer + # is the one the model was trained on regardless of source language. + text_vector = serving.vectorizer.transform([prepare_text(text)]) prediction = serving.model.predict(text_vector) final_output = serving.label_encoder.inverse_transform(prediction)[0] @@ -2283,14 +2296,16 @@ def imap_status(): if not conn_row: return jsonify({"connected": False}) - return jsonify({ - "connected": True, - "host": conn_row["host"], - "imap_username": conn_row["imap_username"], - "scan_interval_minutes": conn_row["scan_interval_minutes"], - "consent_given_at": conn_row["consent_given_at"], - "last_scan_at": conn_row["last_scan_at"], - }) + return jsonify( + { + "connected": True, + "host": conn_row["host"], + "imap_username": conn_row["imap_username"], + "scan_interval_minutes": conn_row["scan_interval_minutes"], + "consent_given_at": conn_row["consent_given_at"], + "last_scan_at": conn_row["last_scan_at"], + } + ) @app.route("/imap/schedule", methods=["PUT"]) @@ -2304,14 +2319,26 @@ def imap_schedule(): scan_interval_minutes = data.get("scan_interval_minutes") if scan_interval_minutes not in imap_store.ALLOWED_INTERVALS: - return jsonify({"error": f"scan_interval_minutes must be one of {imap_store.ALLOWED_INTERVALS}"}), 400 + return ( + jsonify( + { + "error": f"scan_interval_minutes must be one of {imap_store.ALLOWED_INTERVALS}" + } + ), + 400, + ) if not imap_store.get_connection(username): return jsonify({"error": "No connected inbox found for this account"}), 404 imap_store.update_schedule(username, scan_interval_minutes) _schedule_user_job(username, scan_interval_minutes) - return jsonify({"message": "Scan schedule updated", "scan_interval_minutes": scan_interval_minutes}) + return jsonify( + { + "message": "Scan schedule updated", + "scan_interval_minutes": scan_interval_minutes, + } + ) @app.route("/imap/disconnect", methods=["POST"]) @@ -2349,7 +2376,11 @@ def imap_scan_now(): try: emails = imap_connector.fetch_imap_emails( - conn_row["host"], conn_row["port"], conn_row["imap_username"], password, limit=50 + conn_row["host"], + conn_row["port"], + conn_row["imap_username"], + password, + limit=50, ) scan_results = scan_emails_with_model(emails) imap_store.save_scan_results(username, scan_results["emails"]) diff --git a/backend/retrain.py b/backend/retrain.py index c81d373..46bb322 100644 --- a/backend/retrain.py +++ b/backend/retrain.py @@ -12,8 +12,8 @@ 1. Loads the original training dataset (DATASET_PATH env var, default: dataset.csv) 2. Loads feedback_store.csv (the corrected labels submitted via /feedback) 3. Merges them into one training set (feedback's `correct_label` becomes the label) - 4. Encodes labels once with a single LabelEncoder and normalizes text with the - same normalizer api.py uses at inference time. + 4. Encodes labels once with a single LabelEncoder and prepares text with the + shared contract every inference path applies. 5. Fits the vectorizer + LinearSVC ONCE on a held-out train split to report an honest accuracy, then refits ONCE on the full combined data for the artifacts actually written to disk. @@ -46,7 +46,7 @@ from sklearn.preprocessing import LabelEncoder from sklearn.svm import LinearSVC -from utils.text_normalizer import normalizer +from text_preparation import prepare_text VALID_LABELS = {"ham", "spam", "smishing"} @@ -170,13 +170,13 @@ def train( ): """Deterministic training pipeline. - Text is normalized with the same normalizer api.py applies at inference, so - the vectorizer vocabulary matches what serving will see. Labels are encoded - ONCE and the encoded integers are used for every fit -- no raw string labels - leak into any model. The held-out fit and the production fit each happen - exactly once. + Text goes through the shared preparation contract that every inference path + also applies, so the vectorizer vocabulary matches what serving will see. + Labels are encoded ONCE and the encoded integers are used for every fit -- no + raw string labels leak into any model. The held-out fit and the production + fit each happen exactly once. """ - normalized = combined["text"].apply(normalizer.normalize) + normalized = combined["text"].apply(prepare_text) label_encoder = LabelEncoder() y = label_encoder.fit_transform(combined["label"]) diff --git a/backend/tests/test_retrain.py b/backend/tests/test_retrain.py index 55de778..3705193 100644 --- a/backend/tests/test_retrain.py +++ b/backend/tests/test_retrain.py @@ -124,7 +124,7 @@ def test_labels_are_encoded_consistently_everywhere(trained): assert set(trained.label_encoder.classes_) == retrain.VALID_LABELS sample = trained.vectorizer.transform( - [retrain.normalizer.normalize("free prize claim now")] + [retrain.prepare_text("free prize claim now")] ) encoded_pred = trained.model.predict(sample) decoded = trained.label_encoder.inverse_transform(encoded_pred)[0] diff --git a/backend/tests/test_text_preparation.py b/backend/tests/test_text_preparation.py new file mode 100644 index 0000000..870d052 --- /dev/null +++ b/backend/tests/test_text_preparation.py @@ -0,0 +1,51 @@ +"""Parity coverage for the shared text-preparation contract (issue #1037). + +The value of the contract is that training and serving cannot disagree, so these +tests assert the property that matters -- obfuscated input reduces to the same +canonical string as its plain equivalent -- rather than pinning the normalizer's +internal steps, which belong to its own tests. +""" + +import text_preparation +from text_preparation import prepare_text + + +class TestCanonicalForm: + def test_zero_width_characters_are_stripped(self): + assert prepare_text("Free\u200b Prize\u200d") == prepare_text("Free Prize") + + def test_cyrillic_homoglyphs_fold_to_latin(self): + # "claim" spelled with Cyrillic es, a and i. + assert prepare_text("\u0441l\u0430\u0456m") == prepare_text("claim") + + def test_spaced_out_words_are_rejoined(self): + assert prepare_text("f r e e money") == prepare_text("free money") + + def test_repeated_whitespace_collapses(self): + assert prepare_text("win a prize") == prepare_text("win a prize") + + def test_already_canonical_text_is_unchanged(self): + assert prepare_text("claim your free prize") == "claim your free prize" + + def test_preparation_is_idempotent(self): + once = prepare_text("F r e e\u200b m\u043en\u0435y") + assert prepare_text(once) == once + + +class TestNonStringInput: + def test_none_passes_through(self): + assert prepare_text(None) is None + + def test_empty_string_passes_through(self): + assert prepare_text("") == "" + + def test_non_string_passes_through(self): + assert prepare_text(42) == 42 + + +class TestTrainServeParity: + def test_training_and_inference_share_one_entry_point(self): + """Both regimes must import the same callable, not two copies of it.""" + import retrain + + assert retrain.prepare_text is text_preparation.prepare_text diff --git a/backend/text_preparation.py b/backend/text_preparation.py new file mode 100644 index 0000000..4a6bc55 --- /dev/null +++ b/backend/text_preparation.py @@ -0,0 +1,39 @@ +"""The single text-preparation contract shared by training and inference. + +``retrain.py`` fits the TF-IDF vocabulary on text that has been run through +:data:`~utils.text_normalizer.normalizer`, so anything that reaches +``vectorizer.transform`` at serving time must be prepared the same way or the +model is scoring a different alphabet than the one it learned. Homoglyph +substitutions, zero-width joiners and spaced-out words -- exactly the evasions +the normalizer exists to undo -- otherwise survive into the vectorizer and fall +out of vocabulary. + +Every producer of model input calls :func:`prepare_text`; nothing calls the +normalizer directly. Routing both regimes through one function is what makes the +parity checkable rather than a convention that drifts. + +>>> prepare_text("Free\\u200b Prize") +'Free Prize' +>>> prepare_text("\\u0441laim now") +'claim now' +>>> prepare_text("f r e e money") +'free money' +>>> prepare_text("") +'' +>>> prepare_text(None) is None +True +""" + +from utils.text_normalizer import normalizer + +__all__ = ["prepare_text"] + + +def prepare_text(text): + """Return ``text`` in the canonical form the model was trained on. + + Non-string input is handed back untouched: callers upstream of validation + (bulk rows, mailbox payloads) can pass ``None`` or a stray numeric cell, and + a preparation step is the wrong place to decide that is an error. + """ + return normalizer.normalize(text)