Data agent that thinks like a scientist in DABStep

NVIDIA KGMON (NeMo Agent Toolkit) Data Explorer reached first place in the DABStep benchmark by using a strategy that separates heavy learning from fast inference. What’s the key? Build reusable tools during a learning phase, then run responses with a small, agile agent that orchestrates those tools.

What problem it solves

Agents that rely on text search fail when the information lives in tables and requires multi-step reasoning. Complex questions about tabular data aren’t fixed by a web snippet. Have you ever seen a model answer one thing well, then get lost when you cross two CSV files and a set of business rules?

This project was built for that: multi-step questions, stateful tools, and strict validation.

Architecture in three phases

The central idea is to split responsibilities: spend compute once to produce robust tools, then use those tools many times efficiently.

• Fase de Learning (aprendizaje): se usa un modelo pesado (por ejemplo Opus 4.5/4.6) en un loop multi-paso con un conjunto completo de herramientas (intérprete Python stateful, bash, detector de estructura de archivos, retriever). El agente resuelve varios casos representativos, valida contra ground truth y sintetiza soluciones en una biblioteca reutilizable helper.py y ejemplos few-shot.

• Fase de Inference (inferencia): aquí entra un modelo pequeño/rápido (por ejemplo Haiku 4.5) en un loop simple. El contexto se poda agresivamente: el agente recibe solo las firmas de las funciones y un prompt optimizado, y llama a helper.py mediante un intérprete Python mínimo. Resultado: latencia y tokens al mínimo.

• Fase de Offline Reflection (revisión offline): un modelo pesado (Opus o Sonnet 4.6) actúa como revisor no supervisado. Ejecuta dos técnicas clave: reflection (auditoría de código y uso de helper.py) y group-consistency (comparar soluciones para grupos de preguntas similares). Las correcciones se retroalimentan offline para afinar el prompt y las funciones sin ralentizar la inferencia en producción.

Dos loops de agente según caso de uso

• Exploratory EDA: ReAct agent + herramientas de manipulación de notebooks. El agente crea, ejecuta celdas y genera visualizaciones; un handler envía las gráficas a un Vision-Language Model para describirlas y proponer mejoras, devolviendo texto enriquecido al agente.

• Tabular QA orientado a reglas: Tool Calling Agent que orquesta un intérprete Python stateful, un retriever y un detector de estructura. Este enfoque es ideal para tareas que exigen pasos deterministas y validación rígida.

Technical insight: why it works

The big idea is recognizing that many complex questions share basic operations. Instead of redoing scripts per task, the agent builds generalizable functions. For example: computing a transaction fee and listing fee IDs share the same initial steps.

The Learning phase tests function versions across multiple tasks; when a version fails on another case, the agent refines the function until it generalizes.

Moving expensive checks offline lets inference be 30x faster without losing rigor. The system keeps quality thanks to offline reflection and by injecting learned patterns into the light agent’s prompt.

Results and comparison

Key points:

• Dominance on “Hard” tasks: 89.95 versus 66.93 from the baseline that solves everything from scratch.
• 30x speedup in time per task (20 seconds vs 10 minutes) thanks to the helper.py library and the light model at inference.
• Much more compact final code, which reduces human review and format-related errors.

How to replicate or apply it in your projects (brief technical guide)

1. Reúne un dev split representativo para la fase de aprendizaje. Asegúrate de tener ground truth y ejemplos que expongan variaciones operativas.
2. Ejecuta un agente pesado con acceso a un intérprete Python stateful y herramientas de inspección de archivos. Resuelve lotes de tareas y guarda scripts por tarea.
3. Automatiza la síntesis: prueba variaciones de funciones contra múltiples tareas y valida. Extrae las funciones que generalizan bien en helper.py.
4. Diseña un prompt conciso que incluya firmas de funciones, few-shot examples y reglas de formato de salida. Usa un modelo más pequeño para inferencia y un intérprete más limitado.
5. Corre revisiones offline periódicas con un modelo pesado para reflection y group-consistency. Alimenta las mejoras resultantes al prompt y a helper.py.

Practical tips: prioritize automated unit tests for helper.py functions; lengthen the learning cycle when the task distribution is very heterogeneous; use a VLM to convert charts into text when you need to explain visualizations to non-technical users.

Impact and limitations

This approach shows that with architectural design you can make small models solve hard tabular reasoning problems while keeping speed and token economy. It’s especially useful in domains with sensitive or specialized tabular data, where the web helps little.

Limitations to consider: the initial cost of the Learning phase can be high; final quality depends on how representative the dev split is; and it relies on business rules not changing radically between learning and inference.

This proposal isn’t magic: it’s agent engineering and capitalizing on code generalization. If you work with financial databases, transaction logs, or complex catalogs, this architecture can noticeably speed up your analysis pipelines.

Original source

https://huggingface.co/blog/nvidia/nemo-agent-toolkit-data-explorer-dabstep-1st-place