Skip to content

Datasets API

pikit.datasets

Standard benchmark datasets for prompt-injection evaluation.

This module provides a thin loader layer that reads TOML test-case files from the datasets/ directory, converts each case into an :class:~pikit.config.ExperimentConfig, and runs it through the existing :class:~pikit.matrix.MatrixRunner — so dataset evaluation reuses the exact same engine as custom matrix experiments.

Two built-in datasets ship with pikit:

  • direct_injection — attacker payload delivered via the user message (datasets/direct_injection.toml).
  • indirect_injection — attacker payload hidden in tool-returned data (datasets/indirect_injection.toml).

Each TOML file has a [meta] section and an array of [[cases]]. Every case is a flat dict whose keys map directly to ExperimentConfig fields, plus an id and description for reference.

Example

from pikit.datasets import list_datasets, load_dataset, run_dataset list_datasets() ['direct_injection', 'indirect_injection'] results = run_dataset("direct_injection", target_spec="mock") len(results) > 0 True

DatasetCase dataclass 

DatasetCase(id: str, description: str, config: ExperimentConfig)

One test case from a dataset TOML file.

Attributes

id: Short identifier (e.g. "di-001"). description: Human-readable description of what the case tests. config: A fully-built :class:ExperimentConfig ready to run.

Dataset dataclass 

Dataset(name: str, description: str, reference: str, cases: List[DatasetCase] = list(), path: str = '')

A loaded benchmark dataset.

Attributes

name: Dataset name from the [meta] section. description: Dataset description from [meta]. reference: Academic reference string from [meta]. cases: List of :class:DatasetCase objects. path: Filesystem path to the source TOML file.

list_datasets 

list_datasets() -> List[str]

Return the names of all available built-in datasets.

Examples

list_datasets() ['direct_injection', 'indirect_injection']

load_dataset 

load_dataset(name: str) -> Dataset

Load a built-in dataset by name.

Parameters

name: Dataset name (e.g. "direct_injection").

Returns

Dataset The loaded dataset with all cases parsed.

Raises

KeyError If no dataset with that name exists.

run_dataset 

run_dataset(name: str, *, target_spec: Optional[str] = None, judge_type: Optional[str] = None, temperature: Optional[float] = None, repeats: Optional[int] = None, attacks: Optional[List[str]] = None, verbose: bool = False) -> List[ExperimentResult]

Run all cases in a dataset and collect results.

Each case is run independently via :class:MatrixRunner. Results from all cases are concatenated into a single list, with the case_id and case_description stored in each result's reason prefix for traceability.

Parameters

name: Dataset name. target_spec: Override the target spec for all cases (e.g. "openai:gpt-4o-mini"). Defaults to each case's own spec (usually "mock"). judge_type: Override the judge type for all cases. temperature: Override sampling temperature. repeats: Override number of repeats per case. attacks: Override the attack methods for all cases (e.g. ["naive", "escape"]). When provided, every case will use these attacks regardless of what the dataset TOML specifies (or omits). When None, each case falls back to its own attacks field, or ["naive"] if the dataset omits the field entirely (naive = raw payload without any attack-method wrapping). verbose: Print progress to stderr.

Returns

List[ExperimentResult] Results from all cases, in dataset order.