# How Needle Handles Tool Selection vs Argument Grounding During Training

> Discover how Needle unifies tool selection and argument grounding in a single fine tuning stage. Learn how it emits structured tool calls with JSON-compliant arguments.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-18

---

**Needle uses a single‑stage fine‑tuning process that jointly learns tool selection and argument grounding by training the model to emit structured `<tool_call>` blocks containing both the tool name and a JSON schema‑compliant argument payload.**

The `cactus-compute/needle` repository implements a unified training strategy that avoids separate pipelines for choosing tools and filling their parameters. Instead, the model learns to generate a single structured markup block that encodes the decision and the data simultaneously, leveraging JSON schemas produced by the tool decorator system.

## The Single‑Stage Training Architecture

Traditional tool‑learning frameworks often split training into distinct phases: first teaching the model to select the correct tool, then fine‑tuning it to generate valid arguments. Needle rejects this separation in favor of a **coupled generative task**. During training, the model sees complete conversation examples where the assistant response includes a `<tool_call>` XML block. Inside this block, the `"name"` field represents the tool selection decision, while the `"arguments"` object represents the grounded parameter values. By predicting both within the same token sequence, the model learns the statistical relationship between user intent, tool identity, and required arguments.

## Tool Selection: Encoding the Decision

Tool selection is represented as the **first token generation decision** within the `<tool_call>` block. When the model determines a tool is required, it begins emitting JSON after the opening `<tool_call>` token, with the first key‑value pair being `"name": "tool_name"`.

The available tool candidates are defined by the **`@tool` decorator** in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py). This decorator inspects function signatures and builds a JSON schema via the `build_schema` method:

```python
from needle import tool, Field

@tool
def send_email(to: str, subject: str, body: str = Field(default="")) -> str:
    """Send an email."""
    pass

```

The resulting schema—including types, required fields, defaults, and validation constraints—is attached to training examples as the `"parameters"` field. This schema acts as the ground truth against which the model learns to select the correct tool name based on the user query context.

## Argument Grounding: Structured Schema Completion

Once the tool name is selected, the model must populate the JSON object with values derived from the user query. This **argument grounding** task is treated as constrained JSON generation. The model learns to emit a complete JSON object that validates against the schema produced by `build_schema`, filling in the `"arguments"` field with properly typed values.

During dataset creation in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), the `generate_examples` function ensures that every synthetic training example includes:
- The `"tools"` array containing all available tool schemas
- A `"query"` that naturally implies a specific tool invocation
- The correct `<tool_call>` markup with a fully populated JSON payload

The model is trained to maximize the likelihood of generating this structured output, effectively learning both *which* tool to invoke and *how* to parameterize it based on the schema constraints.

## Training Data Pipeline

The training pipeline constructs `(query, answer, tools)` triples using three specific operations:

1. **Schema Collection**: All decorated tools are inspected via `build_schema` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), producing a JSON schema for each function signature.
2. **Synthetic Generation**: The `generate_examples` function in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) prompts a teacher LLM to write user queries that would invoke specific tools, then appends the corresponding `<tool_call>` fragments.
3. **Dataset Assembly**: Rows are written as JSONL lines containing the `"tools"` array (selection candidates) and the target query‑answer pair, enabling the model to learn the mapping from natural language to structured calls.

Each training datum explicitly couples the candidate tool definitions with the execution target, preventing the model from learning selection and grounding as isolated tasks.

## From Training to Inference

At inference time, the **tokenizer** ([`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py)) monitors the output stream for the special tokens `<tool_call>` and `</tool_call>`. Once a complete JSON payload is emitted, the runtime in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) extracts the `"name"` field (tool selection) and the `"arguments"` field (argument grounding), then dispatches the call to the appropriate Python function.

This architecture ensures that the model's behavior during inference mirrors exactly what it learned during training: a unified generation process where tool selection and argument grounding are two aspects of the same structured output.

## Summary

- Needle trains tool selection and argument grounding as a **single generative task** within `<tool_call>` blocks, rather than using separate training phases.
- The **`@tool` decorator** in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) generates JSON schemas via `build_schema` that define valid tool names and argument structures.
- **Training data** is synthesized by `generate_examples` in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), creating examples that pair user queries with complete tool‑call markup.
- The model learns to predict the tool name (selection) and populate JSON arguments (grounding) simultaneously, conditioned on the provided schemas.
- At inference, the **tokenizer** and **runtime** parse the same markup structure, extracting both decisions from the unified output format.

## Frequently Asked Questions

### How does Needle differ from multi‑stage tool learning approaches?

Multi‑stage approaches typically train one model (or adapter) to classify which tool to use, then a second model to fill arguments. Needle performs both operations in a single forward pass by generating a structured `<tool_call>` block, allowing the model to condition argument values on the specific tool selected and vice versa.

### What role does the @tool decorator play in training?

The **`@tool` decorator** in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) introspects function signatures to build JSON schemas using `build_schema`. These schemas are embedded in the training examples as the `"tools"` field, providing the model with the structural constraints it must learn to satisfy during argument grounding.

### How does the tokenizer enforce structured output during inference?

The **tokenizer** defined in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) recognizes the special delimiter tokens `<tool_call>` and `</tool_call>`. While it does not hard‑constrain the JSON syntax during generation (depending on the base model), it parses the emitted text to identify when a complete tool call block has been produced, allowing the runtime to validate and execute the extracted tool name and arguments.

### Where are the training examples generated in the codebase?

Training examples are generated by the **`generate_examples` function** in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py). This function creates synthetic user queries and corresponding assistant responses that include valid `<tool_call>` markup, ensuring the dataset explicitly links natural language to structured tool invocations.