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

1. O Problema

Agentes de codificação operam consumindo fluxos contínuos de observações de ferramentas (leitura de arquivos, saídas de grep, rastreamentos de pilha, logs de build, respostas de API, histórico de controle de versão, etc.). Um desafio central de eficiência nesses sistemas é que, embora apenas uma pequena fração dos tokens de cada observação seja relevante para a próxima decisão do agente, o sistema frequentemente precisa reler a saída completa. Isso consome recursos computacionais e janelas de contexto desnecessariamente.

O problema abordado pelo Squeez é uma versão específica e prática desse desafio: dada uma consulta focada (query) e uma única observação de ferramenta bruta, retornar o menor bloco de evidência textual (verbatim) que o agente deve inspecionar a seguir. Diferente de sumarização abstrativa ou compressão de prompts genérica, o objetivo aqui é preservar a evidência original exata para garantir rastreabilidade e evitar alucinações, descartando apenas o irrelevante.

2. Metodologia

Definição da Tarefa

O modelo recebe um par $(q, o)$, onde $q$ é uma consulta curta e condicionada à tarefa (ex: "Encontre o traceback que explica o ImportError") e $o$ é a observação bruta da ferramenta. O modelo deve identificar e extrair os índices de linha contíguos que formam o bloco de evidência mínimo necessário.

Dataset e Benchmark

Os autores criaram um benchmark robusto composto por 11.477 exemplos, divididos em:

• Dados Reais (SWE-bench): 9.205 exemplos derivados de interações reais em repositórios do SWE-bench, executando 14 tipos de ferramentas (leitura de arquivos, git, testes, linters, etc.).
• Dados Sintéticos: 2.039 exemplos positivos e 575 negativos gerados para cobrir ecossistemas além do Python (TypeScript, Go, Rust, Java, Docker, Terraform, etc.), totalizando 27 famílias de ferramentas.
• Conjunto de Teste: Um conjunto de 618 exemplos foi manualmente curado e validado para garantir a qualidade das anotações (span de ouro).

Modelo e Treinamento

• Arquitetura: Utiliza-se o modelo Qwen 3.5 2B (uma variante pequena e eficiente da família Qwen).
• Ajuste Fino (Fine-tuning): O modelo é ajustado com LoRA (Low-Rank Adaptation) por 3 épocas.
• Objetivo de Treino: O modelo é treinado para emitir o texto extraído exato, envolto em tags XML (<relevant lines>), aprendendo uma política de extração supervisionada específica para a tarefa.
• Infraestrutura: Treinado em uma única GPU NVIDIA A100 80GB e servido via vLLM para baixa latência.

Baselines de Comparação

O modelo foi comparado contra:

1. Modelos Generativos Zero-Shot: Qwen 3.5 35B A3B (18x maior), Kimi K2 e a base Qwen 3.5 2B sem ajuste.
2. Heurísticas Simples: BM25, First-N (primeiras N linhas), Last-N (últimas N linhas) e Aleatório, configurados para manter ~10% do input (nível de compressão similar ao do conjunto de ouro).

3. Principais Contribuições

1. Formulação do Problema: Define e formaliza o "pruning de saída de ferramenta condicionado à tarefa" como uma tarefa de benchmark distinta para agentes de codificação.
2. Dataset Abrangente: Liberação de um dataset misto (real e sintético) cobrindo múltiplos ecossistemas e 27 tipos de ferramentas, com anotações de spans fundamentados (grounded spans).
3. Eficiência de Modelo Pequeno: Demonstra que um modelo compacto (2B parâmetros) ajustado especificamente para esta tarefa supera significativamente modelos muito maiores em modo zero-shot e heurísticas tradicionais.
4. Integração Prática: O modelo é projetado para ser inserido em pilhas de agentes existentes (como Codex ou Claude Code) como um filtro pré-processador via CLI ou API, sem alterar a lógica central de planejamento do agente.

4. Resultados

No conjunto de teste mantido (held-out test set de 618 exemplos), o Squeez-2B obteve desempenho superior:

• Recall: 0.86 (superando o Qwen 35B em 11 pontos e a base 2B em 33 pontos).
• F1 Score: 0.80.
• Compressão: Remove 92% dos tokens de entrada, mantendo a precisão.
• Precisão: 0.80, indicando que o modelo não apenas encontra o que precisa, mas evita incluir ruído excessivo.

Análise Comparativa:

• Modelos Grandes Zero-Shot: O Qwen 35B, embora forte, falha ao selecionar blocos semanticamente adjacentes mas incorretos (ex: escolher um commit relacionado mas errado em logs do Git) e tem dificuldade em retornar vazios quando não há evidência (falso positivo).
• Heurísticas: Desempenham muito mal (Recall ~0.22 para BM25), pois a relevância em saídas de ferramentas depende da consulta e do contexto, não apenas da sobreposição lexical ou posição fixa (início/fim).
• Casos de Sucesso: O modelo ajustado destaca-se em isolar hits únicos em saídas estruturadas, extrair blocos compactos de falha em logs ruidosos e, crucialmente, retornar saída vazia corretamente em exemplos negativos (80% de acerto vs 7% do Qwen 35B).

5. Significado e Conclusão

O trabalho demonstra que saídas de ferramentas em formato misto (código, logs, metadados) não são bem tratadas por modelos generativos zero-shot ou heurísticas de recuperação padrão.

A lição central é que a supervisão estreita e específica da tarefa é fundamental. Um modelo pequeno e barato, quando ajustado para uma função específica de extração de evidência, pode substituir com sucesso modelos massivos ou heurísticas ineficientes dentro do loop de um agente de codificação. Isso permite que sistemas de IA operem com maior eficiência de contexto, reduzindo custos e latência, enquanto mantêm a fidelidade às evidências originais necessárias para o raciocínio do agente.

O código, o modelo e o dataset estão disponíveis publicamente sob licença Apache 2.0, facilitando a adoção imediata pela comunidade de pesquisa e desenvolvimento de agentes autônomos.