Span-level prompt-injection / jailbreak detection

[adaamko]

Problem

Guard models today emit a sequence-level verdict: this input is safe or unsafe. That is not actionable for an agent pipeline — to strip, quote, or refuse the offending content you need to know which spans of an input, tool output, or retrieved document attempt to override instructions, exfiltrate data, or coerce a role change.

LettuceDetect's machinery is already span-native and, importantly, none of it is hard-coded to hallucination labels:

• The sample schema stores arbitrary typed spans: labels is a list of {start, end, label, category, subcategory} dicts over the answer text (lettucedetect/datasets/hallucination_dataset.py:23-25, full schema at hallucination_dataset.py:10-50).
• scripts/train_span_detector.py trains a binary token tagger on any data in that schema (label semantics: prompt tokens -100, target tokens 0/1).
• The label set does not require a fixed classification head: LLMDetector accepts a custom taxonomy as a {category: description} dict (lettucedetect/detectors/llm.py:88-94, handled at llm.py:119-121), and the encoder typing head matches spans to label descriptions by similarity (lettucedetect/detectors/taxonomy_head.py, module docstring), so an injection label set is a data change, not an architecture change.

This issue is the research track for turning that machinery toward prompt-injection / jailbreak localization. It complements the safety-span issues #43 (span-based safety classification) and #44 (span supervision from Aegis 2.0), which target harmful content; this one targets adversarial instructions.

Current behavior

Nothing injection-specific exists in the repo. The pieces to reuse:

• lettucedetect/generation/injection.py — corrupts a known-good text into one with exact character-level spans, driven by a taxonomy (see lettucedetect/generation/__init__.py:1-18 for the pipeline overview). The same mechanism can inject malicious instructions into benign text.
• lettucedetect/generation/classify.py — LLM assigns a (category, subcategory) to an existing untyped span; scripts/classify_psiloqa_spans.py:1-20 is the end-to-end precedent for folding an external, example-level-labeled dataset into the span schema.
• scripts/span_eval_metrics.py — stdlib char-overlap span metrics; scripts/evaluate_span_model.py — per-source / per-language result tables.

What to do

1. Define a compact label set as {name: description} text, e.g. instruction_override, data_exfiltration, roleplay_jailbreak — a starting point, revise as the data demands. Descriptions matter: both the LLM annotator and the label-conditioned typing head consume them.
2. Pick 1-2 public prompt-injection / jailbreak corpora (they are example-level: a text is known-malicious, but the offending substring is not marked). Localize spans with an LLM annotator following the scripts/classify_psiloqa_spans.py pattern; validate every offset by exact substring match.
3. Generate balanced positives by injecting malicious spans into benign carriers (tool outputs, retrieved passages, emails) using the lettucedetect/generation/injection.py approach — injection gives exact offsets by construction.
4. Emit records in the HallucinationSample schema (task_type="prompt_injection", spans in labels with your categories) with train/dev/test splits, so the existing training and eval scripts work unmodified.
5. Hold out a small human-checked test slice (precedent: the expert-reviewed code test split; protocol described in docs/code-hallucination/).
6. Train the binary tagger with scripts/train_span_detector.py, type spans with the label-conditioned head or an LLMDetector(include_taxonomy={...}) baseline, and report span char-F1 and example-level F1 via scripts/span_eval_metrics.py.

Acceptance

This is a research issue: partial contributions are welcome and count. Any of the following is a valid outcome:

• a design note comparing label sets / source datasets, with examples;
• a pilot slice (~200 localized examples) with offset-integrity checks and inter-annotator spot-checks;
• a trained pilot model with a span-F1 / example-F1 table on the held-out slice;
• a negative result with error analysis (e.g. "LLM localization of jailbreak spans is unreliable because X") — documented failures save the next person a month.

For any code that lands: offset-producing functions are unit-tested in tests/test_injection_spans_pytest.py (pattern required by tests/pytest.ini); python -m pytest tests/test_injection_spans_pytest.py -v passes; python tests/run_pytest.py stays green; every stored span satisfies answer[start:end] == text.

Non-goals

• No production guardrail API or release commitment — exploratory research only.
• No example-level (sequence) classifier; that already exists elsewhere, span localization is the point.
• No claims of adversarial robustness against attackers who know the detector.
• Harmful-content safety spans are Span-based safety classification (GliGuard / GuardReasoner-style) #43 / Generate span-level safety supervision from Aegis 2.0 #44, not this issue.

Start here

git clone https://github.com/KRLabsOrg/LettuceDetect.git
cd LettuceDetect
pip install -e ".[dev]"
python tests/run_pytest.py   # should be green before you change anything

# The code to read first:
sed -n '1,40p' lettucedetect/generation/__init__.py    # generation pipeline overview
sed -n '1,25p' scripts/classify_psiloqa_spans.py       # LLM span-annotation precedent
grep -n "labels" lettucedetect/datasets/hallucination_dataset.py | head