How the @needle.tool Decorator Extracts Function Schemas and Docstrings

The @needle.tool decorator automatically generates OpenAI-compatible JSON schemas by inspecting function signatures, resolving type hints, and parsing docstrings to attach a complete _needle_tool attribute to the callable.

The @needle.tool decorator in the cactus-compute/needle repository enables Python functions to be exposed to Large Language Models (LLMs) via structured schemas. When applied to a function, the decorator invokes build_schema() to construct a JSON-compatible representation of the function's interface, storing the result as a private attribute for consumption by LLM backends.

How the Decorator Attaches Schemas

The decorator implementation in needle/agent/tools.py (lines 64-66) acts as a thin wrapper around the build_schema function. When you apply @tool to a function, it immediately calls build_schema(fn) and assigns the returned dictionary to the function object's _needle_tool attribute. This attachment allows LLM clients to retrieve the schema without modifying the function's behavior or call signature.

from needle.agent.tools import tool

@tool
def example():
    pass

print(example._needle_tool)  # Access the generated schema

The Schema Construction Pipeline

Signature Inspection

The build_schema function begins by obtaining the function's inspect.signature and resolving type hints using typing.get_type_hints. According to the source code at lines 11-17 and 111-119, this resolution supports modern Python typing constructs including PEP 604 union types (e.g., str | int), generic collections, and typing.Annotated metadata.

Docstring Parsing

A lightweight internal parser called _parse_doc (lines 95-108) extracts structured information from the function's docstring. The parser returns a clean summary description and a mapping of argument names to their informal descriptions, typically scanning for Args:, Parameters:, or similar sections.

Type-to-JSON Conversion

The _json_type function (lines 57-82) performs the critical translation of Python type annotations into JSON Schema fragments. This includes:

• Primitive mapping: str becomes {"type": "string"}, int becomes {"type": "integer"}
• Collection handling: list[int] transforms to {"type": "array", "items": {"type": "integer"}}
• Advanced types: Supports enum.Enum, typing.Literal, Pydantic models, and optional/union types through recursive resolution

Field Metadata Handling

When parameters use typing.Annotated with a custom Field instance, the _field_of helper (lines 85-93) extracts that metadata. The Field.apply method (lines 36-49) then injects constraints such as description, enum values, minimum/maximum bounds, or default values directly into the generated schema properties.

Required vs Optional Parameters

The logic at lines 30-35 determines parameter necessity. Parameters without default values and not annotated as Optional[...] (or Union[..., None]) are added to the schema's required array. Explicit defaults and Field defaults are respected and included in the property definitions.

The Final Schema Structure

The resulting dictionary, assembled at lines 135-142, conforms to an OpenAI-style function schema:

• name: The function name as a string
• description: The parsed docstring summary
• parameters: A JSON Schema object with type: "object", a properties object for each argument, and an optional required array listing mandatory parameters

Practical Code Examples

Basic function decoration automatically extracts types and docstrings:

from needle.agent.tools import tool

@tool
def multiply(a: int, b: int = 2) -> int:
    """Multiply two numbers.

    Args:
        a: The multiplicand.
        b: The multiplier (default 2).
    """
    return a * b

# Access via multiply._needle_tool

Resulting schema:

{
  "name": "multiply",
  "description": "Multiply two numbers.",
  "parameters": {
    "type": "object",
    "properties": {
      "a": {"type": "integer", "description": "The multiplicand."},
      "b": {"type": "integer", "description": "The multiplier (default 2).", "default": 2}
    },
    "required": ["a"]
  }
}

Advanced usage with Field metadata provides fine-grained control:

from needle.agent.tools import tool, Field

@tool
def greet(name: str, loud: bool = Field(default=False, description="Speak loudly")):
    """Greet a user.

    Parameters:
        name: The person's name.
        loud: Whether to shout.
    """
    return f"Hello, {name}{'!!!' if loud else '.'}"

This generates enhanced property definitions:

{
  "name": "greet",
  "description": "Greet a user.",
  "parameters": {
    "type": "object",
    "properties": {
      "name": {"type": "string", "description": "The person's name."},
      "loud": {"type": "boolean", "description": "Speak loudly", "default": false}
    },
    "required": ["name"]
  }
}

Summary

• The @needle.tool decorator stores generated schemas in the _needle_tool attribute of decorated functions (lines 64-66 in needle/agent/tools.py).
• Type resolution uses inspect.signature and typing.get_type_hints to support modern Python annotations including PEP 604 unions and Annotated types.
• Docstring parsing extracts descriptions via _parse_doc, mapping argument documentation to schema properties.
• JSON Schema generation converts Python types via _json_type, handling primitives, collections, enums, literals, and Pydantic models.
• Field metadata allows explicit schema constraints through Field objects processed by _field_of and Field.apply.
• Required parameters are automatically detected based on the absence of defaults and Optional wrappers.

Frequently Asked Questions

How does @needle.tool handle optional parameters in Python functions?

Parameters annotated as Optional[T] or Union[T, None] are marked as optional in the schema unless they have a default value. The logic at lines 30-35 checks for None types in the annotation; if present and no default exists, the parameter is excluded from the required array, though it remains in the properties object with a null or specified type.

What docstring formats does the _parse_doc function support?

The _parse_doc implementation (lines 95-108) recognizes common docstring conventions including Google-style Args: sections, NumPy-style Parameters: sections, and standard reST param directives. It extracts a summary from the first line and maps subsequent argument descriptions to parameter names for inclusion in the generated schema's property descriptions.

Can @needle.tool convert Pydantic models into JSON schema properties?

Yes. The _json_type function (lines 57-82) specifically checks for Pydantic models using issubclass checks and model_json_schema() calls. When a parameter type is a Pydantic model, the decorator recursively converts its fields into the properties object, preserving nested structures, field descriptions, and validation constraints within the parent function's parameter schema.

Where is the generated schema stored on the decorated function?

The decorator assigns the complete schema dictionary to the private attribute _needle_tool on the function object (lines 64-66 in needle/agent/tools.py). This attribute contains the name, description, and parameters keys necessary for LLM function calling APIs, accessible directly via function._needle_tool without invoking the function itself.

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