# How Needle 2 Ensures Valid Tool Calls with Grammar-Constrained Decoding

> Needle 2 ensures valid tool calls using grammar constrained decoding and JSON schema extraction. Learn how it restricts model vocabulary for type correct generation.

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

---

**Needle 2 guarantees syntactically valid and type-correct tool calls by coupling JSON schema extraction with a grammar-based decoder that restricts the model's vocabulary during generation.**

Needle is an open-source agent framework developed by `cactus-compute/needle`. By embedding tool schemas directly into the decoding process, it eliminates post-generation validation and repair. This article breaks down the three-step pipeline that powers Needle 2's grammar-constrained decoding, from schema extraction to constrained sampling.

## Schema Extraction via the `@tool` Decorator

Every tool in Needle 2 is annotated with the `@tool` decorator. This decorator automatically generates a JSON schema that describes the tool's name, description, and expected parameter types.

In [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), the `build_schema` function inspects Python type hints, `Field` objects, and Pydantic models to produce the schema. The decorator then attaches this metadata directly to the function object under `_needle_tool`.

```python

# Define a tool with type hints – the schema is auto‑generated

from needle import tool, Field

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

# The decorator stores the schema in the function object

print(send_email._needle_tool)

# -> {

#      "name": "send_email",

#      "description": "Send an email.",

#      "parameters": {

#          "type": "object",

#          "properties": {

#              "to": {"type": "string"},

#              "subject": {"type": "string"},

#              "body": {"type": "string", "default": ""}

#          },

#          "required": ["to", "subject"]

#      }

#    }

```

This schema becomes the foundation for the grammar that will constrain the decoder.

## Grammar Construction from JSON Schemas

Once the schemas are collected, Needle 2 constructs a formal grammar for the decoding step. The grammar is anchored by special tokens defined in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) at lines 18–21: `<tool_call>` and `</tool_call>`. These tokens delimit the exact segment of the model output that must conform to the tool schema.

The decoder's vocabulary is restricted so that only strings matching the constructed grammar can be produced inside this segment. This prevents malformed JSON, incorrect argument names, and type mismatches before they are generated.

## Runtime Grammar-Constrained Decoding

The model's sampler invokes the decoder with the generated grammar context active. Because the decoder only yields token sequences that satisfy the grammar, every tool invocation is guaranteed to be well-formed and type-correct according to the declared schema.

As implemented in `cactus-compute/needle`, the orchestration logic in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) ties the schema generation step directly to the sampling loop. There is no need for external JSON repair or retry logic because the constraint is enforced at the token level.

```python

# Running a query – Needle builds a grammar from the schema

agent = Needle(tools=[send_email])
response = agent.complete("Please email the report to alice@example.com")

# The model’s decoder can only emit a <tool_call> block that matches the schema,

# e.g.:

# <tool_call>

# {"name":"send_email","arguments":{"to":"alice@example.com","subject":"Report","body":""}}

# </tool_call>

```

## Bypassing Constraints with `--no-grammar`

By default, grammar-constrained decoding is active for all tool calls. However, Needle 2 exposes a CLI flag to bypass this restriction. In [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) at line 118, the `--no-grammar` flag turns off the grammar constraint, allowing the model to generate unconstrained text within tool-call segments.

Use this override with caution, as it removes the guarantee of syntactically valid output.

## Summary

- **Schema extraction** happens automatically through the `@tool` decorator and `build_schema` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), capturing type hints and Pydantic models as JSON schemas.
- **Grammar construction** embeds these schemas into a formal grammar bounded by `<tool_call>` and `</tool_call>` tokens from [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py).
- **Constrained decoding** restricts the sampler's vocabulary so only grammar-compliant token sequences are emitted, ensuring valid tool calls without post-processing.
- The `--no-grammar` flag in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) (line 118) allows users to disable the safety net when needed.

## Frequently Asked Questions

### What is grammar-constrained decoding in Needle 2?

Grammar-constrained decoding is a mechanism where the model's token sampler is restricted by a formal grammar derived from each tool's JSON schema. According to the `cactus-compute/needle` source code, this ensures that any text emitted between `<tool_call>` and `</tool_call>` tokens is syntactically valid JSON that matches the tool's declared argument types.

### How does Needle 2 build a JSON schema for a tool?

Needle 2 uses the `@tool` decorator to call `build_schema` inside [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py). This function inspects Python type hints, default values, `Field` objects, and Pydantic models to generate a JSON schema, which is then stored on the function as `_needle_tool`.

### Can I disable grammar-constrained decoding?

Yes. The [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) file exposes a `--no-grammar` flag at line 118 that disables the grammar constraint. When this flag is used, the model generates tool calls without vocabulary restrictions, though this may produce malformed output.

### Where are the special tool-call tokens defined?

The special `<tool_call>` and `</tool_call>` tokens are defined in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) between lines 18 and 21. These tokens delimit the section of model output that is governed by the tool schema grammar.