From 246e0671bba88a10ab60a178085b499c2af4b742 Mon Sep 17 00:00:00 2001 From: Enes Deniz <142517728+3nesdeniz@users.noreply.github.com> Date: Mon, 13 Jul 2026 05:39:46 +0300 Subject: [PATCH] feat: add agentic prompt-injection boundary dataset --- doc/bibliography.md | 2 +- doc/code/datasets/1_loading_datasets.ipynb | 38 ++- doc/code/datasets/1_loading_datasets.py | 23 ++ doc/references.bib | 8 + .../datasets/seed_datasets/remote/__init__.py | 12 + ...entic_prompt_injection_boundary_dataset.py | 259 ++++++++++++++++ ...entic_prompt_injection_boundary_dataset.py | 284 ++++++++++++++++++ 7 files changed, 620 insertions(+), 6 deletions(-) create mode 100644 pyrit/datasets/seed_datasets/remote/agentic_prompt_injection_boundary_dataset.py create mode 100644 tests/unit/datasets/test_agentic_prompt_injection_boundary_dataset.py diff --git a/doc/bibliography.md b/doc/bibliography.md index ccbab89d8c..af488bab48 100644 --- a/doc/bibliography.md +++ b/doc/bibliography.md @@ -5,6 +5,6 @@ All academic papers, research blogs, and technical reports referenced throughout :::{dropdown} Citation Keys :class: hidden-citations -[@aakanksha2024multilingual; @adversaai2023universal; @andriushchenko2024tense; @anthropic2024manyshot; @aqrawi2024singleturncrescendo; @atr2026; @bethany2024mathprompt; @bhardwaj2023harmfulqa; @bhardwaj2024homer; @brahman2024coconot; @bryan2025agentictaxonomy; @bullwinkel2025airtlessons; @bullwinkel2025repeng; @bullwinkel2026trigger; @chao2023pair; @chao2024jailbreakbench; @cui2024orbench; @darkbench2025; @derczynski2024garak; @ding2023wolf; @embracethered2024unicode; @embracethered2025sneakybits; @gehman2020realtoxicityprompts; @ghosh2025aegis; @ghosh2025ailuminate; @gong2025figstep; @gupta2024walledeval; @haider2024phi3safety; @han2024medsafetybench; @hines2024spotlighting; @inie2025summon; @ji2023beavertails; @ji2024pkusaferlhf; @jiang2025sosbench; @jones2025computeruse; @kingma2014adam; @li2024mossbench; @li2024saladbench; @li2024wmdp; @lin2023toxicchat; @liu2024flipattack; @liu2024mmsafetybench; @lopez2024pyrit; @luo2024jailbreakv; @lv2024codechameleon; @mazeika2023tdc; @mazeika2024harmbench; @mckee2024transparency; @mehrotra2023tap; @microsoft2024skeletonkey; @odin2024; @palaskar2025vlsu; @pfohl2024equitymedqa; @promptfoo2025ccp; @robustintelligence2024bypass; @roccia2024promptintel; @rottger2023xstest; @rottger2025msts; @russinovich2024crescendo; @russinovich2025price; @scheuerman2025transphobia; @shaikh2022second; @shayegani2025computeruse; @shen2023donotanything; @sheshadri2024lat; @souly2024strongreject; @stok2023ansi; @tan2026comicjailbreak; @tang2025multilingual; @tedeschi2024alert; @vantaylor2024socialbias; @vidgen2023simplesafetytests; @wang2023decodingtrust; @wang2023donotanswer; @wang2025siuo; @wei2023jailbroken; @xie2024sorrybench; @yu2023gptfuzzer; @yuan2023cipherchat; @zeng2024persuasion; @zhang2024cbtbench; @ziems2022mic; @zou2023gcg] +[@aakanksha2024multilingual; @adversaai2023universal; @andriushchenko2024tense; @anthropic2024manyshot; @aqrawi2024singleturncrescendo; @atr2026; @bethany2024mathprompt; @bhardwaj2023harmfulqa; @bhardwaj2024homer; @brahman2024coconot; @bryan2025agentictaxonomy; @bullwinkel2025airtlessons; @bullwinkel2025repeng; @bullwinkel2026trigger; @chao2023pair; @chao2024jailbreakbench; @cui2024orbench; @darkbench2025; @deniz2026agenticboundarypairs; @derczynski2024garak; @ding2023wolf; @embracethered2024unicode; @embracethered2025sneakybits; @gehman2020realtoxicityprompts; @ghosh2025aegis; @ghosh2025ailuminate; @gong2025figstep; @gupta2024walledeval; @haider2024phi3safety; @han2024medsafetybench; @hines2024spotlighting; @inie2025summon; @ji2023beavertails; @ji2024pkusaferlhf; @jiang2025sosbench; @jones2025computeruse; @kingma2014adam; @li2024mossbench; @li2024saladbench; @li2024wmdp; @lin2023toxicchat; @liu2024flipattack; @liu2024mmsafetybench; @lopez2024pyrit; @luo2024jailbreakv; @lv2024codechameleon; @mazeika2023tdc; @mazeika2024harmbench; @mckee2024transparency; @mehrotra2023tap; @microsoft2024skeletonkey; @odin2024; @palaskar2025vlsu; @pfohl2024equitymedqa; @promptfoo2025ccp; @robustintelligence2024bypass; @roccia2024promptintel; @rottger2023xstest; @rottger2025msts; @russinovich2024crescendo; @russinovich2025price; @scheuerman2025transphobia; @shaikh2022second; @shayegani2025computeruse; @shen2023donotanything; @sheshadri2024lat; @souly2024strongreject; @stok2023ansi; @tan2026comicjailbreak; @tang2025multilingual; @tedeschi2024alert; @vantaylor2024socialbias; @vidgen2023simplesafetytests; @wang2023decodingtrust; @wang2023donotanswer; @wang2025siuo; @wei2023jailbroken; @xie2024sorrybench; @yu2023gptfuzzer; @yuan2023cipherchat; @zeng2024persuasion; @zhang2024cbtbench; @ziems2022mic; @zou2023gcg] ::: diff --git a/doc/code/datasets/1_loading_datasets.ipynb b/doc/code/datasets/1_loading_datasets.ipynb index f118067276..00c7f0062b 100644 --- a/doc/code/datasets/1_loading_datasets.ipynb +++ b/doc/code/datasets/1_loading_datasets.ipynb @@ -17,6 +17,7 @@ "0DIN [@odin2024],\n", "Aegis [@ghosh2025aegis],\n", "Agent Threat Rules [@atr2026],\n", + "Agentic Prompt-Injection Boundary Pairs [@deniz2026agenticboundarypairs],\n", "ALERT [@tedeschi2024alert],\n", "BeaverTails [@ji2023beavertails],\n", "CBT-Bench [@zhang2024cbtbench],\n", @@ -66,10 +67,37 @@ "(`garak_audio_achilles_heel`)." ] }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Loading Agentic Prompt-Injection Boundary Pairs\n", + "\n", + "The `agentic_prompt_injection_boundary_pairs` loader defaults to attack\n", + "prompts for red-team use. It can also return the paired legitimate requests,\n", + "preserving each record's pair ID, scenario, attack family, source context,\n", + "target boundary, and expected action. For example:\n", + "\n", + "```python\n", + "from pyrit.datasets.seed_datasets.remote import (\n", + " AgenticPromptInjectionBoundaryLabel,\n", + " AgenticPromptInjectionBoundarySplit,\n", + " _AgenticPromptInjectionBoundaryDataset,\n", + ")\n", + "\n", + "paired_test_loader = _AgenticPromptInjectionBoundaryDataset(\n", + " label=AgenticPromptInjectionBoundaryLabel.ALL,\n", + " split=AgenticPromptInjectionBoundarySplit.TEST,\n", + ")\n", + "paired_test_surface = await paired_test_loader.fetch_dataset_async()\n", + "```" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "1", + "id": "2", "metadata": {}, "outputs": [ { @@ -190,7 +218,7 @@ }, { "cell_type": "markdown", - "id": "2", + "id": "3", "metadata": {}, "source": [ "## Loading Specific Datasets\n", @@ -201,7 +229,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [ { @@ -238,7 +266,7 @@ }, { "cell_type": "markdown", - "id": "4", + "id": "5", "metadata": {}, "source": [ "## Adding Datasets to Memory\n", @@ -255,7 +283,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [ { diff --git a/doc/code/datasets/1_loading_datasets.py b/doc/code/datasets/1_loading_datasets.py index 271a46731f..c956083404 100644 --- a/doc/code/datasets/1_loading_datasets.py +++ b/doc/code/datasets/1_loading_datasets.py @@ -21,6 +21,7 @@ # 0DIN [@odin2024], # Aegis [@ghosh2025aegis], # Agent Threat Rules [@atr2026], +# Agentic Prompt-Injection Boundary Pairs [@deniz2026agenticboundarypairs], # ALERT [@tedeschi2024alert], # BeaverTails [@ji2023beavertails], # CBT-Bench [@zhang2024cbtbench], @@ -69,6 +70,28 @@ # `garak_tm_system_prompts`), and an audio jailbreak set # (`garak_audio_achilles_heel`). +# %% [markdown] +# ## Loading Agentic Prompt-Injection Boundary Pairs +# +# The `agentic_prompt_injection_boundary_pairs` loader defaults to attack +# prompts for red-team use. It can also return the paired legitimate requests, +# preserving each record's pair ID, scenario, attack family, source context, +# target boundary, and expected action. For example: +# +# ```python +# from pyrit.datasets.seed_datasets.remote import ( +# AgenticPromptInjectionBoundaryLabel, +# AgenticPromptInjectionBoundarySplit, +# _AgenticPromptInjectionBoundaryDataset, +# ) +# +# paired_test_loader = _AgenticPromptInjectionBoundaryDataset( +# label=AgenticPromptInjectionBoundaryLabel.ALL, +# split=AgenticPromptInjectionBoundarySplit.TEST, +# ) +# paired_test_surface = await paired_test_loader.fetch_dataset_async() +# ``` + # %% from pyrit.datasets import SeedDatasetProvider from pyrit.memory import CentralMemory diff --git a/doc/references.bib b/doc/references.bib index c9b29c00b3..a08bddece5 100644 --- a/doc/references.bib +++ b/doc/references.bib @@ -39,6 +39,14 @@ @article{tedeschi2024alert url = {https://arxiv.org/abs/2404.08676}, } +@misc{deniz2026agenticboundarypairs, + title = {Agentic Prompt-Injection Boundary Pairs}, + author = {Enes Deniz}, + year = {2026}, + url = {https://huggingface.co/datasets/3nesdeniz/agentic-prompt-injection-boundary-pairs}, + note = {CC BY 4.0 dataset. Source repository: https://github.com/3nesdeniz/agentic-prompt-injection-boundary-pairs}, +} + @article{derczynski2024garak, title = {garak: A Framework for Security Probing Large Language Models}, author = {Leon Derczynski and Erick Galinkin and Jeffrey Martin and Subho Majumdar and Nanna Inie}, diff --git a/pyrit/datasets/seed_datasets/remote/__init__.py b/pyrit/datasets/seed_datasets/remote/__init__.py index f625c4f4f5..489822907a 100644 --- a/pyrit/datasets/seed_datasets/remote/__init__.py +++ b/pyrit/datasets/seed_datasets/remote/__init__.py @@ -17,6 +17,13 @@ ATRVariationType, _AgentThreatRulesDataset, ) +from pyrit.datasets.seed_datasets.remote.agentic_prompt_injection_boundary_dataset import ( + AgenticPromptInjectionBoundaryAttackFamily, + AgenticPromptInjectionBoundaryLabel, + AgenticPromptInjectionBoundarySourceContext, + AgenticPromptInjectionBoundarySplit, + _AgenticPromptInjectionBoundaryDataset, +) from pyrit.datasets.seed_datasets.remote.aya_redteaming_dataset import _AyaRedteamingDataset from pyrit.datasets.seed_datasets.remote.babelscape_alert_dataset import _BabelscapeAlertDataset from pyrit.datasets.seed_datasets.remote.beaver_tails_dataset import _BeaverTailsDataset @@ -124,6 +131,10 @@ __all__ = [ "AegisHarmCategory", + "AgenticPromptInjectionBoundaryAttackFamily", + "AgenticPromptInjectionBoundaryLabel", + "AgenticPromptInjectionBoundarySourceContext", + "AgenticPromptInjectionBoundarySplit", "CoCoNotCategory", "CoCoNotSplit", "DecodingTrustToxicitySubset", @@ -148,6 +159,7 @@ "ATRDetectionField", "ATRVariationType", "_AgentThreatRulesDataset", + "_AgenticPromptInjectionBoundaryDataset", "_AyaRedteamingDataset", "_BabelscapeAlertDataset", "_BeaverTailsDataset", diff --git a/pyrit/datasets/seed_datasets/remote/agentic_prompt_injection_boundary_dataset.py b/pyrit/datasets/seed_datasets/remote/agentic_prompt_injection_boundary_dataset.py new file mode 100644 index 0000000000..3277a779d8 --- /dev/null +++ b/pyrit/datasets/seed_datasets/remote/agentic_prompt_injection_boundary_dataset.py @@ -0,0 +1,259 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import logging +from enum import Enum +from typing import Any + +from typing_extensions import override + +from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import _RemoteDatasetLoader +from pyrit.models import Modality, SeedDataset, SeedPrompt, SeedUnion + +logger = logging.getLogger(__name__) + + +class AgenticPromptInjectionBoundaryLabel(Enum): + """Filter records by the published security decision label.""" + + BENIGN = 0 + ATTACK = 1 + ALL = "all" + + +class AgenticPromptInjectionBoundarySplit(Enum): + """Select a published split or combine the complete dataset.""" + + TRAIN = "train" + VALIDATION = "validation" + TEST = "test" + ALL = "all" + + +class AgenticPromptInjectionBoundaryAttackFamily(Enum): + """Prompt-injection families represented by the controlled pairs.""" + + DIRECT_INSTRUCTION_OVERRIDE = "direct_instruction_override" + SYSTEM_PROMPT_EXTRACTION = "system_prompt_extraction" + ROLEPLAY_JAILBREAK = "roleplay_jailbreak" + AUTHORITY_CLAIM_BYPASS = "authority_claim_bypass" + SENSITIVE_DATA_EXFILTRATION = "sensitive_data_exfiltration" + TOOL_ACTION_ABUSE = "tool_action_abuse" + INDIRECT_CONTENT_INJECTION = "indirect_content_injection" + RAG_CONTEXT_POISONING = "rag_context_poisoning" + MEMORY_CONTEXT_POISONING = "memory_context_poisoning" + OBFUSCATION_ENCODING = "obfuscation_encoding" + MULTI_AGENT_TRUST_SPOOFING = "multi_agent_trust_spoofing" + APPROVAL_WORKFLOW_BYPASS = "approval_workflow_bypass" + + +class AgenticPromptInjectionBoundarySourceContext(Enum): + """Input surfaces through which a record reaches an LLM or agent.""" + + AGENT_MESSAGE = "agent_message" + AGENT_TOOL_REQUEST = "agent_tool_request" + CALENDAR_EVENT = "calendar_event" + CHAT_MESSAGE = "chat_message" + CODE_COMMENT = "code_comment" + CONVERSATION_MEMORY = "conversation_memory" + DIRECT_USER = "direct_user" + DOCUMENT = "document" + EMAIL = "email" + IMAGE_OCR = "image_ocr" + ISSUE_DESCRIPTION = "issue_description" + MCP_TOOL_RESPONSE = "mcp_tool_response" + PDF_DOCUMENT = "pdf_document" + RETRIEVED_DOCUMENT = "retrieved_document" + RUNBOOK = "runbook" + SPREADSHEET_CELL = "spreadsheet_cell" + SUPPORT_TICKET = "support_ticket" + TOOL_OUTPUT = "tool_output" + WEB_PAGE = "web_page" + + +class _AgenticPromptInjectionBoundaryDataset(_RemoteDatasetLoader): + """ + Loader for Agentic Prompt-Injection Boundary Pairs. + + The dataset contains 1,200 English records arranged as 600 controlled + benign/attack pairs. Each pair holds the workflow, asset, role, tool, and + topic constant while changing whether the request crosses an instruction, + authorization, confidentiality, trust, tool-use, or approval boundary. + The 50 scenarios are isolated across train, validation, and test splits. + + Attack records are returned by default for red-team use. Set ``label`` to + ``ALL`` to reconstruct the controlled pairs or ``BENIGN`` to load the hard + negative surface. Family filters use ``pair_family`` rather than + ``attack_family`` so benign rows remain paired with their corresponding + attacks. + + Reference: [@deniz2026agenticboundarypairs]. + License: CC BY 4.0. + """ + + HF_DATASET_NAME: str = "3nesdeniz/agentic-prompt-injection-boundary-pairs" + harm_categories: list[str] = ["benign_boundary", "prompt_injection"] + modalities: tuple[Modality, ...] = (Modality.TEXT,) + size: str = "large" # 1,200 records + tags: frozenset[str] = frozenset({"safety", "agent_security", "prompt_injection", "refusal", "synthetic"}) + + def __init__( + self, + *, + label: AgenticPromptInjectionBoundaryLabel = AgenticPromptInjectionBoundaryLabel.ATTACK, + split: AgenticPromptInjectionBoundarySplit = AgenticPromptInjectionBoundarySplit.ALL, + attack_families: list[AgenticPromptInjectionBoundaryAttackFamily] | None = None, + source_contexts: list[AgenticPromptInjectionBoundarySourceContext] | None = None, + ) -> None: + """ + Initialize the Agentic Prompt-Injection Boundary Pairs loader. + + Args: + label: Security decision label to load. Defaults to attack records. + split: Published split to load, or all splits. Defaults to all. + attack_families: Optional non-empty list of paired attack families. + source_contexts: Optional non-empty list of input surfaces. + + Raises: + ValueError: If an enum value is invalid or a filter list is empty. + """ + self._validate_enum(label, AgenticPromptInjectionBoundaryLabel, "label") + self._validate_enum(split, AgenticPromptInjectionBoundarySplit, "split") + + if attack_families is not None: + if not attack_families: + raise ValueError("`attack_families` must be a non-empty list (pass None to include all families)") + self._validate_enums( + attack_families, + AgenticPromptInjectionBoundaryAttackFamily, + "attack_family", + ) + + if source_contexts is not None: + if not source_contexts: + raise ValueError("`source_contexts` must be a non-empty list (pass None to include all contexts)") + self._validate_enums( + source_contexts, + AgenticPromptInjectionBoundarySourceContext, + "source_context", + ) + + self._label = label + self._split = split + self._attack_families = {family.value for family in attack_families} if attack_families else None + self._source_contexts = {context.value for context in source_contexts} if source_contexts else None + + @property + @override + def dataset_name(self) -> str: + """The registered dataset name.""" + return "agentic_prompt_injection_boundary_pairs" + + @override + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: + """ + Fetch the selected records from Hugging Face. + + Args: + cache: Whether Hugging Face may reuse its local cache. + + Returns: + SeedDataset: Literal prompts with pair and boundary metadata. + + Raises: + ValueError: If a row is malformed, split metadata conflicts with + the requested split, or filters produce no records. + """ + split_names = ( + ["train", "validation", "test"] + if self._split is AgenticPromptInjectionBoundarySplit.ALL + else [self._split.value] + ) + + rows: list[dict[str, Any]] = [] + for split_name in split_names: + split_rows = await self._fetch_from_huggingface_async( + dataset_name=self.HF_DATASET_NAME, + config="default", + split=split_name, + cache=cache, + ) + rows.extend(dict(row) for row in split_rows) + + required_keys = { + "id", + "pair_id", + "scenario_id", + "text", + "label", + "category", + "attack_family", + "pair_family", + "source_context", + "risk_domain", + "target_boundary", + "expected_action", + "source_type", + "language", + "split", + } + + source_url = f"https://huggingface.co/datasets/{self.HF_DATASET_NAME}" + authors = ["Enes Deniz"] + groups = ["AltaySec"] + seeds: list[SeedUnion] = [] + + for row in rows: + missing = required_keys - row.keys() + if missing: + raise ValueError(f"Missing keys in Agentic Boundary Pairs entry: {', '.join(sorted(missing))}") + + row_split = str(row["split"]) + if self._split is not AgenticPromptInjectionBoundarySplit.ALL and row_split != self._split.value: + raise ValueError( + f"Dataset row split '{row_split}' does not match requested split '{self._split.value}'" + ) + + row_label = int(row["label"]) + if row_label not in {0, 1}: + raise ValueError(f"Invalid label in Agentic Boundary Pairs entry: {row_label}") + if self._label is not AgenticPromptInjectionBoundaryLabel.ALL and row_label != self._label.value: + continue + + pair_family = str(row["pair_family"]) + if self._attack_families and pair_family not in self._attack_families: + continue + + source_context = str(row["source_context"]) + if self._source_contexts and source_context not in self._source_contexts: + continue + + category = str(row["category"]) + description_prefix = "Prompt-injection" if row_label == 1 else "Paired legitimate" + description = ( + f"{description_prefix} boundary case for {pair_family.replace('_', ' ')} " + f"delivered through {source_context.replace('_', ' ')}." + ) + + metadata = {key: row[key] for key in sorted(required_keys - {"text"})} + seeds.append( + SeedPrompt( + value=str(row["text"]), + name=str(row["id"]), + data_type="text", + dataset_name=self.dataset_name, + harm_categories=[category], + description=description, + source=source_url, + authors=authors, + groups=groups, + metadata=metadata, + prompt_group_alias=str(row["pair_id"]), + ) + ) + + if not seeds: + raise ValueError("SeedDataset cannot be empty. Check your filter criteria.") + + logger.info(f"Loaded {len(seeds)} records from {self.dataset_name}") + return SeedDataset(seeds=seeds, dataset_name=self.dataset_name) diff --git a/tests/unit/datasets/test_agentic_prompt_injection_boundary_dataset.py b/tests/unit/datasets/test_agentic_prompt_injection_boundary_dataset.py new file mode 100644 index 0000000000..d5d25cf3dc --- /dev/null +++ b/tests/unit/datasets/test_agentic_prompt_injection_boundary_dataset.py @@ -0,0 +1,284 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from collections.abc import Callable +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest + +from pyrit.datasets.seed_datasets.remote.agentic_prompt_injection_boundary_dataset import ( + AgenticPromptInjectionBoundaryAttackFamily, + AgenticPromptInjectionBoundaryLabel, + AgenticPromptInjectionBoundarySourceContext, + AgenticPromptInjectionBoundarySplit, + _AgenticPromptInjectionBoundaryDataset, +) +from pyrit.models import SeedDataset, SeedPrompt + + +def _row( + *, + identifier: str, + pair_id: str, + scenario_id: str, + label: int, + split: str, + pair_family: str = "tool_action_abuse", + source_context: str = "agent_tool_request", +) -> dict[str, Any]: + category = "prompt_injection" if label == 1 else "benign_boundary" + return { + "id": identifier, + "pair_id": pair_id, + "scenario_id": scenario_id, + "text": f"A sufficiently detailed {'attack' if label else 'legitimate'} workflow request for {pair_id}.", + "label": label, + "category": category, + "attack_family": pair_family if label == 1 else "none", + "pair_family": pair_family, + "source_context": source_context, + "risk_domain": "finance", + "target_boundary": "tool_permissions", + "expected_action": "block_or_review" if label == 1 else "allow", + "source_type": "synthetic_curated", + "language": "en", + "split": split, + } + + +@pytest.fixture +def mock_rows_by_split() -> dict[str, list[dict[str, Any]]]: + return { + "train": [ + _row(identifier="apibp_0001_a", pair_id="pair_0001", scenario_id="scn_001", label=1, split="train"), + _row(identifier="apibp_0001_b", pair_id="pair_0001", scenario_id="scn_001", label=0, split="train"), + ], + "validation": [ + _row( + identifier="apibp_0002_a", + pair_id="pair_0002", + scenario_id="scn_002", + label=1, + split="validation", + pair_family="direct_instruction_override", + source_context="direct_user", + ), + _row( + identifier="apibp_0002_b", + pair_id="pair_0002", + scenario_id="scn_002", + label=0, + split="validation", + pair_family="direct_instruction_override", + source_context="direct_user", + ), + ], + "test": [ + _row(identifier="apibp_0003_a", pair_id="pair_0003", scenario_id="scn_003", label=1, split="test"), + _row(identifier="apibp_0003_b", pair_id="pair_0003", scenario_id="scn_003", label=0, split="test"), + _row( + identifier="apibp_0004_a", + pair_id="pair_0004", + scenario_id="scn_004", + label=1, + split="test", + pair_family="indirect_content_injection", + source_context="retrieved_document", + ), + _row( + identifier="apibp_0004_b", + pair_id="pair_0004", + scenario_id="scn_004", + label=0, + split="test", + pair_family="indirect_content_injection", + source_context="retrieved_document", + ), + ], + } + + +def _fetch_side_effect( + rows_by_split: dict[str, list[dict[str, Any]]], +) -> Callable[..., Any]: + async def _fetch(**kwargs: Any) -> list[dict[str, Any]]: + return rows_by_split[str(kwargs["split"])] + + return _fetch + + +def test_dataset_name() -> None: + loader = _AgenticPromptInjectionBoundaryDataset() + assert loader.dataset_name == "agentic_prompt_injection_boundary_pairs" + + +async def test_default_loads_attacks_from_all_splits( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset() + mock_fetch = AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)) + + with patch.object(loader, "_fetch_from_huggingface_async", new=mock_fetch): + dataset = await loader.fetch_dataset_async(cache=False) + + assert isinstance(dataset, SeedDataset) + assert len(dataset.seeds) == 4 + assert all(isinstance(seed, SeedPrompt) for seed in dataset.seeds) + assert all(seed.metadata["label"] == 1 for seed in dataset.seeds) + assert {call.kwargs["split"] for call in mock_fetch.await_args_list} == {"train", "validation", "test"} + assert all(call.kwargs["config"] == "default" for call in mock_fetch.await_args_list) + assert all(call.kwargs["cache"] is False for call in mock_fetch.await_args_list) + + first = dataset.seeds[0] + assert first.name == "apibp_0001_a" + assert first.prompt_group_alias == "pair_0001" + assert first.metadata["scenario_id"] == "scn_001" + assert first.metadata["target_boundary"] == "tool_permissions" + assert first.source == ("https://huggingface.co/datasets/3nesdeniz/agentic-prompt-injection-boundary-pairs") + + +async def test_label_all_reconstructs_pairs( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset( + label=AgenticPromptInjectionBoundaryLabel.ALL, + split=AgenticPromptInjectionBoundarySplit.TEST, + ) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)), + ): + dataset = await loader.fetch_dataset_async() + + assert len(dataset.seeds) == 4 + assert {seed.metadata["label"] for seed in dataset.seeds} == {0, 1} + assert {seed.prompt_group_alias for seed in dataset.seeds} == {"pair_0003", "pair_0004"} + + +async def test_benign_filter_loads_hard_negative_surface( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset( + label=AgenticPromptInjectionBoundaryLabel.BENIGN, + split=AgenticPromptInjectionBoundarySplit.TRAIN, + ) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)), + ): + dataset = await loader.fetch_dataset_async() + + assert len(dataset.seeds) == 1 + assert dataset.seeds[0].harm_categories == ["benign_boundary"] + assert dataset.seeds[0].metadata["expected_action"] == "allow" + + +async def test_family_filter_uses_pair_family_and_keeps_both_sides( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset( + label=AgenticPromptInjectionBoundaryLabel.ALL, + split=AgenticPromptInjectionBoundarySplit.TEST, + attack_families=[AgenticPromptInjectionBoundaryAttackFamily.INDIRECT_CONTENT_INJECTION], + ) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)), + ): + dataset = await loader.fetch_dataset_async() + + assert len(dataset.seeds) == 2 + assert {seed.metadata["label"] for seed in dataset.seeds} == {0, 1} + assert all(seed.metadata["pair_family"] == "indirect_content_injection" for seed in dataset.seeds) + + +async def test_source_context_filter( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset( + label=AgenticPromptInjectionBoundaryLabel.ALL, + split=AgenticPromptInjectionBoundarySplit.TEST, + source_contexts=[AgenticPromptInjectionBoundarySourceContext.RETRIEVED_DOCUMENT], + ) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)), + ): + dataset = await loader.fetch_dataset_async() + + assert len(dataset.seeds) == 2 + assert all(seed.metadata["source_context"] == "retrieved_document" for seed in dataset.seeds) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"label": "attack"}, "Expected AgenticPromptInjectionBoundaryLabel"), + ({"split": "test"}, "Expected AgenticPromptInjectionBoundarySplit"), + ({"attack_families": []}, "non-empty"), + ({"source_contexts": []}, "non-empty"), + ({"attack_families": ["tool_action_abuse"]}, "Expected AgenticPromptInjectionBoundaryAttackFamily"), + ({"source_contexts": ["direct_user"]}, "Expected AgenticPromptInjectionBoundarySourceContext"), + ], +) +def test_invalid_filters_raise(kwargs: dict[str, Any], message: str) -> None: + with pytest.raises(ValueError, match=message): + _AgenticPromptInjectionBoundaryDataset(**kwargs) + + +async def test_empty_after_filter_raises( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + loader = _AgenticPromptInjectionBoundaryDataset( + split=AgenticPromptInjectionBoundarySplit.TRAIN, + attack_families=[AgenticPromptInjectionBoundaryAttackFamily.APPROVAL_WORKFLOW_BYPASS], + ) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(side_effect=_fetch_side_effect(mock_rows_by_split)), + ): + with pytest.raises(ValueError, match="SeedDataset cannot be empty"): + await loader.fetch_dataset_async() + + +async def test_missing_required_key_raises( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + malformed = dict(mock_rows_by_split["train"][0]) + malformed.pop("pair_id") + loader = _AgenticPromptInjectionBoundaryDataset(split=AgenticPromptInjectionBoundarySplit.TRAIN) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(return_value=[malformed]), + ): + with pytest.raises(ValueError, match="pair_id"): + await loader.fetch_dataset_async() + + +async def test_conflicting_split_metadata_raises( + mock_rows_by_split: dict[str, list[dict[str, Any]]], +) -> None: + mismatched = dict(mock_rows_by_split["test"][0]) + mismatched["split"] = "train" + loader = _AgenticPromptInjectionBoundaryDataset(split=AgenticPromptInjectionBoundarySplit.TEST) + + with patch.object( + loader, + "_fetch_from_huggingface_async", + new=AsyncMock(return_value=[mismatched]), + ): + with pytest.raises(ValueError, match="does not match requested split"): + await loader.fetch_dataset_async()