# How Needle 2 Enforces Byte-Level Grammar Constrained Decoding for Tool Calls

> Needle 2 uses special tokens, tokenizer mappings, and logit masking for byte-level grammar constrained decoding, ensuring tool calls match JSON schemas before token emission.

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

---

**Needle 2 enforces byte-level grammar constrained decoding by combining special token markers, tokenizer mappings, and real-time logit masking to ensure every tool call conforms to a JSON schema before the token is even emitted.**

Needle 2 is an open-source LLM framework that implements rigorous **byte-level grammar constrained decoding** to guarantee syntactically valid tool calls. By baking grammar rules directly into the tokenizer and decoding pipeline, the system prevents malformed JSON or stray characters from ever appearing in tool-related outputs. This article examines the three-component architecture and specific source implementations that make this enforcement possible.

## The Three Pillars of Grammar-Constrained Decoding

Needle 2's constrained decoding relies on three integrated components working in concert: special token markers that delimit tool sections, a tokenizer that maps these markers to unique IDs, and a masking system that filters logits at generation time.

### Special Token Markers

The framework defines explicit XML-like delimiters to bracket tool-related content within the model's text stream. In [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py), the following constants establish the boundary markers at lines 16-21:

- `TOOLS_START` / `TOOLS_END` (`<tools>` / `</tools>`)
- `TOOL_CALL_START` / `TOOL_CALL_END` (`<tool_call>` / `</tool_call>`)
- `TOOL_RESULT_START` / `TOOL_RESULT_END` (`<tool_result>` / `</tool_result>`)

These markers serve as anchors for the grammar state machine. When the model generates `TOOL_CALL_START`, the decoder enters a constrained mode where subsequent tokens must conform to the tool's JSON schema until the matching `TOOL_CALL_END` token appears.

### Tokenizer Mapping

The tokenizer registers each marker as a literal byte-sequence token within the model's vocabulary. Located in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py), this mapping assigns unique token IDs to the delimiter constants defined in [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py). By treating these markers as ordinary tokens rather than post-processing annotations, the model can emit them naturally during generation while the decoder tracks grammatical state based on their presence in the token stream.

### Constrained Decoding Logic

The core enforcement mechanism resides in the generation functions `generate_cached` and `batched_generate` within [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py). During each forward pass, the model produces logits for the next token, but before sampling occurs, the framework applies a **grammar mask** derived from the current parsing state.

The masking process works as follows:

1. The decoder maintains a finite-state automaton tracking position within the tool call grammar (e.g., inside JSON keys, values, or closing braces).
2. Based on this state and the tool's schema, it constructs a mask identifying which tokens would violate the grammar.
3. Invalid tokens have their logits set to a very low value (effectively negative infinity), preventing the sampler from selecting them.
4. Valid tokens retain their original probabilities, allowing temperature-based sampling to proceed within grammatical constraints.

Because this mask is recomputed after every single token, the constraint operates at the byte level, ensuring that partial JSON strings or malformed structures cannot emerge even with stochastic sampling enabled.

## The Constrained Decoding Pipeline

The complete workflow for tool call generation follows a deterministic path from prompt construction to constrained output:

1. **Prompt Construction** – The `build_prompt` function in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) detects registered tools and injects the `<tools>` block into the prompt context.
2. **Token Generation** – The model invokes `forward_cached` followed by `generate_cached` (or `batched_generate` for batch processing).
3. **Grammar Mask Calculation** – Upon detecting `TOOL_CALL_START`, the decoder initializes a state machine based on the tool's schema (attached via `_needle_tool` by the `@tool` decorator).
4. **Masked Sampling** – Before each softmax operation, the grammar mask filters the logits, allowing only tokens that advance valid JSON syntax for the specific tool signature.
5. **Completion Detection** – Once the model emits `TOOL_CALL_END`, the decoder exits constrained mode and resumes standard generation for the remainder of the response.

This pipelined approach guarantees that content between `<tool_call>` and `</tool_call>` tags is always syntactically valid according to the schema defined in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py).

## Implementing Schema-Aware Tool Definitions

The grammar constraints originate from tool definitions decorated with `@tool`. When a function is decorated, the framework attaches a `_needle_tool` attribute containing the JSON schema derived from type hints and docstrings.

```python
from needle import tool, Field

@tool
def send_email(to: str, subject: str, body: str):
    """Send an email to the specified address."""
    pass

# Initialize agent with constrained decoding support

agent = Needle(tools=[send_email])

# Generation automatically triggers grammar constraints

response = agent("Please email the team the daily report.")

```

During the call to `agent()`, the decoder extracts the schema from `send_email._needle_tool` and builds a deterministic state machine for the allowed byte sequences. As the model generates tokens, it must produce valid JSON keys matching the schema parameters (`to`, `subject`, `body`) and properly formatted string values, with the grammar mask rejecting any deviation.

For example, the constrained output will follow this exact structure:

```json
<tool_call>
{"name":"send_email","arguments":{"to":"team@example.com","subject":"Daily Report","body":"..."}}
</tool_call>

```

The decoder ensures that the JSON is well-formed, keys match the schema, and the structure is properly closed before `</tool_call>` terminates the constrained region.

## Summary

- **Special token markers** (`<tool_call>`, `</tool_call>`, etc.) defined in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) delimit grammar-constrained regions.
- **Tokenizer integration** in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) maps these markers to unique token IDs, enabling the model to emit them as standard tokens.
- **Real-time logit masking** in `generate_cached` and `batched_generate` prevents invalid tokens from being sampled by setting their probabilities to negligible values.
- **Schema extraction** via the `@tool` decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) provides the grammatical rules enforced during generation.
- **Byte-level enforcement** occurs at every generation step, ensuring valid JSON even with temperature > 0 or other stochastic sampling methods.

## Frequently Asked Questions

### What is byte-level grammar constrained decoding?

Byte-level grammar constrained decoding is a generation technique that filters model outputs at the token prediction stage to ensure they conform to a formal grammar or schema. Unlike post-processing validation that rejects invalid outputs after generation, this approach masks invalid tokens during the softmax step, guaranteeing that every emitted byte contributes to a syntactically valid structure.

### How does Needle 2 prevent invalid JSON in tool calls?

Needle 2 prevents invalid JSON through real-time logit masking in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py). When the decoder enters a tool call region (triggered by the `TOOL_CALL_START` token), it maintains a state machine representing valid JSON transitions for that specific tool's schema. Before each token is sampled, the mask eliminates logits that would produce invalid JSON syntax, ensuring only grammatically correct tokens can be selected.

### Where are the special tool tokens defined in the Needle codebase?

The special tool tokens are defined as Python constants in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) at lines 16-21. The constants `TOOLS_START`, `TOOLS_END`, `TOOL_CALL_START`, `TOOL_CALL_END`, `TOOL_RESULT_START`, and `TOOL_RESULT_END` correspond to the string markers `<tools>`, `</tools>`, `<tool_call>`, `</tool_call>`, `<tool_result>`, and `</tool_result>` respectively. These are then registered as actual tokens in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py).

### Can grammar constraints work with temperature-based sampling?

Yes, grammar constraints in Needle 2 are fully compatible with temperature-based sampling and other stochastic methods. Because the constraints are applied as a mask to the logits before the softmax operation, they only restrict the *set* of possible tokens to those that maintain grammatical validity. Within that valid set, temperature scaling, top-k, and top-p sampling operate normally, allowing for creative variations in phrasing or string values while preserving structural correctness.