# How to Define Tools for Needle 2 Using Python Decorators

> Learn to define tools for Needle 2 with Python decorators. Automatically generate JSON schemas from type annotations and docstrings for efficient model creation.

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

---

**Use the `@needle.tool` decorator on type-annotated Python functions to register them as callable tools; Needle 2 automatically extracts JSON schemas from signatures and docstrings to constrain model generation.**

The cactus-compute/needle repository provides a lightweight inference engine where external capabilities are exposed through Python functions. When you define tools for Needle 2 using Python decorators, the framework inspects your function signatures to generate constrained grammars that guide the model's output toward valid executable calls.

## How the @needle.tool Decorator Works

The `@needle.tool` decorator, implemented in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), transforms ordinary Python functions into registered tools by extracting two critical pieces of metadata:

1. **Argument types** – Derived directly from type-annotated parameters to build the JSON schema for tool inputs.
2. **Tool description** – Taken from the function’s docstring to provide the natural-language description shown to the model.

When you apply the decorator, the function is entered into a global registry that the `Needle` constructor reads when building its tool catalogue. This registry is exposed publicly through [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), allowing you to import `needle.tool` directly.

## Creating Your First Tool with Type Annotations

Define a tool by decorating any function with standard Python type hints. The following example from the `cactus-compute/needle` source demonstrates a simple weather lookup:

```python
import needle

@needle.tool
def get_weather(city: str):
    """Get the current weather for a city."""
    # In a real implementation you might call an API here.

    return {"city": city, "temp_c": 27, "sky": "clear"}

```

Supported primitive types include `str`, `int`, `float`, `bool`, and `list`. The decorator automatically converts these annotations into a JSON schema that constrains what arguments the model can generate.

## Handling Complex Inputs with Pydantic Models

For nested or structured data, Needle 2 accepts `pydantic.BaseModel` classes as type annotations. When the model predicts a call, Needle instantiates the Pydantic model from the JSON arguments before invoking your function:

```python
from pydantic import BaseModel
import needle

class Invoice(BaseModel):
    vendor: str
    total: float
    due_date: str

@needle.tool
def create_invoice(data: Invoice):
    """Create an invoice from structured data."""
    # Process the invoice here…

    return {"status": "created", "id": 1234}

```

This pattern ensures complex objects are validated and parsed according to your schema definitions before execution.

## Tool Execution Flow During Inference

During inference, the execution pipeline follows a strict contract to bridge model predictions with Python execution. As documented in [`README.md`](https://github.com/cactus-compute/needle/blob/main/README.md) (lines 35‑48) and [`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md), the process works as follows:

- The model predicts a tool call based on the conversation context.
- Needle retrieves the matching Python function by name from the tool catalogue.
- Predicted JSON arguments are parsed and converted to the annotated Python types.
- The function executes, and its return value (typically a JSON-serializable dict) is placed under the `results` key in the final response.

Crucially, the global registry supplies the JSON schema to the byte-level grammar compiler, ensuring the model can only emit well-formed calls that respect the declared types and constraints.

## Initializing the Needle Agent with Your Tools

After defining your tools, pass them explicitly to the `Needle` constructor or allow auto-discovery from the global registry:

```python

# Create agent with explicit tool list

agent = needle.Needle(tools=[get_weather, create_invoice])

# Run a query; the model decides which tool to call

response = agent.run("What's the weather like in Lagos right now?")

# Access execution results

print(response["results"])

# → [{'city': 'Lagos', 'temp_c': 27, 'sky': 'clear'}]

```

The `Needle` class, exposed in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), manages the tool catalogue and orchestrates the inference loop.

## Summary

- Decorate functions with `@needle.tool` to register them in the global registry defined in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py).
- Use Python type annotations to automatically generate JSON schemas for tool inputs; include docstrings to provide natural-language descriptions.
- Leverage `pydantic.BaseModel` for complex, nested argument structures.
- Pass tool functions to the `Needle` constructor to make them available during inference.
- Needle 2 uses a byte-level grammar compiler to ensure model outputs conform strictly to your declared types before execution.

## Frequently Asked Questions

### What file implements the @needle.tool decorator?

The decorator logic resides in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), while the public API is exposed through [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) for direct import as `needle.tool`.

### Does Needle 2 support Pydantic models for tool arguments?

Yes, you can annotate parameters with `pydantic.BaseModel` subclasses for nested object inputs. Needle automatically instantiates these models from the JSON arguments predicted by the model before calling your function.

### How does Needle 2 ensure the model generates valid tool call syntax?

The framework feeds your function signatures into a byte-level grammar compiler that constrains the model's token generation during inference. This guarantees that the model produces only well-formed JSON that respects your declared types and required parameters.

### Where does Needle 2 store tool execution results?

Function return values appear under the `results` key in the response dictionary returned by `agent.run()`, allowing you to access structured output from multiple tool calls within a single inference session.