Tool Output Extraction Dataset

Paper | Code

Training data for squeez — a small model that prunes verbose coding agent tool output to only the evidence the agent needs next.

Task

Task-conditioned context pruning of a single tool observation for coding agents.

Given a focused extraction query and one verbose tool output, return the smallest verbatim evidence block(s) the agent should read next.

The model copies lines from the tool output — it never rewrites, summarizes, or invents content. Every line in the target exists verbatim in the source.

Dataset Summary

Data sources:

Why This Dataset Exists

LLM coding agents waste 80–95% of context tokens on irrelevant tool output. When an agent reads a 500-line file to find one function, or runs pytest to find a failing test, most of the output is noise. Squeez trains small models to compress this output before it enters the agent's context window.

Why v3?

The previous dataset (v2) had three problems:

1. Line-number leakage. The synthetic data generation pipeline showed numbered output to the teacher LLM, and the teacher leaked those numbers into the target annotations. 87% of synthetic response lines contained line number prefixes (2: npm WARN...) that did not exist in the raw tool output. This trained the model to hallucinate formatting.

2. No canonical truth format. The dataset stored XML-wrapped ChatML prompts as the ground truth. This tangled the benchmark representation with one specific model's training format, making cross-model evaluation fragile.

3. Task drift. The labeling task drifted from "context pruning for an agent" toward "recover one arbitrary teacher-selected subset." Queries were either too vague (full issue descriptions) or too literal (grep-able string searches), neither matching the actual product use case.

v3 fixes all three by introducing a canonical span-based representation, regenerating all labels with a focus on verbatim extraction, and manually curating the test set.

How The Data Was Created

Step 1: Generate Tool Calls (SWE-bench)

Starting from 2,294 instances in the SWE-bench test split, we simulated what a coding agent would do for each issue: read source files, grep for symbols, run tests, check git history, install packages. We generated 3–7 tool calls per instance, weighted toward common agent actions (read_file 28%, grep 18%, test_output 8%, etc.).

Step 2: Execute on Real Repos

Every tool call was executed against the actual repository checked out at the base commit (before the fix). git grep ran real searches, pytest ran real tests, pip install -e . did real installs. This produces authentic tool output with real file paths, error messages, formatting, and noise — things an LLM cannot reliably generate from scratch.

Step 3: Generate Focused Queries

For each (issue, tool_output) pair, a teacher LLM generated a focused extraction query — a short, concrete request for evidence rather than a full issue description or a literal grep pattern.

Good queries: "Find the traceback block that explains the import failure", "Find the diff hunk that changes CSV parsing" Bad queries: "Fix the bug", "Find all lines containing 'raise AttributeError'"

Step 4: Label Gold Spans

A teacher LLM selected gold spans — contiguous blocks of lines in the raw tool output that answer the extraction query. The teacher saw numbered output as a reference interface, but the canonical target maps line numbers back to the raw unnumbered text. Every target line must exist verbatim in the source.

The canonical representation is a list of {start_line, end_line} spans over the raw tool output. XML wrappers, ChatML formatting, and line number prefixes only appear in derived training files, never in the ground truth.

Step 5: Synthetic Multi-Ecosystem Data

SWE-bench only covers Python repositories, but real coding agents work across all ecosystems. We generated synthetic data for 15 tool types that SWE-bench cannot provide: npm_install, npm_build, tsc, eslint, cargo_build, go_build, mvn_gradle, make_cmake, docker_build, docker_logs, terraform, kubectl, pip_install, mypy_pyright, and curl.

Two-pass generation:

1. An LLM generates a realistic task description and tool output in XML markers. Each tool type has a config with scenarios (e.g., "peer dependency conflict", "missing native module") and seed examples guiding realistic formatting.
2. Given the numbered output, a second LLM call selects relevant line numbers as JSON. These are mapped back to raw unnumbered source lines, validated (must exist verbatim, must be non-empty, reasonable compression ratio), and stored as canonical spans.

Step 6: Hard Negatives

To teach the model when to output nothing, 575 samples use intentionally mismatched task+output pairs — e.g., a React authentication task paired with a Rust borrow-checker error. The correct answer is an empty extraction.

Step 7: Quality Filtering and Assembly

Samples with empty spans (where the teacher found nothing relevant) were capped at 10% per tool type. Overly broad annotations (>60% of lines selected) were reviewed. Train/dev/test splits were assembled:

• SWE-bench: split by repository (test: xarray, flask; dev: requests; train: all others)
• Synthetic: split per tool type (10% test, 5% dev, 85% train)
• Negatives: capped at ~10% of positives per tool type in test

Step 8: Test Set Curation

The held-out test set was manually reviewed and 111 samples were excluded:

• Near-duplicate np.unicode_ errors (63): The xarray repo on NumPy 2.0 produces the same AttributeError: np.unicode_ was removed on every import xarray. These identical errors across 65 different xarray instances were deduplicated to 2 representative samples.
• Trivial tiny outputs (39): Samples with 1–2 line output (e.g., lint "All checks passed!", "Python 3.12.9", single-line curl errors). Nothing to filter — not a meaningful benchmark.
• Overly broad spans (5): Samples selecting >50% of a large output, or spanning the entire top half of a file.
• Wrong annotations (4): Mislabeled tool types, spans pointing to wrong content, or vague queries without task context.

The exclusion list is tracked in test_exclusions.json with per-sample reasons.

Formats

The dataset ships in three parallel formats, all derived from the same canonical spans:

Canonical (canonical_train/dev/test.jsonl)

The source of truth. Model-agnostic, no XML, no formatting artifacts.

{
  "instance_id": "django__django-11270",
  "source": "swe",
  "tool_type": "read_file",
  "query": "Find the code block that validates the referer in CsrfViewMiddleware",
  "background_task": "Fix CSRF validation bug when referer URL contains port number...",
  "tool_output": "raw output exactly as shown to the agent",
  "gold_spans": [
    {"start_line": 41, "end_line": 52}
  ],
  "is_irrelevant": false,
  "command": "django/middleware/csrf.py"
}

• gold_spans reference 1-indexed line numbers in tool_output
• is_irrelevant: true means no lines are relevant (hard negative)
• query is the focused extraction request; background_task is the full issue for provenance

Generative / Qwen (train/dev/test.jsonl)

ChatML-formatted for SFT training with Qwen or similar models.

{
  "prompt": "<|im_start|>system
You prune verbose tool output...<|im_end|>
<|im_start|>user
<query>
Find the code block...
</query>
<tool_output>
1: class CsrfViewMiddleware:
2:     def _check_referer(self, request):
...
</tool_output><|im_end|>
<|im_start|>assistant
",
  "response": "<relevant_lines>
41: referer = request.META.get('HTTP_REFERER')
42: if referer is None:
...
</relevant_lines>",
  "metadata": {
    "instance_id": "django__django-11270",
    "tool_type": "read_file",
    "source": "swe",
    "num_total_lines": 84,
    "num_relevant_lines": 12,
    "compression_ratio": 0.857
  }
}

Encoder (encoder_train/dev/test.jsonl)

For token/line classification models (mmBERT, etc.).

Tool Types

27 tool types across multiple ecosystems:

Splits

SWE-bench data is split by repository (zero instance overlap):

• Test: pydata/xarray, pallets/flask
• Dev: psf/requests
• Train: all others (django, sympy, scikit-learn, sphinx, matplotlib, pytest, astropy, pylint, seaborn)

Synthetic data is split per tool type: 10% test, 5% dev, 85% train. Hard negatives are capped at ~10% per tool type in test.

Usage

from datasets import load_dataset

ds = load_dataset("KRLabsOrg/tool-output-extraction-swebench")

# Generative training splits
print(ds)
# DatasetDict({
#     train: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 10508})
#     dev: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 240})
#     test: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 618})
# })

Citation

@misc{kovács2026squeeztaskconditionedtooloutputpruning,
      title={Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents}, 
      author={Ádám Kovács},
      year={2026},
      eprint={2604.04979},
      archivePrefix={arXiv},
      primaryClass={cs.SE},
      url={https://arxiv.org/abs/2604.04979}, 
}