💻 computer science

Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents

The paper introduces "Squeez," a task-conditioned tool-output pruning method that uses a fine-tuned Qwen 3.5 2B model to identify and retain only the most relevant evidence from long tool observations, achieving high recall and F1 scores while removing 92% of input tokens and significantly outperforming larger zero-shot models and heuristic baselines.

Original authors: Ádám Kovács

Published 2026-04-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ádám Kovács

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a complex crime, but instead of getting a neat summary of the evidence, you are handed a giant, chaotic warehouse filled with millions of documents, audio tapes, and video clips. Every time you ask a question, you have to read everything in the warehouse before you can make your next move. It's slow, exhausting, and you often miss the one crucial clue buried under a mountain of irrelevant paperwork.

This is exactly the problem coding agents (AI programmers) face today. When they try to fix a bug, they get flooded with massive logs, error messages, and code files. Most of this data is "noise," but the AI has to process it all, which is incredibly inefficient.

Enter Squeez: The "Smart Shredder" for AI.

The Core Idea: Finding the Needle in the Haystack

The authors of this paper realized that AI doesn't need to read the whole haystack to find the needle. They created a system called Squeez that acts like a super-smart filter.

Here is how it works in simple terms:

The Input: You give the AI a specific question (e.g., "Why did the login fail?") and a massive wall of text (the tool output, like a 500-line error log).
The Magic: Instead of reading the whole thing, Squeez instantly scans the text and cuts out only the tiny, specific lines that actually answer the question.
The Output: It hands the AI just those few lines, throwing away 92% of the junk.

The "Chef" Analogy

Think of a coding agent as a chef trying to make a soup.

The Old Way: The chef is given a truckload of groceries (the tool output). To find the salt, the chef has to open every single bag, taste every vegetable, and read every label before they can start cooking. It takes forever.
The Squeez Way: A sous-chef (Squeez) stands by the truck. The head chef asks, "Where is the salt?" The sous-chef instantly grabs only the salt shaker and hands it over, leaving the truck full of carrots and potatoes behind. The head chef can now cook immediately.

How They Built It

To teach their AI this skill, the researchers didn't just guess. They built a massive training gym with over 11,000 practice scenarios.

They used real-world coding problems (from a famous dataset called SWE-bench) and created fake ones for different programming languages.
They taught a small, efficient AI model (a 2-billion parameter version of Qwen) to look at a question and a messy log, then highlight the exact lines of code or error messages that matter.
They tested it against "big brain" models (like a 35-billion parameter AI) and simple rules (like "just read the first 10 lines").

The Results: Small is Beautiful

The results were surprising and impressive:

The Small Model Won: The tiny, specialized AI (Squeez) was much better at finding the right clues than the massive, expensive 35-billion parameter AI.
Speed and Efficiency: Squeez removed 92% of the unnecessary text. It found the right evidence 86% of the time (Recall), beating the giant models by a wide margin.
Why the Big Models Failed: The big, "zero-shot" models (models not specifically trained for this) were too chatty. They often tried to summarize or rewrite the text, or they got distracted by similar-looking but wrong information. Squeez, however, learned to be a strict "extractor"—it just grabs the exact lines needed, verbatim, without rewriting them.

Why This Matters

In the world of AI coding, attention is a limited resource. If an AI spends too much time reading irrelevant logs, it gets slower, costs more money to run, and might even get confused.

Squeez is like a laser-guided spotlight. It shines only on the part of the story that matters, letting the AI agent move faster and smarter. It proves that you don't need a "super-genius" AI to filter data; you just need a small, well-trained specialist who knows exactly what to look for.

In a nutshell: Squeez is the ultimate "Ctrl+F" for AI agents, but instead of just finding a word, it understands the context, cuts out the garbage, and hands the agent the exact evidence it needs to solve the problem next.

1. Problem Definition

Coding agents (e.g., those interacting with file systems, version control, and build tools) frequently process long streams of tool observations (logs, stack traces, file contents, API responses). However, only a tiny fraction of tokens in these observations is relevant to the agent's next decision.

The Bottleneck: Agents often re-read entire outputs, leading to high token consumption, latency, and context window bloat.
The Gap: Existing solutions like generic prompt compression (e.g., LLMLingua) or document-level retrieval compression are designed for prose or retrieved passages. They fail to handle mixed-format tool outputs where code, logs, shell traces, and structured metadata are interleaved.
The Task: The paper defines Task-Conditioned Tool-Output Pruning. Given a focused query ( $q$ ) and a single raw tool observation ( $o$ ), the system must return the smallest verbatim evidence block (contiguous line spans) necessary for the next reasoning step. The goal is not to solve the bug, but to extract the minimal evidence required to do so.

2. Methodology

A. Dataset Construction

The authors created a comprehensive benchmark of 11,477 examples derived from two sources:

SWE-bench Derived (9,205 examples): Real repository snapshots from SWE-bench. The authors executed 14 different tool types (e.g., read file, grep, git log, pytest, pip install) to generate raw observations.
Synthetic Multi-Ecosystem (2,272 examples): Generated using openai/gpt-oss-120b to cover languages and ecosystems not dominant in SWE-bench (TypeScript, Go, Rust, Java, Docker, Terraform, Kubernetes). This includes 575 explicit negative examples (mismatched query/output pairs where the correct answer is an empty selection).

Annotation Pipeline:

A "teacher" model generates focused, tool-aware queries and selects the minimal contiguous spans (gold standard) that answer the query.
The final dataset includes span coordinates mapped back to the original raw text, ensuring the output is a verbatim subset.
Split: 10,508 training, 240 development, and a manually curated 618-example test set (excluding duplicates, trivial outputs, or broad spans).

B. Model Architecture

Base Model: Qwen 3.5 2B. The authors chose a small model to demonstrate that a specialized, lightweight policy is more efficient than large zero-shot models for this specific extraction task.
Training: Fine-tuned using LoRA (Low-Rank Adaptation) for 3 epochs.
Input/Output: The model takes the query and raw tool output and generates the extracted text wrapped in <relevant lines> tags.
Deployment: The LoRA adapter is merged into the base model for serving via vLLM or as a CLI filter, allowing seamless integration into existing agent stacks (e.g., Codex, Claude Code).

C. Evaluation Metrics

Recall: The primary metric, measuring the fraction of gold lines retrieved.
F1 Score: A tolerant line-matching score (fuzzy substring similarity > 0.5) to account for minor formatting differences in generative output.
Compression: The fraction of input tokens removed.
Baselines:
- Generative: Qwen 3.5 35B A3B, Kimi K2, and the un-tuned Qwen 3.5 2B base.
- Heuristic: BM25, First-N, Last-N, and Random (all keeping ~10% of input lines).

3. Key Contributions

Task Formulation: Defined a specific, narrow benchmark for verbatim evidence extraction from mixed-format tool outputs, distinct from generic summarization or bug solving.
Dataset Release: Released a large-scale, multi-ecosystem dataset (11k+ examples) covering 27 tool types, including critical negative examples for training "do nothing" behavior.
Efficient Model: Demonstrated that a 2B parameter model fine-tuned for this specific task significantly outperforms much larger zero-shot models (18x larger) and heuristic baselines.
Practical Integration: Provided a deployable solution (CLI and vLLM) that acts as a lightweight pre-processing step, requiring minimal changes to existing agent loops.

4. Results

On the held-out test set (618 examples), Squeez-2B achieved state-of-the-art performance:

Recall: 0.86 (outperforming the 35B zero-shot model by 11 points and the un-tuned 2B base by 33 points).
F1 Score: 0.80 (tolerant) / 0.49 (strict exact match).
Compression: Removed 92% of input tokens while maintaining high recall.
Comparison:
- vs. Zero-Shot 35B: The 35B model struggled with repetitive logs and often selected semantically adjacent but incorrect blocks. Squeez-2B learned specific extraction policies.
- vs. Heuristics: Heuristics (BM25, First-N, etc.) performed poorly (Recall ~0.10–0.22) because relevant evidence can appear anywhere in the output, and relevance is query-dependent, not just lexical.
- Negative Handling: Squeez-2B correctly returned empty output for 80% of negative examples, whereas the 35B model only managed 7% (often hallucinating explanations).

Qualitative Analysis:
The model excels at:

Isolating single hits in structured outputs (e.g., git log, grep).
Extracting compact failure blocks from noisy logs.
Recognizing when no relevant evidence exists (True Negatives).
Residual Errors: Mostly "adjacent over-selection" (e.g., including a nearby unrelated commit or a correct error block with extra context), rather than selecting wholly irrelevant content.

5. Significance and Conclusion

The paper argues that mixed-format tool output is a unique challenge that generic large language models (LLMs) and retrieval heuristics cannot solve efficiently via zero-shot prompting.

Efficiency: A small, task-specific model is superior to massive generalist models for pruning, offering a 92% token reduction with higher accuracy.
Scalability: By reducing the context window requirements for downstream reasoning steps, this approach makes complex coding agents more scalable and cost-effective.
Future Work: The authors note limitations in evaluating full agent trajectories (the current benchmark focuses on single-step pruning) and the varying noise levels across different tool families.

In summary, Squeez provides a practical, high-performance solution for the "needle in a haystack" problem within coding agents, proving that specialized, lightweight supervision is the key to efficient tool-output processing.