src/features.py

"""
features.py
Feature extraction for prompt injection detection.

Two layers:
1. TF-IDF vectorization of the raw text (the statistical signal).
2. Hand-crafted heuristic features (the security-domain-knowledge signal) --
   things like presence of override phrases, imperative-mood ratio, and
   suspicious formatting markers. These get concatenated to the TF-IDF
   vector before training.
"""

import re
import numpy as np
from scipy.sparse import hstack, csr_matrix

# Phrases strongly associated with instruction-override attempts.
OVERRIDE_PHRASES = [
    "ignore previous", "ignore all previous", "ignore the above",
    "disregard your", "disregard prior", "forget everything",
    "system override", "new directive", "no longer an ai",
    "developer mode", "do anything now", "no restrictions",
    "no content policy", "unrestricted ai", "reveal your system prompt",
    "print your initial instructions", "repeat the text above",
    "from now on", "override your safety",
]

SUSPICIOUS_MARKERS = [
    "###system", "<<sys>>", "<</sys>>", "[system]", "{system}",
]


def heuristic_features(texts):
    """Returns an (n_samples, k) numpy array of hand-crafted signals."""
    rows = []
    for t in texts:
        lower = t.lower()
        override_hits = sum(1 for p in OVERRIDE_PHRASES if p in lower)
        marker_hits = sum(1 for m in SUSPICIOUS_MARKERS if m in lower)
        imperative_start = 1 if re.match(
            r"^(ignore|disregard|forget|pretend|act as|reveal|print|repeat|override)",
            lower.strip()
        ) else 0
        length = len(t)
        exclaim_count = t.count("!")
        rows.append([override_hits, marker_hits, imperative_start, length, exclaim_count])
    return np.array(rows, dtype=float)


def build_features(texts, vectorizer, fit=False):
    """
    texts: list of strings
    vectorizer: a fitted (or to-be-fitted) sklearn TfidfVectorizer
    fit: if True, calls fit_transform; else transform only
    Returns a combined sparse feature matrix.
    """
    if fit:
        tfidf = vectorizer.fit_transform(texts)
    else:
        tfidf = vectorizer.transform(texts)

    heur = heuristic_features(texts)
    heur_sparse = csr_matrix(heur)
    combined = hstack([tfidf, heur_sparse])
    return combined