# How to Implement Grammar-Constrained Decoding with Byte-Level Schemas in Needle

> Implement grammar-constrained decoding in Needle using byte-level schemas. Needle 2 ensures syntactically valid outputs matching your JSON structure, guaranteeing correct data formats.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-14

---

**Needle 2 automatically compiles JSON schemas into byte-level grammars that constrain the decoder at every token, guaranteeing syntactically valid outputs that match your declared structure.**

The **grammar-constrained decoding** system in Needle transforms your tool schemas and Pydantic models into deterministic finite automata (DFAs) that operate at the byte level. This constraint is applied during generation—not after—so the model can only emit tokens that preserve valid JSON structure.

---

## How Needle Builds Byte-Level Grammars from Schemas

Needle's pipeline follows a five-step process from Python declaration to enforced constraint:

| Step | Action | Source Location |
|------|--------|-----------------|
| **1. Schema creation** | `@needle.tool` decorator or `needle.extract` with Pydantic models generate JSON-Schema dictionaries via `build_schema` and `pydantic_schema` | [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) — [`build_schema`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L110), [`pydantic_schema`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L149) |
| **2. Schema registration** | Schemas passed to `Needle(tools=[...])` or `--tools` CLI flag are stored and prepared | [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) — [`Needle.__init__`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L90) |
| **3. Grammar compilation** | JSON-Schemas transformed into byte-level DFAs by the C++ runtime | Triggered via [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) → `decode_cfg` → engine initialization |
| **4. Constrained generation** | Decoder masks tokens that would violate the DFA after each emitted token | [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) — [`generate_cached`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py#L54), [`batched_generate`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py#L50) |
| **5. Optional override** | `--no-grammar` flag disables constraint for free-form output | [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) — [`--no-grammar`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py#L118) |

The **byte-level** nature of this grammar is critical: the engine can reject individual bytes that would lead to invalid JSON, not just complete tokens. This enables precise structural guarantees even with variable-length token encodings.

---

## Implementing Grammar-Constrained Decoding in Python

### Simple Tool with Automatic Schema Extraction

```python
import needle

@needle.tool
def add(a: int, b: int):
    """Add two numbers."""
    return a + b

agent = needle.Needle(tools=[add])
result = agent.run("Add 4 and 7")["results"]
print(result)  # → [{'name': 'add', 'arguments': {'a': 4, 'b': 7}}]

```

**What the grammar enforces:**

- The `build_schema` function in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) generates `{"name":"add","parameters":{"a":{"type":"integer"},"b":{"type":"integer"}}}`.
- The compiled **byte-level grammar** permits only token sequences representing valid JSON objects containing integer-valued keys `"a"` and `"b"`.
- During `generate_cached`, any token that would produce a floating-point number, string, or missing key is **masked out** before softmax.

---

### Pydantic Model Extraction with Schema Constraints

```python
from pydantic import BaseModel
import needle

class Contact(BaseModel):
    name: str
    email: str

text = "John Doe can be reached at john@doe.com"
contact = needle.extract(text, Contact)
print(contact)  # → name='John Doe' email='john@doe.com'

```

**What the grammar enforces:**

- `pydantic_schema(Contact)` creates a JSON-Schema requiring `name` and `email` as strings.
- Needle creates an internal tool containing only this schema, compiles the **constrained grammar**, and executes single-turn generation.
- The `_is_pydantic_model(schema)` check in the extraction path ensures the returned dict is instantiated as `Contact(**result)`.

---

## Using Grammar-Constrained Decoding from the CLI

### With External Schema Files

```bash

# tools.json contains an array of JSON-Schema objects

needle playground --tools tools.json

```

The CLI loader passes the schema array to `NeedleEngine`, triggering the same compilation pipeline as the Python API.

---

### Disabling Grammar Constraints

```bash
needle playground --no-grammar

```

This skips the grammar mask entirely. The model generates unrestricted text. Use this only for debugging or when **grammar-constrained decoding** would prevent valid outputs for non-structured prompts.

---

## Core Source Files for Grammar Implementation

| File | Purpose |
|------|---------|
| [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) | Schema building: `build_schema()` for functions, `pydantic_schema()` for models |
| [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) | `Needle` class, `extract()` helper, schema-to-engine registration |
| [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) | Decoding loops `generate_cached()` and `batched_generate()` with grammar masking |
| [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) | Command-line interface, `--tools` and `--no-grammar` flags |
| [`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md) | Public API documentation for tool creation and extraction |

---

## Performance and Compatibility Considerations

- **Compilation cost:** Grammar DFAs are built once at engine initialization, not per-generation. This overhead is negligible for multi-turn conversations.
- **Token efficiency:** Byte-level constraints can reject partial tokens, reducing the search space earlier than character-level approaches.
- **Schema complexity:** Deeply nested schemas or extensive `anyOf` unions increase DFA state count. Monitor memory if defining 100+ tool schemas.
- **Compatibility:** Grammar compilation occurs in the C++ runtime. Python-side schema validation uses standard JSON-Schema Draft 7.

---

## Summary

- **Declare schemas** with `@needle.tool` or Pydantic models—[`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) handles JSON-Schema generation.
- **Register schemas** via `Needle(tools=[...])` or `--tools`—[`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) passes them to the engine.
- **Enforce constraints** automatically—the C++ runtime compiles byte-level grammars that `generate_cached` applies at every token.
- **Control behavior** with `--no-grammar` when free-form output is required.

---

## Frequently Asked Questions

### What is byte-level grammar compilation in Needle?

Byte-level grammar compilation transforms JSON schemas into deterministic finite automata that validate individual bytes rather than complete tokens. This allows the decoder to reject partial tokens that would lead to invalid JSON structure, enabling stricter constraints than character or token-level approaches.

### How does Needle guarantee valid JSON outputs?

Needle masks the model's output distribution at every generation step using the compiled DFA. If a token's byte sequence would transition the automaton to an invalid state, its probability is set to zero. This **grammar-constrained decoding** continues until the JSON object is complete and valid.

### Can I use custom JSON schemas not derived from Python functions?

Yes. Pass a JSON file containing an array of JSON-Schema objects to the `--tools` CLI flag, or construct the schema dictionary manually and pass it to `Needle(tools=[{"name": "custom", "parameters": {...}}])`. The compilation pipeline accepts any valid JSON-Schema Draft 7.

### What happens if my schema is too complex for real-time constraint?

Extremely complex schemas with deep nesting or many `anyOf` branches increase DFA state count and memory usage. Needle's C++ runtime handles most practical schemas efficiently, but if you encounter performance issues, simplify the schema or split functionality across multiple focused tools.