How to Declare Tools for the Needle Agent Using Decorators

Apply the @tool decorator from needle.agent.tools to any Python function to automatically generate OpenAI-compatible JSON schemas from type hints and docstrings, enabling seamless agent integration without manual registration.

The cactus-compute/needle repository provides a lightweight, decorator-based system for extending AI agent capabilities. When you declare tools for the Needle Agent using decorators, the framework automatically handles schema generation, type validation, and runtime discovery. This eliminates the need to hand-write JSON schemas or maintain separate registration boilerplate.

The @tool Decorator Mechanism

The core implementation lives in needle/agent/tools.py at lines 68-71. The decorator wraps any callable and stores the generated schema in a private attribute called _needle_tool on the function object.

When the agent initializes, it can collect all decorated functions automatically. Because the schema attaches directly to the function, you avoid explicit registration calls or configuration files. The decorator delegates the heavy lifting to the build_schema helper function (lines 15-46), which inspects the function signature and produces a complete OpenAI "function calling" compatible description.

Schema Generation with build_schema

The build_schema function (lines 15-46 in needle/agent/tools.py) performs deep introspection on the decorated callable. It extracts:

• Parameter types including primitives, unions, and container types
• Optional values marked with typing.Optional or default arguments
• Literal types and Enum members for constrained choices
• Pydantic models for complex nested structures
• Documentation from the function docstring and Field descriptions

The resulting schema follows this structure:

{
  "name": "<function_name>",
  "description": "<docstring-derived description>",
  "parameters": {
    "type": "object",
    "properties": {  },
    "required": [  ]
  }
}

Declaring Your First Tool

Basic Function Declaration

To declare a tool for the Needle Agent using decorators, import the @tool decorator and apply it to a regularly typed Python function:

from needle.agent.tools import tool

@tool
def echo(message: str) -> str:
    """Return the same text that was sent.

    Args:
        message: The text to be echoed back.
    """
    return message

The decorator automatically converts the type hint str and the docstring into the corresponding JSON schema. You can access the generated schema directly through the _needle_tool attribute:

>>> echo._needle_tool
{
  "name": "echo",
  "description": "Return the same text that was sent.",
  "parameters": {
    "type": "object",
    "properties": {
      "message": {"type": "string"}
    },
    "required": ["message"]
  }
}

Advanced Type Hints and Pydantic Models

The decorator supports complex type systems including Enums, Literals, and Pydantic Field metadata for richer tool definitions:

from enum import Enum
from typing import Literal, Optional
from needle.agent.tools import tool, Field

class Color(Enum):
    RED = "red"
    GREEN = "green"
    BLUE = "blue"

@tool
def paint(
    surface: str,
    color: Color = Color.RED,
    intensity: Optional[int] = Field(default=5, description="Brightness 1-10")
) -> bool:
    """Paint a given surface with a chosen colour.

    Parameters:
        surface: Name of the surface to paint.
        color: The colour to apply.
        intensity: Optional brightness level.
    """
    return True

In this example, the build_schema function correctly interprets the Color Enum choices, the default value, and the Field description for the intensity parameter.

Runtime Tool Discovery

Since the @tool decorator stores the schema on the function object itself, you can collect all tools from a module using standard Python introspection:

import inspect
import needle.agent.tools as tools_mod

def load_all_tools(module):
    return {
        name: obj._needle_tool
        for name, obj in inspect.getmembers(module, inspect.isfunction)
        if hasattr(obj, "_needle_tool")
    }

all_tool_schemas = load_all_tools(tools_mod)

The Needle runtime uses this exact pattern to inject discovered tools into the agent's toolset at initialization. No additional registration step is required because the decorator has already done the work of schema generation and attachment.

Summary

• The @tool decorator in needle/agent/tools.py (lines 68-71) automatically transforms Python functions into OpenAI-compatible tool schemas.
• The build_schema helper (lines 15-46) extracts type hints, docstrings, and Pydantic models to generate complete JSON schemas.
• Decorated functions store their schemas in the _needle_tool attribute, enabling runtime discovery via inspect.getmembers.
• The system supports complex types including Enums, Literals, Optional fields, and Pydantic models with metadata.
• No manual registration is required; the Needle Agent discovers tools automatically by scanning for the _needle_tool attribute.

Frequently Asked Questions

What file contains the @tool decorator implementation?

The @tool decorator is implemented in needle/agent/tools.py at lines 68-71, with the underlying schema generation logic in the build_schema function at lines 15-46 according to the cactus-compute/needle source code.

Does the @tool decorator support Pydantic models for complex inputs?

Yes, the build_schema function inspects Pydantic models and converts them into nested JSON schema properties. You can use Field objects to provide descriptions and constraints, and mix Pydantic models with standard Python type hints like Optional and Literal.

How does the Needle Agent discover functions decorated with @tool?

The agent discovers tools by scanning modules for functions that possess the _needle_tool attribute. This attribute contains the pre-generated JSON schema, allowing the runtime to collect and inject tools without explicit registration calls or configuration files.

What JSON schema format does the @tool decorator produce?

The decorator generates schemas compatible with the OpenAI function calling format. Each schema includes the function name, a description derived from the docstring, a "parameters" object defining properties and required fields, and type mappings that reflect the original Python type hints.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}