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

> Learn how to define tools for Needle 2 with the @tool decorator. Automatically register Python functions as LLM accessible tools by inspecting type hints and docstrings.

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

---

**The `@tool` decorator automatically registers Python functions as LLM-accessible tools by inspecting type hints and docstrings, storing callable references in the global `TOOL_REGISTRY` defined in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py).**

Needle 2 is an open-source framework that bridges large language models (LLMs) with arbitrary Python functionality. By using the `@tool` decorator, developers can expose any callable—synchronous or asynchronous—to the LLM runtime without manual API scaffolding or explicit registration calls.

## Understanding the @tool Decorator Architecture

The `@tool` decorator operates entirely at import time, transforming standard Python functions into structured tool definitions that the Needle runtime can discover and invoke.

### Global Registration in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)

The core implementation resides in [[`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py). When you apply `@tool` to a function, the decorator immediately adds the callable to the **`TOOL_REGISTRY`**, a global dictionary keyed by the function name. This registry acts as the source of truth for the [**[`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py)**](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py) module, which resolves LLM tool requests to actual Python implementations at runtime.

Because registration happens during module import, simply importing a file containing decorated functions makes those tools available to the agent. No additional registration boilerplate is required.

### Automatic Schema Generation

Under the hood, the decorator performs **signature inspection** on the wrapped function. It extracts:

- **Type hints** from parameters and return values, converting them into JSON Schema-compatible type definitions that the LLM uses to construct valid arguments.
- **Docstrings**, which are parsed and stored as the tool's description. This text appears in the system prompt, helping the model determine when to invoke the tool.
- **Return type metadata**, ensuring outputs are automatically serialized (e.g., dataclasses convert to dictionaries, tensors to lists) before being sent back to the LLM.

## Defining Tools with the @tool Decorator

To expose a function, import the decorator from `needle.agent.tools` and apply it to any top-level callable.

### Basic Function Decoration

Define a simple tool by adding the decorator and providing type hints and a docstring:

```python
from needle.agent.tools import tool

@tool
def add(a: int, b: int) -> int:
    """Add two integers together and return the sum."""
    return a + b

```

When this module loads, `add` is registered under the key `"add"` in `TOOL_REGISTRY` with a JSON schema describing the integer parameters `a` and `b`.

### Async Tool Support

The decorator handles asynchronous functions without additional configuration. Needle 2 automatically detects `async def` signatures and awaits them during execution:

```python
import httpx
from needle.agent.tools import tool

@tool
async def fetch_url(url: str) -> str:
    """Fetch the raw HTML content from a given URL."""
    async with httpx.AsyncClient() as client:
        response = await client.get(url)
        return response.text

```

The runtime in [`fetch.py`](https://github.com/cactus-compute/needle/blob/main/fetch.py) detects the coroutine and manages the event loop, ensuring the LLM receives the serialized string result.

### Overriding Metadata

For cases where the Python function name or docstring is not ideal for the LLM, the decorator accepts optional arguments to customize the tool's public interface:

- **`name`**: Overrides the registry key.
- **`description`**: Provides a custom description instead of the docstring.
- **`examples`**: Supplies sample input/output pairs for better LLM grounding.

```python
from datetime import datetime
from needle.agent.tools import tool

@tool(
    name="current_time",
    description="Return the current UTC timestamp in ISO format.",
    examples=[{"input": {}, "output": "2024-01-15T12:00:00Z"}]
)
def get_time() -> str:
    return datetime.utcnow().isoformat() + "Z"

```

These metadata fields are merged into the tool definition before insertion into `TOOL_REGISTRY`.

## Runtime Resolution and Invocation

When an LLM decides to call a tool, it sends a request containing the tool name and arguments. The [**[`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py)**](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py) module handles the lookup by querying `TOOL_REGISTRY` for the matching callable. It validates the incoming arguments against the pre-computed JSON schema, executes the function (awaiting it if necessary), and returns the serialized result to the model.

This architecture keeps the tool-definition surface minimal while ensuring type safety and clear separation between definition (in [`tools.py`](https://github.com/cactus-compute/needle/blob/main/tools.py)) and execution (in [`fetch.py`](https://github.com/cactus-compute/needle/blob/main/fetch.py)).

## Testing and Validation

You can verify that your tools are correctly registered by inspecting `TOOL_REGISTRY` directly in your test suite, as demonstrated in the repository's [**[`tests/test_tools.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_tools.py)**](https://github.com/cactus-compute/needle/blob/main/tests/test_tools.py):

```python
from needle.agent.tools import TOOL_REGISTRY

def test_add_tool_registration():
    assert "add" in TOOL_REGISTRY
    result = TOOL_REGISTRY["add"](a=3, b=4)
    assert result == 7

def test_custom_name_registration():
    assert "current_time" in TOOL_REGISTRY
    assert "get_time" not in TOOL_REGISTRY

```

This approach allows unit testing of tool logic without launching the full Needle runtime.

## Summary

- The **`@tool`** decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) registers functions into the global **`TOOL_REGISTRY`** during module import.
- It automatically converts **type hints** to JSON schemas and extracts **docstrings** for LLM-facing descriptions.
- Both synchronous and **asynchronous** functions are supported without additional configuration.
- Optional parameters (`name`, `description`, `examples`) allow customization of the tool's public interface.
- Runtime invocation is handled by [`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py), which resolves tool names from the registry, validates arguments, and serializes outputs.

## Frequently Asked Questions

### What Python file contains the @tool decorator implementation?

The **`@tool`** decorator and the **`TOOL_REGISTRY`** are implemented in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py). This file handles all registration logic, signature inspection, and metadata extraction at import time.

### How does Needle 2 convert Python type hints into LLM-compatible schemas?

The decorator uses Python’s `inspect` module to parse function signatures at definition time. It maps primitive types (int, str, float, bool) and complex types (dataclasses, lists, dicts) to JSON Schema formats, allowing the LLM to receive a structured description of required arguments and return types.

### Can I register a tool under a different name than the function name?

Yes. Pass the **`name`** parameter to the decorator: `@tool(name="custom_name")`. This overrides the default behavior of using the function’s `__name__` attribute as the registry key, which is useful when function names are internal or not descriptive enough for the LLM.

### How do I verify that my custom tool is available to the Needle agent?

Import `TOOL_REGISTRY` from `needle/agent/tools` in your test file and assert that your tool’s name exists as a key in the dictionary. You can also call the registered function directly via `TOOL_REGISTRY["tool_name"]()` to verify its behavior without initializing the full agent loop.