# How to Define Custom Tools for Needle 2 Using the @needle.tool Decorator

> Learn to define custom tools for Needle 2 with the @needle.tool decorator. Safely validate and execute Python code using type hints and docstrings.

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

---

**The `@needle.tool` decorator converts any Python function into a JSON-Schema-backed tool by inspecting type hints, docstrings, and `typing.Annotated` fields, enabling the Needle 2 agent to validate and execute your code safely.**

The `cactus-compute/needle` repository provides a lightweight framework for building AI agents that interact with user-defined Python functions. Defining custom tools for Needle 2 requires only a single decorator that automatically generates JSON Schema descriptions from your function signatures. This schema-driven approach ensures that language models produce valid arguments that conform to your specified types and constraints.

## The `@needle.tool` Mechanism

Internally, the decorator operates as a schema factory that bridges Python functions and the Needle inference engine. When you apply `@needle.tool` to a function, the implementation in **[`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)** executes a three-step process:

1. The `tool()` function (lines 64-67) intercepts the function definition
2. It attaches a schema payload via `fn._needle_tool = build_schema(fn)` (lines 64-66)
3. The `build_schema()` function (lines 11-43) introspects the signature, extracting type hints through `_json_type`, default values, and docstring descriptions via `_parse_doc`

The resulting schema complies with the JSON-Schema specification and is stored on the function object. When you instantiate a `Needle` agent with a list of tools, the constructor reads these schemas to construct the system prompt, as documented in **[`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md)** (lines 9-12).

## Structuring Tool Functions

Needle 2 generates schemas automatically from four function characteristics. You do not need to write JSON manually; instead, follow these Python conventions:

- **Type hints** – Parameters with types like `int`, `str`, or `Literal` map directly to JSON Schema types (e.g., `int → integer`, `str → string`)
- **Default values** – Any parameter with a default becomes optional in the schema
- **Docstrings** – The function’s docstring becomes the tool’s description; an `Args:` block provides per-parameter documentation
- **Return values** – The return type is not enforced by the schema but is serialized and returned to the model under the `"results"` key

## Adding Validation Constraints

For fine-grained control over parameter constraints, combine **`typing.Annotated`** with **`needle.Field`**. This pattern allows you to specify ranges, regex patterns, enums, and string lengths directly in the type signature.

The `needle.Field` helper supports validation keywords such as:

- **`gt`** and **`le`** – Greater than and less than or equal for numeric bounds
- **`pattern`** – Regular expression validation for strings
- **`max_length`** – Maximum string length
- **`description`** – Human-readable explanation of the constraint

When the model outputs a tool call, Needle validates the JSON against these constraints before executing your function, ensuring that `Annotated[float, needle.Field(gt=0)]` truly receives positive numbers.

## Complete Working Example

The following example demonstrates three tiers of tool definition: basic docstring-driven schemas, `Literal` enums, and strict validation with `Annotated` fields.

```python
import needle
from typing import Annotated, Literal

# Simple example – only the signature and docstring are needed.

@needle.tool
def get_weather(city: str):
    """Get the current weather for a city."""
    # In a real scenario you would query a weather service here.

    return {"city": city, "temp_c": 22, "sky": "partly cloudy"}

# Using `Literal` for a fixed set of options.

@needle.tool
def set_thermostat(
    temperature: int,
    mode: Literal["heat", "cool", "auto"] = "auto"
):
    """
    Set the thermostat.

    Args:
        temperature: Target temperature in Celsius.
        mode: Heating strategy to use.
    """
    return {"temperature": temperature, "mode": mode}

# Adding validation constraints with `needle.Field`.

@needle.tool
def send_money(
    amount: Annotated[float, needle.Field(gt=0, le=10000,
                                         description="USD, up to 10 000")],
    to:     Annotated[str,   needle.Field(pattern=r"^@[a-z0-9_]+$",
                                         description="recipient handle")],
    memo:   Annotated[str,   needle.Field(max_length=80)] = ""
):
    """Send money to a handle."""
    return {"sent": amount, "to": to, "memo": memo}

# Create an agent with the custom tools.

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

# Run a query – the model will pick the appropriate tool.

response = agent.run("What’s the weather in Berlin?")
print(response["results"])

# → [{'city': 'Berlin', 'temp_c': 22, 'sky': 'partly cloudy'}]

```

In this workflow, `Needle.run()` handles the orchestration: it prompts the model with the tool schemas, validates the model’s JSON output against those schemas, executes the corresponding Python function, and feeds the return value back to the model for the final response.

## Summary

- **`@needle.tool`** in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) attaches a JSON Schema to functions via `fn._needle_tool = build_schema(fn)`
- **Schema generation** is automatic, deriving types from signatures, constraints from `typing.Annotated[...]`, and descriptions from docstrings
- **Validation** occurs at inference time; the model must produce JSON that satisfies all `needle.Field` constraints (ranges, patterns, lengths) before your function executes
- **Results** are returned under the `"results"` key in the final response dictionary

## Frequently Asked Questions

### How does Needle 2 discover which tools are available?

Needle 2 discovers tools through the schema attached by `@needle.tool`. When you pass a list of decorated functions to `needle.Needle(tools=[...])`, the constructor extracts the `_needle_tool` attribute (set by the decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) lines 64-66) and includes these schemas in the system prompt sent to the language model.

### What Python types can I use in tool function signatures?

You can use standard JSON-serializable types: `str`, `int`, `float`, `bool`, and `list`. For enumerated values, use `typing.Literal` (e.g., `Literal["heat", "cool"]`). Complex objects are not supported directly; return dictionaries or lists that can serialize to JSON for the final response under the `"results"` key.

### How do I make a tool parameter optional?

Declare the parameter with a default value in the function signature. The `build_schema()` function in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) inspects defaults and marks those arguments as optional in the generated JSON Schema, allowing the model to omit them when calling the tool.

### Can I add regex validation to string parameters?

Yes. Use `typing.Annotated` combined with `needle.Field(pattern=r"your_regex")`. The pattern is embedded in the JSON Schema, and the inference engine validates the model's output against it before executing your function, ensuring the string matches the specified regex.